Change the code to save all listen params so we will know how to remove them when uninstument

Author: shalevr

Diff for: instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py

@@ -160,22 +160,17 @@ def _instrument(self, **kwargs):
                 "create_async_engine",
                 _wrap_create_async_engine(tracer_provider, enable_commenter),
             )
-
-        self.engines = []
         if kwargs.get("engine") is not None:
-            self.engines.append(
-                EngineTracer(
-                    _get_tracer(tracer_provider),
-                    kwargs.get("engine"),
-                    kwargs.get("enable_commenter", False),
-                    kwargs.get("commenter_options", {}),
-                )
+            return EngineTracer(
+                _get_tracer(tracer_provider),
+                kwargs.get("engine"),
+                kwargs.get("enable_commenter", False),
+                kwargs.get("commenter_options", {}),
             )
-            return self.engines[0]
         if kwargs.get("engines") is not None and isinstance(
             kwargs.get("engines"), Sequence
         ):
-            self.engines = [
+            return [
                 EngineTracer(
                     _get_tracer(tracer_provider),
                     engine,
@@ -184,7 +179,6 @@ def _instrument(self, **kwargs):
                 )
                 for engine in kwargs.get("engines")
             ]
-            return self.engines
 
         return None
 
@@ -194,5 +188,4 @@ def _uninstrument(self, **kwargs):
         unwrap(Engine, "connect")
         if parse_version(sqlalchemy.__version__).release >= (1, 4):
             unwrap(sqlalchemy.ext.asyncio, "create_async_engine")
-        for engine in self.engines:
-            engine.remove_event_listeners()
+        EngineTracer.remove_all_event_listeners()

Diff for: instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py

@@ -98,6 +98,8 @@ def _wrap_connect_internal(func, module, args, kwargs):
 
 
 class EngineTracer:
+    _removeEventListenerParams = []
+
     def __init__(
         self, tracer, engine, enable_commenter=False, commenter_options=None
     ):
@@ -108,11 +110,22 @@ 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, fn, *args, **kw):
+        listen(target, identifier, fn, *args, **kw)
+        cls._removeEventListenerParams.append((target, identifier, fn))
+
+    @classmethod
+    def remove_all_event_listeners(cls):
+        for removeParams in cls._removeEventListenerParams:
+            remove(*removeParams)
+        cls._removeEventListenerParams.clear()
 
     def remove_event_listeners(self):
         remove(self.engine, "before_cursor_execute", self._before_cur_exec)

Diff for: instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py

@@ -255,6 +255,22 @@ def test_uninstrument(self):
         spans = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans), 0)
 
+    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(