gRPC integration and aio interceptors

sentry_sdk/integrations/grpc/__init__.py

@@ -1,2 +1,89 @@
-from .server import ServerInterceptor  # noqa: F401
-from .client import ClientInterceptor  # noqa: F401
+from typing import List, Optional, Sequence, Callable, ParamSpec
+from functools import wraps
+
+import grpc
+from grpc import Channel, Server, intercept_channel
+from grpc.aio import Channel as AsyncChannel
+from grpc.aio import Server as AsyncServer
+
+from sentry_sdk.integrations import Integration
+from .client import ClientInterceptor
+from .server import ServerInterceptor
+from .aio.server import ServerInterceptor as AsyncServerInterceptor
+from .aio.client import AsyncClientInterceptor
+
+P = ParamSpec("P")
+
+
+def _wrap_channel_sync(func: Callable[P, Channel]) -> Callable[P, Channel]:
+    "Wrapper for synchronous secure and insecure channel."
+
+    @wraps(func)
+    def patched_channel(*args, **kwargs) -> Channel:
+        channel = func(*args, **kwargs)
+        return intercept_channel(channel, ClientInterceptor())
+
+    return patched_channel
+
+
+def _wrap_channel_async(func: Callable[P, AsyncChannel]) -> Callable[P, AsyncChannel]:
+    "Wrapper for asynchronous secure and insecure channel."
+
+    @wraps(func)
+    def patched_channel(
+        *args,
+        interceptors: Optional[Sequence[grpc.aio.ClientInterceptor]] = None,
+        **kwargs
+    ) -> Channel:
+        interceptor = AsyncClientInterceptor()
+        interceptors = [interceptor, *(interceptors or [])]
+        return func(*args, interceptors=interceptors, **kwargs)
+
+    return patched_channel
+
+
+def _wrap_sync_server(func: Callable[P, Server]) -> Callable[P, Server]:
+    """Wrapper for synchronous server."""
+
+    @wraps(func)
+    def patched_server(
+        *args, interceptors: Optional[List[grpc.ServerInterceptor]] = None, **kwargs
+    ) -> Server:
+        server_interceptor = ServerInterceptor()
+        interceptors = [server_interceptor, *(interceptors or [])]
+        return func(*args, interceptors=interceptors, **kwargs)
+
+    return patched_server
+
+
+def _wrap_async_server(func: Callable[P, AsyncServer]) -> Callable[P, AsyncServer]:
+    """Wrapper for asynchronous server."""
+
+    @wraps(func)
+    def patched_aio_server(
+        *args, interceptors: Optional[List[grpc.ServerInterceptor]] = None, **kwargs
+    ) -> Server:
+        server_interceptor = AsyncServerInterceptor(
+            find_name=lambda request: request.__class__
+        )
+        interceptors = [server_interceptor, *(interceptors or [])]
+        return func(*args, interceptors=interceptors, **kwargs)
+
+    return patched_aio_server
+
+
+class GRPCIntegration(Integration):
+    identifier = "grpc"
+
+    @staticmethod
+    def setup_once() -> None:
+        import grpc
+
+        grpc.insecure_channel = _wrap_channel_sync(grpc.insecure_channel)
+        grpc.secure_channel = _wrap_channel_sync(grpc.secure_channel)
+
+        grpc.aio.insecure_channel = _wrap_channel_async(grpc.aio.insecure_channel)
+        grpc.aio.secure_channel = _wrap_channel_async(grpc.aio.secure_channel)
+
+        grpc.server = _wrap_sync_server(grpc.server)
+        grpc.aio.server = _wrap_async_server(grpc.aio.server)

sentry_sdk/integrations/grpc/aio/__init__.py

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

sentry_sdk/integrations/grpc/aio/client.py

