Respect suppress instrumentation key in gRPC client

open-telemetry · Jul 1, 2021 · afe84d2 · afe84d2
1 parent c100b21
commit afe84d2
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 2 deletions.
diff --git a/...tion/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_client.py b/...tion/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_client.py
@@ -24,9 +24,10 @@
 
 import grpc
 
-from opentelemetry import trace
+from opentelemetry import context, trace
 from opentelemetry.instrumentation.grpc import grpcext
 from opentelemetry.instrumentation.grpc._utilities import RpcInfo
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.propagate import inject
 from opentelemetry.propagators.textmap import Setter
 from opentelemetry.semconv.trace import SpanAttributes
@@ -122,6 +123,9 @@ def _start_guarded_span(self, *args, **kwargs):
         return _GuardedSpan(self._start_span(*args, **kwargs))
 
     def intercept_unary(self, request, metadata, client_info, invoker):
+        if context.get_value( _SUPPRESS_INSTRUMENTATION_KEY):
+            return invoker(request, metadata)
+
         if not metadata:
             mutable_metadata = OrderedDict()
         else:
@@ -189,6 +193,9 @@ def _intercept_server_stream(
     def intercept_stream(
         self, request_or_iterator, metadata, client_info, invoker
     ):
+        if context.get_value( _SUPPRESS_INSTRUMENTATION_KEY):
+            return invoker(request_or_iterator, metadata)
+
         if client_info.is_server_stream:
             return self._intercept_server_stream(
                 request_or_iterator, metadata, client_info, invoker

diff --git a/instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor.py b/instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor.py
@@ -18,14 +18,15 @@
 )
 
 import opentelemetry.instrumentation.grpc
-from opentelemetry import trace
+from opentelemetry import context, trace
 from opentelemetry.instrumentation.grpc import GrpcInstrumentorClient
 from opentelemetry.instrumentation.grpc._client import (
     OpenTelemetryClientInterceptor,
 )
 from opentelemetry.instrumentation.grpc.grpcext._interceptor import (
     _UnaryClientInfo,
 )
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.propagate import get_global_textmap, set_global_textmap
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.test.mock_textmap import MockTextMapPropagator
@@ -178,6 +179,7 @@ def test_stream_unary(self):
             },
         )
 
+
     def test_stream_stream(self):
         bidirectional_streaming_method(self._stub)
         spans = self.memory_exporter.get_finished_spans()
@@ -286,3 +288,47 @@ def invoker(request, metadata):
 
         finally:
             set_global_textmap(previous_propagator)
+
+    def test_unary_unary_with_suppress_key(self):
+        token = context.attach(
+            context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)
+        )
+        try:
+            simple_method(self._stub)
+            spans = self.memory_exporter.get_finished_spans()
+        finally:
+            context.detach(token)
+        self.assertEqual(len(spans), 0)
+
+    def test_unary_stream_with_suppress_key(self):
+        token = context.attach(
+            context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)
+        )
+        try:
+            server_streaming_method(self._stub)
+            spans = self.memory_exporter.get_finished_spans()
+        finally:
+            context.detach(token)
+        self.assertEqual(len(spans), 0)
+
+    def test_stream_unary_with_suppress_key(self):
+        token = context.attach(
+            context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)
+        )
+        try:
+            client_streaming_method(self._stub)
+            spans = self.memory_exporter.get_finished_spans()
+        finally:
+            context.detach(token)
+        self.assertEqual(len(spans), 0)
+
+    def test_stream_stream_with_suppress_key(self):
+        token = context.attach(
+            context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)
+        )
+        try:
+            bidirectional_streaming_method(self._stub)
+            spans = self.memory_exporter.get_finished_spans()
+        finally:
+            context.detach(token)
+        self.assertEqual(len(spans), 0)