from typing import Any, AsyncIterable, Callable, Union import sentry_sdk from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN from sentry_sdk.tracing_utils import has_span_streaming_enabled try: from google.protobuf.message import Message from grpc.aio import ( ClientCallDetails, Metadata, UnaryStreamCall, UnaryStreamClientInterceptor, UnaryUnaryCall, UnaryUnaryClientInterceptor, ) except ImportError: raise DidNotEnable("grpcio is not installed") 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 span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) if span_streaming: with sentry_sdk.traces.start_span( name="unary unary call to %s" % method.decode(), attributes={ "sentry.op": OP.GRPC_CLIENT, "sentry.origin": SPAN_ORIGIN, SPANDATA.RPC_METHOD: method.decode(), }, ) as span: 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_attribute(SPANDATA.RPC_RESPONSE_STATUS_CODE, status_code.name) return response else: 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 span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) if span_streaming: with sentry_sdk.traces.start_span( name="unary stream call to %s" % method.decode(), attributes={ "sentry.op": OP.GRPC_CLIENT, "sentry.origin": SPAN_ORIGIN, SPANDATA.RPC_METHOD: method.decode(), }, ) as span: client_call_details = ( self._update_client_call_details_metadata_from_scope( client_call_details ) ) response = await continuation(client_call_details, request) return response else: 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