@@ -0,0 +1,56 @@
+from typing import Callable, Union
+
+from grpc.aio import UnaryUnaryClientInterceptor
+from grpc.aio._call import UnaryUnaryCall
+from grpc.aio._interceptor import ClientCallDetails
+from google.protobuf.message import Message
+
+from sentry_sdk import Hub
+from sentry_sdk.consts import OP
+
+
+class AsyncClientInterceptor(UnaryUnaryClientInterceptor):
+    async def intercept_unary_unary(
+        self,
+        continuation: Callable[[ClientCallDetails, Message], UnaryUnaryCall],
+        client_call_details: ClientCallDetails,
+        request: Message,
+    ) -> Union[UnaryUnaryCall, Message]:
+        hub = Hub.current
+        method = client_call_details.method
+
+        with hub.start_span(
+            op=OP.GRPC_CLIENT, description="unary unary call to %s" % method.decode()
+        ) as span:
+            span.set_data("type", "unary unary")
+            span.set_data("method", method)
+
+            client_call_details = self._update_client_call_details_metadata_from_hub(
+                client_call_details, hub
+            )
+
+            response = await continuation(client_call_details, request)
+            status_code = await response.code()
+            span.set_data("code", status_code.name)
+
+            return response
+
+    @staticmethod
+    def _update_client_call_details_metadata_from_hub(
+        client_call_details: ClientCallDetails, hub: Hub
+    ):
+        metadata = (
+            list(client_call_details.metadata) if client_call_details.metadata else []
+        )
+        for key, value in hub.iter_trace_propagation_headers():
+            metadata.append((key, value))
+
+        client_call_details = ClientCallDetails(
+            method=client_call_details.method,
+            timeout=client_call_details.timeout,
+            metadata=metadata,
+            credentials=client_call_details.credentials,
+            wait_for_ready=client_call_details.wait_for_ready,
+        )
+
+        return client_call_details

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

tests/integrations/grpc/test_aio_grpc.py

@@ -0,0 +1,169 @@
+from __future__ import absolute_import
+
+import asyncio
+import os
+
+import grpc
+import pytest
+import pytest_asyncio
+import sentry_sdk
+
+from sentry_sdk import Hub, start_transaction
+from sentry_sdk.consts import OP
+from sentry_sdk.integrations.grpc import GRPCIntegration
+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,
+)
+
+AIO_PORT = 50052
+AIO_PORT += os.getpid() % 100  # avoid port conflicts when running tests in parallel
+
+
+@pytest.fixture(scope="function")
+def event_loop(request):
+    """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(sentry_init, event_loop):
+    sentry_init(traces_sample_rate=1.0, integrations=[GRPCIntegration()])
+    server = grpc.aio.server()
+    server.add_insecure_port(f"[::]:{AIO_PORT}")
+    add_gRPCTestServiceServicer_to_server(TestService, server)
+
+    await event_loop.create_task(server.start())
+
+    try:
+        yield server
+    finally:
+        await server.stop(None)
+
+
+@pytest.mark.asyncio
+async def test_grpc_server_starts_transaction(capture_events, grpc_server):
+    events = capture_events()
+
+    async with grpc.aio.insecure_channel(f"localhost:{AIO_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(capture_events, grpc_server):
+    events = capture_events()
+
+    async with grpc.aio.insecure_channel(f"localhost:{AIO_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:{AIO_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.TestException"
+    assert event["exception"]["values"][0]["value"] == "test"
+    assert event["exception"]["values"][0]["mechanism"]["handled"] is False
+    assert event["exception"]["values"][0]["mechanism"]["type"] == "gRPC"
+
+
+@pytest.mark.asyncio
+async def test_grpc_client_starts_span(grpc_server, capture_events_forksafe):
+    events = capture_events_forksafe()
+
+    async with grpc.aio.insecure_channel(f"localhost:{AIO_PORT}") as channel:
+        stub = gRPCTestServiceStub(channel)
+        with start_transaction():
+            await stub.TestServe(gRPCTestMessage(text="test"))
+
+    events.write_file.close()
+    events.read_event()
+    local_transaction = events.read_event()
+    span = local_transaction["spans"][0]
+
+    assert len(local_transaction["spans"]) == 1
+    assert span["op"] == OP.GRPC_CLIENT
+    assert (
+        span["description"]
+        == "unary unary call to /grpc_test_server.gRPCTestService/TestServe"
+    )
+    assert span["data"] == {
+        "type": "unary unary",
+        "method": "/grpc_test_server.gRPCTestService/TestServe",
+        "code": "OK",
+    }
+
+
+class TestService(gRPCTestServiceServicer):
+    class TestException(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.TestException()
+
+        return gRPCTestMessage(text=request.text)