diff --git a/sentry_sdk/integrations/grpc/aio/__init__.py b/sentry_sdk/integrations/grpc/aio/__init__.py new file mode 100644 index 0000000000..c3695e948e --- /dev/null +++ b/sentry_sdk/integrations/grpc/aio/__init__.py @@ -0,0 +1 @@ +from .server import ServerInterceptor # noqa: F401 diff --git a/sentry_sdk/integrations/grpc/aio/server.py b/sentry_sdk/integrations/grpc/aio/server.py new file mode 100644 index 0000000000..34554d866f --- /dev/null +++ b/sentry_sdk/integrations/grpc/aio/server.py @@ -0,0 +1,70 @@ +from functools import wraps + +from sentry_sdk import Hub +from sentry_sdk._types import MYPY +from sentry_sdk.consts import OP +from sentry_sdk.integrations import DidNotEnable +from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_CUSTOM +from sentry_sdk.utils import event_from_exception + +if MYPY: + from collections.abc import Awaitable, Callable + from typing import Any + + +try: + import grpc + from grpc import HandlerCallDetails, RpcMethodHandler + from grpc.aio import ServicerContext + from grpc.experimental import wrap_server_method_handler +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 ServerInterceptor._find_name + + super(ServerInterceptor, self).__init__() + + async def intercept_service(self, continuation, handler_call_details): + # type: (ServerInterceptor, Callable[[HandlerCallDetails], Awaitable[RpcMethodHandler]], HandlerCallDetails) -> Awaitable[RpcMethodHandler] + handler = await continuation(handler_call_details) + return wrap_server_method_handler(self.wrapper, handler) + + def wrapper(self, handler): + # type: (ServerInterceptor, Callable[[Any, ServicerContext], Awaitable[Any]]) -> Callable[[Any, ServicerContext], Awaitable[Any]] + @wraps(handler) + async def wrapped(request, context): + # type: (Any, ServicerContext) -> Any + name = self._find_method_name(context) + if not name: + return await handler(request, context) + + hub = Hub(Hub.current) + + transaction = Transaction.continue_from_headers( + dict(context.invocation_metadata()), + op=OP.GRPC_SERVER, + name=name, + source=TRANSACTION_SOURCE_CUSTOM, + ) + + with hub.start_transaction(transaction=transaction): + try: + return await handler(request, context) + except Exception as exc: + event, hint = event_from_exception( + exc, + mechanism={"type": "gRPC", "handled": False}, + ) + hub.capture_event(event, hint=hint) + raise + + return wrapped + + @staticmethod + def _find_name(context): + # type: (ServicerContext) -> str + return context._rpc_event.call_details.method.decode() diff --git a/setup.py b/setup.py index 9dd2e13b79..e6a6893e65 100644 --- a/setup.py +++ b/setup.py @@ -67,7 +67,7 @@ def get_file_text(file_name): "fastapi": ["fastapi>=0.79.0"], "pymongo": ["pymongo>=3.1"], "opentelemetry": ["opentelemetry-distro>=0.35b0"], - "grpcio": ["grpcio>=1.21.1"] + "grpcio": ["grpcio>=1.21.1"], }, classifiers=[ "Development Status :: 5 - Production/Stable", diff --git a/tests/integrations/grpc/test_aio_grpc.py b/tests/integrations/grpc/test_aio_grpc.py new file mode 100644 index 0000000000..50ad3f7698 --- /dev/null +++ b/tests/integrations/grpc/test_aio_grpc.py @@ -0,0 +1,148 @@ +from __future__ import absolute_import + +import asyncio +import os + +import grpc +import pytest +import pytest_asyncio +import sentry_sdk + +from sentry_sdk import Hub +from sentry_sdk.consts import OP +from sentry_sdk.integrations.grpc.aio.server import ServerInterceptor +from tests.integrations.grpc.grpc_test_service_pb2 import gRPCTestMessage +from tests.integrations.grpc.grpc_test_service_pb2_grpc import ( + gRPCTestServiceServicer, + add_gRPCTestServiceServicer_to_server, + gRPCTestServiceStub, +) + +PORT = 50151 +PORT += os.getpid() % 100 # avoid port conflicts when running tests in parallel + + +@pytest.fixture(scope="module") +def event_loop(): + """Create an instance of the default event loop for each test case.""" + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +@pytest_asyncio.fixture(scope="function") +async def grpc_server(event_loop): + server = grpc.aio.server( + interceptors=[ServerInterceptor(find_name=lambda request: request.__class__)] + ) + server.add_insecure_port(f"[::]:{PORT}") + add_gRPCTestServiceServicer_to_server(TestService, server) + + await server.start() + + try: + yield server + finally: + await server.stop(None) + await server.wait_for_termination() + + +@pytest.mark.asyncio +async def test_grpc_server_starts_transaction(sentry_init, capture_events, grpc_server): + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + async with grpc.aio.insecure_channel(f"localhost:{PORT}") as channel: + stub = gRPCTestServiceStub(channel) + await stub.TestServe(gRPCTestMessage(text="test")) + + (event,) = events + span = event["spans"][0] + + assert event["type"] == "transaction" + assert event["transaction_info"] == { + "source": "custom", + } + assert event["contexts"]["trace"]["op"] == OP.GRPC_SERVER + assert span["op"] == "test" + + +@pytest.mark.asyncio +async def test_grpc_server_continues_transaction( + sentry_init, capture_events, grpc_server +): + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + async with grpc.aio.insecure_channel(f"localhost:{PORT}") as channel: + stub = gRPCTestServiceStub(channel) + + with sentry_sdk.start_transaction() as transaction: + metadata = ( + ( + "baggage", + "sentry-trace_id={trace_id},sentry-environment=test," + "sentry-transaction=test-transaction,sentry-sample_rate=1.0".format( + trace_id=transaction.trace_id + ), + ), + ( + "sentry-trace", + "{trace_id}-{parent_span_id}-{sampled}".format( + trace_id=transaction.trace_id, + parent_span_id=transaction.span_id, + sampled=1, + ), + ), + ) + + await stub.TestServe(gRPCTestMessage(text="test"), metadata=metadata) + + (event, _) = events + span = event["spans"][0] + + assert event["type"] == "transaction" + assert event["transaction_info"] == { + "source": "custom", + } + assert event["contexts"]["trace"]["op"] == OP.GRPC_SERVER + assert event["contexts"]["trace"]["trace_id"] == transaction.trace_id + assert span["op"] == "test" + + +@pytest.mark.asyncio +async def test_grpc_server_exception(sentry_init, capture_events, grpc_server): + sentry_init(traces_sample_rate=1.0) + events = capture_events() + + async with grpc.aio.insecure_channel(f"localhost:{PORT}") as channel: + stub = gRPCTestServiceStub(channel) + try: + await stub.TestServe(gRPCTestMessage(text="exception")) + raise AssertionError() + except Exception: + pass + + (event, _) = events + + assert event["exception"]["values"][0]["type"] == "TestService.ExceptionForTest" + assert event["exception"]["values"][0]["value"] == "test" + assert event["exception"]["values"][0]["mechanism"]["handled"] is False + assert event["exception"]["values"][0]["mechanism"]["type"] == "gRPC" + + +class TestService(gRPCTestServiceServicer): + class ExceptionForTest(Exception): + def __init__(self): + super().__init__("test") + + @classmethod + async def TestServe(cls, request, context): # noqa: N802 + hub = Hub.current + with hub.start_span(op="test", description="test"): + pass + + if request.text == "exception": + raise cls.ExceptionForTest() + + return gRPCTestMessage(text=request.text) diff --git a/tox.ini b/tox.ini index 7632af225f..cc880ca7c7 100644 --- a/tox.ini +++ b/tox.ini @@ -303,6 +303,7 @@ deps = grpc: protobuf grpc: mypy-protobuf grpc: types-protobuf + grpc: pytest-asyncio # HTTPX httpx: pytest-httpx