feat(integrations): Add integerations for grpc.aio

sentry_sdk/integrations/grpc/aio/__init__.py

@@ -0,0 +1 @@
+from .server import ServerInterceptor  # noqa: F401

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()

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",

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)

tox.ini

@@ -303,6 +303,7 @@ deps =
     grpc: protobuf
     grpc: mypy-protobuf
     grpc: types-protobuf
+    grpc: pytest-asyncio
 
     # HTTPX
     httpx: pytest-httpx