diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/sentry_sdk/integrations/grpc/aio')
3 files changed, 201 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/grpc/aio/__init__.py b/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/grpc/aio/__init__.py new file mode 100644 index 00000000..5b9e3b99 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/grpc/aio/__init__.py @@ -0,0 +1,7 @@ +from .server import ServerInterceptor +from .client import ClientInterceptor + +__all__ = [ + "ClientInterceptor", + "ServerInterceptor", +] diff --git a/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/grpc/aio/client.py b/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/grpc/aio/client.py new file mode 100644 index 00000000..ff3c2131 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/grpc/aio/client.py @@ -0,0 +1,94 @@ +from typing import Callable, Union, AsyncIterable, Any + +from grpc.aio import ( + UnaryUnaryClientInterceptor, + UnaryStreamClientInterceptor, + ClientCallDetails, + UnaryUnaryCall, + UnaryStreamCall, + Metadata, +) +from google.protobuf.message import Message + +import sentry_sdk +from sentry_sdk.consts import OP +from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN + + +class ClientInterceptor: + @staticmethod + def _update_client_call_details_metadata_from_scope( + client_call_details: ClientCallDetails, + ) -> ClientCallDetails: + if client_call_details.metadata is None: + client_call_details = client_call_details._replace(metadata=Metadata()) + elif not isinstance(client_call_details.metadata, Metadata): + # This is a workaround for a GRPC bug, which was fixed in grpcio v1.60.0 + # See https://github.com/grpc/grpc/issues/34298. + client_call_details = client_call_details._replace( + metadata=Metadata.from_tuple(client_call_details.metadata) + ) + for ( + key, + value, + ) in sentry_sdk.get_current_scope().iter_trace_propagation_headers(): + client_call_details.metadata.add(key, value) + return client_call_details + + +class SentryUnaryUnaryClientInterceptor(ClientInterceptor, UnaryUnaryClientInterceptor): # type: ignore + async def intercept_unary_unary( + self, + continuation: Callable[[ClientCallDetails, Message], UnaryUnaryCall], + client_call_details: ClientCallDetails, + request: Message, + ) -> Union[UnaryUnaryCall, Message]: + method = client_call_details.method + + with sentry_sdk.start_span( + op=OP.GRPC_CLIENT, + name="unary unary call to %s" % method.decode(), + origin=SPAN_ORIGIN, + ) as span: + span.set_data("type", "unary unary") + span.set_data("method", method) + + client_call_details = self._update_client_call_details_metadata_from_scope( + client_call_details + ) + + response = await continuation(client_call_details, request) + status_code = await response.code() + span.set_data("code", status_code.name) + + return response + + +class SentryUnaryStreamClientInterceptor( + ClientInterceptor, UnaryStreamClientInterceptor # type: ignore +): + async def intercept_unary_stream( + self, + continuation: Callable[[ClientCallDetails, Message], UnaryStreamCall], + client_call_details: ClientCallDetails, + request: Message, + ) -> Union[AsyncIterable[Any], UnaryStreamCall]: + method = client_call_details.method + + with sentry_sdk.start_span( + op=OP.GRPC_CLIENT, + name="unary stream call to %s" % method.decode(), + origin=SPAN_ORIGIN, + ) as span: + span.set_data("type", "unary stream") + span.set_data("method", method) + + client_call_details = self._update_client_call_details_metadata_from_scope( + client_call_details + ) + + response = await continuation(client_call_details, request) + # status_code = await response.code() + # span.set_data("code", status_code) + + return response diff --git a/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/grpc/aio/server.py b/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/grpc/aio/server.py new file mode 100644 index 00000000..381c6310 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/grpc/aio/server.py @@ -0,0 +1,100 @@ +import sentry_sdk +from sentry_sdk.consts import OP +from sentry_sdk.integrations import DidNotEnable +from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN +from sentry_sdk.tracing import Transaction, TransactionSource +from sentry_sdk.utils import event_from_exception + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + from typing import Any, Optional + + +try: + import grpc + from grpc import HandlerCallDetails, RpcMethodHandler + from grpc.aio import AbortError, ServicerContext +except ImportError: + raise DidNotEnable("grpcio is not installed") + + +class ServerInterceptor(grpc.aio.ServerInterceptor): # type: ignore + def __init__(self, find_name=None): + # type: (ServerInterceptor, Callable[[ServicerContext], str] | None) -> None + self._find_method_name = find_name or self._find_name + + super().__init__() + + async def intercept_service(self, continuation, handler_call_details): + # type: (ServerInterceptor, Callable[[HandlerCallDetails], Awaitable[RpcMethodHandler]], HandlerCallDetails) -> Optional[Awaitable[RpcMethodHandler]] + self._handler_call_details = handler_call_details + handler = await continuation(handler_call_details) + if handler is None: + return None + + if not handler.request_streaming and not handler.response_streaming: + handler_factory = grpc.unary_unary_rpc_method_handler + + async def wrapped(request, context): + # type: (Any, ServicerContext) -> Any + name = self._find_method_name(context) + if not name: + return await handler(request, context) + + # What if the headers are empty? + transaction = Transaction.continue_from_headers( + dict(context.invocation_metadata()), + op=OP.GRPC_SERVER, + name=name, + source=TransactionSource.CUSTOM, + origin=SPAN_ORIGIN, + ) + + with sentry_sdk.start_transaction(transaction=transaction): + try: + return await handler.unary_unary(request, context) + except AbortError: + raise + except Exception as exc: + event, hint = event_from_exception( + exc, + mechanism={"type": "grpc", "handled": False}, + ) + sentry_sdk.capture_event(event, hint=hint) + raise + + elif not handler.request_streaming and handler.response_streaming: + handler_factory = grpc.unary_stream_rpc_method_handler + + async def wrapped(request, context): # type: ignore + # type: (Any, ServicerContext) -> Any + async for r in handler.unary_stream(request, context): + yield r + + elif handler.request_streaming and not handler.response_streaming: + handler_factory = grpc.stream_unary_rpc_method_handler + + async def wrapped(request, context): + # type: (Any, ServicerContext) -> Any + response = handler.stream_unary(request, context) + return await response + + elif handler.request_streaming and handler.response_streaming: + handler_factory = grpc.stream_stream_rpc_method_handler + + async def wrapped(request, context): # type: ignore + # type: (Any, ServicerContext) -> Any + async for r in handler.stream_stream(request, context): + yield r + + return handler_factory( + wrapped, + request_deserializer=handler.request_deserializer, + response_serializer=handler.response_serializer, + ) + + def _find_name(self, context): + # type: (ServicerContext) -> str + return self._handler_call_details.method |