Implement checking for suppression key in pymongo instrumentation (open-telemetry#736)

* Implement checking for suppression key


Add test for suppression key with pymongo


changelog update

* Linting

* pylint

* Bad reformat

* Update changelog with PR

Co-authored-by: Diego Hurtado <[email protected]>

Author: andrew-matteson

CHANGELOG.md

@@ -26,6 +26,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 - `opentelemetry-instrumentation-jinja2` Allow instrumentation of newer Jinja2 versions.
 
+### Changed
+- `opentelemetry-instrumentation-pymongo` Add check for suppression key in PyMongo.
+- ([#736](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/736))
+
 ### Added
 - `opentelemetry-instrumentation-elasticsearch` Added `response_hook` and `request_hook` callbacks
   ([#670](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/670))

instrumentation/opentelemetry-instrumentation-pymongo/src/opentelemetry/instrumentation/pymongo/__init__.py

@@ -41,9 +41,11 @@
 
 from pymongo import monitoring
 
+from opentelemetry import context
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.pymongo.package import _instruments
 from opentelemetry.instrumentation.pymongo.version import __version__
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.semconv.trace import DbSystemValues, SpanAttributes
 from opentelemetry.trace import SpanKind, get_tracer
 from opentelemetry.trace.status import Status, StatusCode
@@ -57,7 +59,9 @@ def __init__(self, tracer):
 
     def started(self, event: monitoring.CommandStartedEvent):
         """ Method to handle a pymongo CommandStartedEvent """
-        if not self.is_enabled:
+        if not self.is_enabled or context.get_value(
+            _SUPPRESS_INSTRUMENTATION_KEY
+        ):
             return
         command = event.command.get(event.command_name, "")
         name = event.command_name
@@ -92,7 +96,9 @@ def started(self, event: monitoring.CommandStartedEvent):
 
     def succeeded(self, event: monitoring.CommandSucceededEvent):
         """ Method to handle a pymongo CommandSucceededEvent """
-        if not self.is_enabled:
+        if not self.is_enabled or context.get_value(
+            _SUPPRESS_INSTRUMENTATION_KEY
+        ):
             return
         span = self._pop_span(event)
         if span is None:
@@ -101,7 +107,9 @@ def succeeded(self, event: monitoring.CommandSucceededEvent):
 
     def failed(self, event: monitoring.CommandFailedEvent):
         """ Method to handle a pymongo CommandFailedEvent """
-        if not self.is_enabled:
+        if not self.is_enabled or context.get_value(
+            _SUPPRESS_INSTRUMENTATION_KEY
+        ):
             return
         span = self._pop_span(event)
         if span is None:

instrumentation/opentelemetry-instrumentation-pymongo/tests/test_pymongo.py

@@ -14,11 +14,13 @@
 
 from unittest import mock
 
+from opentelemetry import context
 from opentelemetry import trace as trace_api
 from opentelemetry.instrumentation.pymongo import (
     CommandTracer,
     PymongoInstrumentor,
 )
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.test.test_base import TestBase
 
@@ -90,6 +92,31 @@ def test_not_recording(self):
         self.assertFalse(mock_span.set_attribute.called)
         self.assertFalse(mock_span.set_status.called)
 
+    def test_suppression_key(self):
+        mock_tracer = mock.Mock()
+        mock_span = mock.Mock()
+        mock_span.is_recording.return_value = True
+        mock_tracer.start_span.return_value = mock_span
+        mock_event = MockEvent({})
+        mock_event.command.get = mock.Mock()
+        mock_event.command.get.return_value = "dummy"
+
+        token = context.attach(
+            context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)
+        )
+
+        try:
+            command_tracer = CommandTracer(mock_tracer)
+            command_tracer.started(event=mock_event)
+            command_tracer.succeeded(event=mock_event)
+        finally:
+            context.detach(token)
+
+        # if suppression key is set, CommandTracer methods return immediately, so command.get is not invoked.
+        self.assertFalse(
+            mock_event.command.get.called  # pylint: disable=no-member
+        )
+
     def test_failed(self):
         mock_event = MockEvent({})
         command_tracer = CommandTracer(self.tracer)