Add support to instrument multiple sqlalchemy engines (#1132)

* Override __new__ to instrument multiple engines * add possibility to instrument mutiple engines * update changelog * update return type in doc * fix CHANGELOG * format * fix test * Check if engines isinstance Sequence Co-authored-by: Diego Hurtado <[email protected]>
open-telemetry · Jun 16, 2022 · f03bef2 · f03bef2
1 parent 4ba6214
commit f03bef2
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -23,6 +23,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 
 ### Added
+- `opentelemetry-instrumentation-sqlalchemy` add support to instrument multiple engines
+  ([#1132](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1132))
 - `opentelemetry-instrumentation-logging` add log hook support
   ([#1117](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1117))
 - `opentelemetry-instrumentation-remoulade` Initial release

diff --git a/...metry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py b/...metry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py
@@ -51,6 +51,7 @@
 API
 ---
 """
+from collections.abc import Sequence
 from typing import Collection
 
 import sqlalchemy
@@ -83,10 +84,11 @@ def _instrument(self, **kwargs):
         Args:
             **kwargs: Optional arguments
                 ``engine``: a SQLAlchemy engine instance
+                ``engines``: a list of SQLAlchemy engine instances
                 ``tracer_provider``: a TracerProvider, defaults to global
 
         Returns:
-            An instrumented engine if passed in as an argument, None otherwise.
+            An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise.
         """
         tracer_provider = kwargs.get("tracer_provider")
         _w("sqlalchemy", "create_engine", _wrap_create_engine(tracer_provider))
@@ -108,6 +110,18 @@ def _instrument(self, **kwargs):
                 kwargs.get("engine"),
                 kwargs.get("enable_commenter", False),
             )
+        if kwargs.get("engines") is not None and isinstance(
+            kwargs.get("engines"), Sequence
+        ):
+            return [
+                EngineTracer(
+                    _get_tracer(tracer_provider),
+                    engine,
+                    kwargs.get("enable_commenter", False),
+                )
+                for engine in kwargs.get("engines")
+            ]
+
         return None
 
     def _uninstrument(self, **kwargs):

diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py
@@ -50,6 +50,24 @@ def test_trace_integration(self):
         self.assertEqual(spans[0].name, "SELECT :memory:")
         self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
 
+    def test_instrument_two_engines(self):
+        engine_1 = create_engine("sqlite:///:memory:")
+        engine_2 = create_engine("sqlite:///:memory:")
+
+        SQLAlchemyInstrumentor().instrument(
+            engines=[engine_1, engine_2],
+            tracer_provider=self.tracer_provider,
+        )
+
+        cnx_1 = engine_1.connect()
+        cnx_1.execute("SELECT	1 + 1;").fetchall()
+        cnx_2 = engine_2.connect()
+        cnx_2.execute("SELECT	1 + 1;").fetchall()
+
+        spans = self.memory_exporter.get_finished_spans()
+
+        self.assertEqual(len(spans), 2)
+
     @pytest.mark.skipif(
         not sqlalchemy.__version__.startswith("1.4"),
         reason="only run async tests for 1.4",