feat(tracing) Implement unary-unary and unary-stream interceptors for grpc stub

Implement gRPC UnaryUnaryClientInterceptor and grpc.UnaryStreamClientInterceptor; starts a span for grpc call and sends propagation headers in metadata

Author: hossein-raeisi

sentry_sdk/consts.py

@@ -80,6 +80,7 @@ class OP:
     SOCKET_CONNECTION = "socket.connection"
     SOCKET_DNS = "socket.dns"
     GRPC_SERVER = "grpc.server"
+    GRPC_CLIENT = "grpc.client"
 
 
 # This type exists to trick mypy and PyCharm into thinking `init` and `Client`

sentry_sdk/integrations/grpc/__init__.py

@@ -1 +1,2 @@
 from .server import ServerInterceptor
+from .client import ClientInterceptor

sentry_sdk/integrations/grpc/client.py

@@ -0,0 +1,73 @@
+from sentry_sdk import Hub
+from sentry_sdk._types import MYPY
+from sentry_sdk.consts import OP
+from sentry_sdk.integrations import DidNotEnable
+
+if MYPY:
+    pass
+
+try:
+    import grpc
+except ImportError:
+    raise DidNotEnable("grpcio is not installed")
+
+
+class ClientInterceptor(
+    grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor
+):
+    def intercept_unary_unary(self, continuation, client_call_details, request):
+        hub = Hub.current
+        method = client_call_details.method
+
+        with hub.start_span(
+            op=OP.GRPC_CLIENT, description="unary unary call to %s" % method
+        ) 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 = continuation(client_call_details, request)
+            span.set_data("code", response.code().name)
+
+            return response
+
+    def intercept_unary_stream(self, continuation, client_call_details, request):
+        hub = Hub.current
+        method = client_call_details.method
+
+        with hub.start_span(
+            op=OP.GRPC_CLIENT, description="unary stream call to %s" % method
+        ) as span:
+            span.set_data("type", "unary stream")
+            span.set_data("method", method)
+
+            client_call_details = self._update_client_call_details_metadata_from_hub(
+                client_call_details, hub
+            )
+
+            response = continuation(client_call_details, request)
+            span.set_data("code", response.code().name)
+
+            return response
+
+    @staticmethod
+    def _update_client_call_details_metadata_from_hub(client_call_details, 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 = grpc._interceptor._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,
+            compression=client_call_details.compression,
+        )
+
+        return client_call_details

tests/integrations/grpc/test_server.py renamed to tests/integrations/grpc/test_grpc.py

@@ -7,6 +7,7 @@
 
 from sentry_sdk import Hub, start_transaction
 from sentry_sdk.consts import OP
+from sentry_sdk.integrations.grpc.client import ClientInterceptor
 from sentry_sdk.integrations.grpc.server import ServerInterceptor
 from tests.integrations.grpc.test_service_pb2 import TestMessage
 from tests.integrations.grpc.test_service_pb2_grpc import (
@@ -68,8 +69,8 @@ def test_grpc_server_continues_transaction(sentry_init, capture_events_forksafe)
                         trace_id=transaction.trace_id,
                         parent_span_id=transaction.span_id,
                         sampled=1,
-                    )
-                )
+                    ),
+                ),
             )
             stub.TestServe(TestMessage(text="test"), metadata=metadata)
 
@@ -88,6 +89,70 @@ def test_grpc_server_continues_transaction(sentry_init, capture_events_forksafe)
     assert span["op"] == "test"
 
 
+@pytest.mark.forked
+def test_grpc_client_starts_span(sentry_init, capture_events_forksafe):
+    sentry_init(traces_sample_rate=1.0)
+    events = capture_events_forksafe()
+    interceptors = [ClientInterceptor()]
+
+    server = _set_up()
+
+    with grpc.insecure_channel(f"localhost:{PORT}") as channel:
+        channel = grpc.intercept_channel(channel, *interceptors)
+        stub = TestServiceStub(channel)
+
+        with start_transaction():
+            stub.TestServe(TestMessage(text="test"))
+
+    _tear_down(server=server)
+
+    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 /test_grpc_server.TestService/TestServe"
+    )
+    assert span["data"] == {
+        "type": "unary unary",
+        "method": "/test_grpc_server.TestService/TestServe",
+        "code": "OK",
+    }
+
+
+@pytest.mark.forked
+def test_grpc_client_and_servers_interceptors_integration(
+    sentry_init, capture_events_forksafe
+):
+    sentry_init(traces_sample_rate=1.0)
+    events = capture_events_forksafe()
+    interceptors = [ClientInterceptor()]
+
+    server = _set_up()
+
+    with grpc.insecure_channel(f"localhost:{PORT}") as channel:
+        channel = grpc.intercept_channel(channel, *interceptors)
+        stub = TestServiceStub(channel)
+
+        with start_transaction():
+            stub.TestServe(TestMessage(text="test"))
+
+    _tear_down(server=server)
+
+    events.write_file.close()
+    server_transaction = events.read_event()
+    local_transaction = events.read_event()
+
+    assert (
+        server_transaction["contexts"]["trace"]["trace_id"]
+        == local_transaction["contexts"]["trace"]["trace_id"]
+    )
+
+
 def _set_up():
     server = grpc.server(
         futures.ThreadPoolExecutor(max_workers=2),