From f03bef257998c4109072316d21210ed168ce2a4b Mon Sep 17 00:00:00 2001 From: hemma Date: Thu, 16 Jun 2022 13:44:55 +0200 Subject: [PATCH] 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 --- CHANGELOG.md | 2 ++ .../instrumentation/sqlalchemy/__init__.py | 16 +++++++++++++++- .../tests/test_sqlalchemy.py | 18 ++++++++++++++++++ 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index eb7d5b1ba..d7ca4a32b 100644 --- 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/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py index 6aa60383f..0e1d7b266 100644 --- a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py @@ -51,6 +51,7 @@ Usage API --- """ +from collections.abc import Sequence from typing import Collection import sqlalchemy @@ -83,10 +84,11 @@ class SQLAlchemyInstrumentor(BaseInstrumentor): 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 @@ class SQLAlchemyInstrumentor(BaseInstrumentor): 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 index bf60e6e6a..3b3c6d735 100644 --- a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py +++ b/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py @@ -50,6 +50,24 @@ class TestSqlalchemyInstrumentation(TestBase): 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",