Fix sqlalchemy uninstrument (#1581)

open-telemetry · Jan 29, 2023 · bbe7578 · bbe7578
1 parent a300d65
commit bbe7578
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 4 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fix SQLAlchemy uninstrumentation
+  ([#1581](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1581))
 - `opentelemetry-instrumentation-grpc` Fix code()/details() of _OpentelemetryServicerContext.
   ([#1578](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1578))
 - Fix aiopg instrumentation to work with aiopg < 2.0.0

diff --git a/...metry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py b/...metry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py
@@ -186,3 +186,4 @@ def _uninstrument(self, **kwargs):
         unwrap(Engine, "connect")
         if parse_version(sqlalchemy.__version__).release >= (1, 4):
             unwrap(sqlalchemy.ext.asyncio, "create_async_engine")
+        EngineTracer.remove_all_event_listeners()
diff --git a/...lemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py b/...lemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py
@@ -14,7 +14,10 @@
 import os
 import re
 
-from sqlalchemy.event import listen  # pylint: disable=no-name-in-module
+from sqlalchemy.event import (  # pylint: disable=no-name-in-module
+    listen,
+    remove,
+)
 
 from opentelemetry import trace
 from opentelemetry.instrumentation.sqlalchemy.package import (
@@ -95,6 +98,8 @@ def _wrap_connect_internal(func, module, args, kwargs):
 
 
 class EngineTracer:
+    _remove_event_listener_params = []
+
     def __init__(
         self, tracer, engine, enable_commenter=False, commenter_options=None
     ):
@@ -105,11 +110,24 @@ def __init__(
         self.commenter_options = commenter_options if commenter_options else {}
         self._leading_comment_remover = re.compile(r"^/\*.*?\*/")
 
-        listen(
+        self._register_event_listener(
             engine, "before_cursor_execute", self._before_cur_exec, retval=True
         )
-        listen(engine, "after_cursor_execute", _after_cur_exec)
-        listen(engine, "handle_error", _handle_error)
+        self._register_event_listener(
+            engine, "after_cursor_execute", _after_cur_exec
+        )
+        self._register_event_listener(engine, "handle_error", _handle_error)
+
+    @classmethod
+    def _register_event_listener(cls, target, identifier, func, *args, **kw):
+        listen(target, identifier, func, *args, **kw)
+        cls._remove_event_listener_params.append((target, identifier, func))
+
+    @classmethod
+    def remove_all_event_listeners(cls):
+        for remove_params in cls._remove_event_listener_params:
+            remove(*remove_params)
+        cls._remove_event_listener_params.clear()
 
     def _operation_name(self, db_name, statement):
         parts = []

diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py
@@ -248,12 +248,41 @@ def test_uninstrument(self):
 
         self.memory_exporter.clear()
         SQLAlchemyInstrumentor().uninstrument()
+        cnx.execute("SELECT	1 + 1;").fetchall()
         engine2 = create_engine("sqlite:///:memory:")
         cnx2 = engine2.connect()
         cnx2.execute("SELECT	2 + 2;").fetchall()
         spans = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans), 0)
 
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+        )
+        cnx = engine.connect()
+        cnx.execute("SELECT	1 + 1;").fetchall()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 2)
+
+    def test_uninstrument_without_engine(self):
+        SQLAlchemyInstrumentor().instrument(
+            tracer_provider=self.tracer_provider
+        )
+        from sqlalchemy import create_engine
+
+        engine = create_engine("sqlite:///:memory:")
+
+        cnx = engine.connect()
+        cnx.execute("SELECT	1 + 1;").fetchall()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 2)
+
+        self.memory_exporter.clear()
+        SQLAlchemyInstrumentor().uninstrument()
+        cnx.execute("SELECT	1 + 1;").fetchall()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 0)
+
     def test_no_op_tracer_provider(self):
         engine = create_engine("sqlite:///:memory:")
         SQLAlchemyInstrumentor().instrument(