diff --git a/ext/opentelemetry-ext-docker-tests/tests/pymongo/test_pymongo_functional.py b/ext/opentelemetry-ext-docker-tests/tests/pymongo/test_pymongo_functional.py index 1b2e2694829..f85a32d450e 100644 --- a/ext/opentelemetry-ext-docker-tests/tests/pymongo/test_pymongo_functional.py +++ b/ext/opentelemetry-ext-docker-tests/tests/pymongo/test_pymongo_functional.py @@ -17,7 +17,7 @@ from pymongo import MongoClient from opentelemetry import trace as trace_api -from opentelemetry.ext.pymongo import trace_integration +from opentelemetry.ext.pymongo import PymongoInstrumentor from opentelemetry.test.test_base import TestBase MONGODB_HOST = os.getenv("MONGODB_HOST ", "localhost") @@ -31,7 +31,7 @@ class TestFunctionalPymongo(TestBase): def setUpClass(cls): super().setUpClass() cls._tracer = cls.tracer_provider.get_tracer(__name__) - trace_integration(cls.tracer_provider) + PymongoInstrumentor().instrument() client = MongoClient( MONGODB_HOST, MONGODB_PORT, serverSelectionTimeoutMS=2000 ) @@ -94,3 +94,23 @@ def test_delete(self): with self._tracer.start_as_current_span("rootSpan"): self._collection.delete_one({"name": "testName"}) self.validate_spans() + + def test_uninstrument(self): + # check that integration is working + self._collection.find_one() + spans = self.memory_exporter.get_finished_spans() + self.memory_exporter.clear() + self.assertEqual(len(spans), 1) + + # uninstrument and check not new spans are created + PymongoInstrumentor().uninstrument() + self._collection.find_one() + spans = self.memory_exporter.get_finished_spans() + self.memory_exporter.clear() + self.assertEqual(len(spans), 0) + + # re-enable and check that it works again + PymongoInstrumentor().instrument() + self._collection.find_one() + spans = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans), 1) diff --git a/ext/opentelemetry-ext-pymongo/CHANGELOG.md b/ext/opentelemetry-ext-pymongo/CHANGELOG.md index d947fff944a..c8661e548d1 100644 --- a/ext/opentelemetry-ext-pymongo/CHANGELOG.md +++ b/ext/opentelemetry-ext-pymongo/CHANGELOG.md @@ -2,6 +2,9 @@ ## Unreleased +- Implement instrumentor interface ([#612](https://github.com/open-telemetry/opentelemetry-python/pull/612)) + + ## 0.4a0 Released 2020-02-21 diff --git a/ext/opentelemetry-ext-pymongo/setup.cfg b/ext/opentelemetry-ext-pymongo/setup.cfg index 790ea8bff2b..5c65a395511 100644 --- a/ext/opentelemetry-ext-pymongo/setup.cfg +++ b/ext/opentelemetry-ext-pymongo/setup.cfg @@ -41,6 +41,7 @@ package_dir= packages=find_namespace: install_requires = opentelemetry-api == 0.7.dev0 + opentelemetry-auto-instrumentation == 0.7.dev0 pymongo ~= 3.1 [options.extras_require] @@ -49,3 +50,7 @@ test = [options.packages.find] where = src + +[options.entry_points] +opentelemetry_instrumentor = + pymongo = opentelemetry.ext.pymongo:PymongoInstrumentor diff --git a/ext/opentelemetry-ext-pymongo/src/opentelemetry/ext/pymongo/__init__.py b/ext/opentelemetry-ext-pymongo/src/opentelemetry/ext/pymongo/__init__.py index b85bf423795..ee61f6d4f97 100644 --- a/ext/opentelemetry-ext-pymongo/src/opentelemetry/ext/pymongo/__init__.py +++ b/ext/opentelemetry-ext-pymongo/src/opentelemetry/ext/pymongo/__init__.py @@ -13,8 +13,8 @@ # limitations under the License. """ -The integration with MongoDB supports the `pymongo`_ library and is specified -to ``trace_integration`` using ``'pymongo'``. +The integration with MongoDB supports the `pymongo`_ library, it can be +enabled using the ``PymongoInstrumentor``. .. _pymongo: https://pypi.org/project/pymongo @@ -26,11 +26,11 @@ from pymongo import MongoClient from opentelemetry import trace from opentelemetry.trace import TracerProvider - from opentelemetry.trace.ext.pymongo import trace_integration + from opentelemetry.trace.ext.pymongo import PymongoInstrumentor trace.set_tracer_provider(TracerProvider()) - trace_integration() + PymongoInstrumentor().instrument() client = MongoClient() db = client["MongoDB_Database"] collection = db["MongoDB_Collection"] @@ -42,6 +42,8 @@ from pymongo import monitoring +from opentelemetry import trace +from opentelemetry.auto_instrumentation.instrumentor import BaseInstrumentor from opentelemetry.ext.pymongo.version import __version__ from opentelemetry.trace import SpanKind, get_tracer from opentelemetry.trace.status import Status, StatusCanonicalCode @@ -50,27 +52,16 @@ COMMAND_ATTRIBUTES = ["filter", "sort", "skip", "limit", "pipeline"] -def trace_integration(tracer_provider=None): - """Integrate with pymongo to trace it using event listener. - https://api.mongodb.com/python/current/api/pymongo/monitoring.html - - Args: - tracer_provider: The `TracerProvider` to use. If none is passed the - current configured one is used. - """ - - tracer = get_tracer(__name__, __version__, tracer_provider) - - monitoring.register(CommandTracer(tracer)) - - class CommandTracer(monitoring.CommandListener): def __init__(self, tracer): self._tracer = tracer self._span_dict = {} + self.is_enabled = True def started(self, event: monitoring.CommandStartedEvent): """ Method to handle a pymongo CommandStartedEvent """ + if not self.is_enabled: + return command = event.command.get(event.command_name, "") name = DATABASE_TYPE + "." + event.command_name statement = event.command_name @@ -103,38 +94,70 @@ def started(self, event: monitoring.CommandStartedEvent): if span is not None: span.set_status(Status(StatusCanonicalCode.INTERNAL, str(ex))) span.end() - self._remove_span(event) + self._pop_span(event) def succeeded(self, event: monitoring.CommandSucceededEvent): """ Method to handle a pymongo CommandSucceededEvent """ - span = self._get_span(event) - if span is not None: - span.set_attribute( - "db.mongo.duration_micros", event.duration_micros - ) - span.set_status(Status(StatusCanonicalCode.OK, event.reply)) - span.end() - self._remove_span(event) + if not self.is_enabled: + return + span = self._pop_span(event) + if span is None: + return + span.set_attribute("db.mongo.duration_micros", event.duration_micros) + span.set_status(Status(StatusCanonicalCode.OK, event.reply)) + span.end() def failed(self, event: monitoring.CommandFailedEvent): """ Method to handle a pymongo CommandFailedEvent """ - span = self._get_span(event) - if span is not None: - span.set_attribute( - "db.mongo.duration_micros", event.duration_micros - ) - span.set_status(Status(StatusCanonicalCode.UNKNOWN, event.failure)) - span.end() - self._remove_span(event) - - def _get_span(self, event): - return self._span_dict.get(_get_span_dict_key(event)) + if not self.is_enabled: + return + span = self._pop_span(event) + if span is None: + return + span.set_attribute("db.mongo.duration_micros", event.duration_micros) + span.set_status(Status(StatusCanonicalCode.UNKNOWN, event.failure)) + span.end() - def _remove_span(self, event): - self._span_dict.pop(_get_span_dict_key(event)) + def _pop_span(self, event): + return self._span_dict.pop(_get_span_dict_key(event), None) def _get_span_dict_key(event): if event.connection_id is not None: return (event.request_id, event.connection_id) return event.request_id + + +class PymongoInstrumentor(BaseInstrumentor): + _commandtracer_instance = None # type CommandTracer + # The instrumentation for PyMongo is based on the event listener interface + # https://api.mongodb.com/python/current/api/pymongo/monitoring.html. + # This interface only allows to register listeners and does not provide + # an unregister API. In order to provide a mechanishm to disable + # instrumentation an enabled flag is implemented in CommandTracer, + # it's checked in the different listeners. + + def _instrument(self, **kwargs): + """Integrate with pymongo to trace it using event listener. + https://api.mongodb.com/python/current/api/pymongo/monitoring.html + + Args: + tracer_provider: The `TracerProvider` to use. If none is passed the + current configured one is used. + """ + + tracer_provider = kwargs.get("tracer_provider") + + # Create and register a CommandTracer only the first time + if self._commandtracer_instance is None: + tracer = get_tracer(__name__, __version__, tracer_provider) + + self._commandtracer_instance = CommandTracer(tracer) + monitoring.register(self._commandtracer_instance) + + # If already created, just enable it + self._commandtracer_instance.is_enabled = True + + def _uninstrument(self, **kwargs): + if self._commandtracer_instance is not None: + self._commandtracer_instance.is_enabled = False diff --git a/ext/opentelemetry-ext-pymongo/tests/test_pymongo.py b/ext/opentelemetry-ext-pymongo/tests/test_pymongo.py index d85b23bfc95..abb4e8ab50a 100644 --- a/ext/opentelemetry-ext-pymongo/tests/test_pymongo.py +++ b/ext/opentelemetry-ext-pymongo/tests/test_pymongo.py @@ -15,7 +15,7 @@ from unittest import mock from opentelemetry import trace as trace_api -from opentelemetry.ext.pymongo import CommandTracer, trace_integration +from opentelemetry.ext.pymongo import CommandTracer, PymongoInstrumentor from opentelemetry.test.test_base import TestBase @@ -24,13 +24,13 @@ def setUp(self): super().setUp() self.tracer = self.tracer_provider.get_tracer(__name__) - def test_trace_integration(self): + def test_pymongo_instrumentor(self): mock_register = mock.Mock() patch = mock.patch( "pymongo.monitoring.register", side_effect=mock_register ) with patch: - trace_integration(self.tracer_provider) + PymongoInstrumentor().instrument() self.assertTrue(mock_register.called) @@ -50,7 +50,7 @@ def test_started(self): # the memory exporter can't be used here because the span isn't ended # yet # pylint: disable=protected-access - span = command_tracer._get_span(mock_event) + span = command_tracer._pop_span(mock_event) self.assertIs(span.kind, trace_api.SpanKind.CLIENT) self.assertEqual(span.name, "mongodb.command_name.find") self.assertEqual(span.attributes["component"], "mongodb") diff --git a/tox.ini b/tox.ini index 198ca773ef3..ef8b6334032 100644 --- a/tox.ini +++ b/tox.ini @@ -177,6 +177,7 @@ commands_pre = prometheus: pip install {toxinidir}/ext/opentelemetry-ext-prometheus + pymongo: pip install {toxinidir}/opentelemetry-auto-instrumentation pymongo: pip install {toxinidir}/ext/opentelemetry-ext-pymongo[test] psycopg2: pip install {toxinidir}/ext/opentelemetry-ext-dbapi