server.py 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566
  1. import sentry_sdk
  2. from sentry_sdk.consts import OP
  3. from sentry_sdk.integrations import DidNotEnable
  4. from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN
  5. from sentry_sdk.tracing import Transaction, TransactionSource
  6. from typing import TYPE_CHECKING
  7. if TYPE_CHECKING:
  8. from typing import Callable, Optional
  9. from google.protobuf.message import Message
  10. try:
  11. import grpc
  12. from grpc import ServicerContext, HandlerCallDetails, RpcMethodHandler
  13. except ImportError:
  14. raise DidNotEnable("grpcio is not installed")
  15. class ServerInterceptor(grpc.ServerInterceptor): # type: ignore
  16. def __init__(self, find_name=None):
  17. # type: (ServerInterceptor, Optional[Callable[[ServicerContext], str]]) -> None
  18. self._find_method_name = find_name or ServerInterceptor._find_name
  19. super().__init__()
  20. def intercept_service(self, continuation, handler_call_details):
  21. # type: (ServerInterceptor, Callable[[HandlerCallDetails], RpcMethodHandler], HandlerCallDetails) -> RpcMethodHandler
  22. handler = continuation(handler_call_details)
  23. if not handler or not handler.unary_unary:
  24. return handler
  25. def behavior(request, context):
  26. # type: (Message, ServicerContext) -> Message
  27. with sentry_sdk.isolation_scope():
  28. name = self._find_method_name(context)
  29. if name:
  30. metadata = dict(context.invocation_metadata())
  31. transaction = Transaction.continue_from_headers(
  32. metadata,
  33. op=OP.GRPC_SERVER,
  34. name=name,
  35. source=TransactionSource.CUSTOM,
  36. origin=SPAN_ORIGIN,
  37. )
  38. with sentry_sdk.start_transaction(transaction=transaction):
  39. try:
  40. return handler.unary_unary(request, context)
  41. except BaseException as e:
  42. raise e
  43. else:
  44. return handler.unary_unary(request, context)
  45. return grpc.unary_unary_rpc_method_handler(
  46. behavior,
  47. request_deserializer=handler.request_deserializer,
  48. response_serializer=handler.response_serializer,
  49. )
  50. @staticmethod
  51. def _find_name(context):
  52. # type: (ServicerContext) -> str
  53. return context._rpc_event.call_details.method.decode()