From 6c655d36bf8d7cab84da6896b530911e7d8263fd Mon Sep 17 00:00:00 2001 From: Bas Schoenmaeckers <7943856+bschoenmaeckers@users.noreply.github.com> Date: Mon, 20 Feb 2023 18:17:18 +0100 Subject: [PATCH] Add connection attributes to sqlalchemy connect span (#1608) * Add connection attributes to sqlalchemy connect span * Fix formatting * Move changelog entry to unreleased section --------- Co-authored-by: Srikanth Chekuri Co-authored-by: Shalev Roda <65566801+shalevr@users.noreply.github.com> --- CHANGELOG.md | 3 +++ .../instrumentation/sqlalchemy/engine.py | 8 ++++++- .../tests/test_sqlalchemy.py | 22 +++++++++++++++---- .../tests/sqlalchemy_tests/test_mssql.py | 2 ++ .../tests/sqlalchemy_tests/test_mysql.py | 1 + 5 files changed, 31 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0a6fe9c56e..d93f7bab1a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased +- Add connection attributes to sqlalchemy connect span + ([#1608](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1608)) + ## Version 1.16.0/0.37b0 (2023-02-17) ### Added diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py index f95751d9c6..2afbd7a82f 100644 --- a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py +++ b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py @@ -91,7 +91,13 @@ def _wrap_connect(tracer_provider=None): def _wrap_connect_internal(func, module, args, kwargs): with tracer.start_as_current_span( "connect", kind=trace.SpanKind.CLIENT - ): + ) as span: + if span.is_recording(): + attrs, _ = _get_attributes_from_url(module.url) + span.set_attributes(attrs) + span.set_attribute( + SpanAttributes.DB_SYSTEM, _normalize_vendor(module.name) + ) return func(*args, **kwargs) return _wrap_connect_internal diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py index e00148320c..981da107db 100644 --- a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py +++ b/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py @@ -22,6 +22,7 @@ from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import TracerProvider, export +from opentelemetry.semconv.trace import SpanAttributes from opentelemetry.test.test_base import TestBase @@ -128,11 +129,12 @@ async def run(): def test_not_recording(self): mock_tracer = mock.Mock() mock_span = mock.Mock() + mock_context = mock.Mock() mock_span.is_recording.return_value = False - mock_span.__enter__ = mock.Mock(return_value=(mock.Mock(), None)) - mock_span.__exit__ = mock.Mock(return_value=None) - mock_tracer.start_span.return_value = mock_span - mock_tracer.start_as_current_span.return_value = mock_span + mock_context.__enter__ = mock.Mock(return_value=mock_span) + mock_context.__exit__ = mock.Mock(return_value=None) + mock_tracer.start_span.return_value = mock_context + mock_tracer.start_as_current_span.return_value = mock_context with mock.patch("opentelemetry.trace.get_tracer") as tracer: tracer.return_value = mock_tracer engine = create_engine("sqlite:///:memory:") @@ -159,6 +161,12 @@ def test_create_engine_wrapper(self): self.assertEqual(len(spans), 2) # first span - the connection to the db self.assertEqual(spans[0].name, "connect") + self.assertEqual( + spans[0].attributes[SpanAttributes.DB_NAME], ":memory:" + ) + self.assertEqual( + spans[0].attributes[SpanAttributes.DB_SYSTEM], "sqlite" + ) self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT) # second span - the query self.assertEqual(spans[1].name, "SELECT :memory:") @@ -217,6 +225,12 @@ async def run(): self.assertEqual(len(spans), 2) # first span - the connection to the db self.assertEqual(spans[0].name, "connect") + self.assertEqual( + spans[0].attributes[SpanAttributes.DB_NAME], ":memory:" + ) + self.assertEqual( + spans[0].attributes[SpanAttributes.DB_SYSTEM], "sqlite" + ) self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT) # second span - the query self.assertEqual(spans[1].name, "SELECT :memory:") diff --git a/tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mssql.py b/tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mssql.py index fef1bf7b7b..713be24198 100644 --- a/tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mssql.py +++ b/tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mssql.py @@ -71,6 +71,7 @@ def test_engine_execute_errors(self): spans = self.memory_exporter.get_finished_spans() # one span for the connection and one for the query self.assertEqual(len(spans), 2) + self.check_meta(spans[0]) span = spans[1] # span fields self.assertEqual(span.name, "SELECT opentelemetry-tests") @@ -99,6 +100,7 @@ def test_orm_insert(self): spans = self.memory_exporter.get_finished_spans() # connect, identity insert on before the insert, insert, and identity insert off after the insert self.assertEqual(len(spans), 4) + self.check_meta(spans[0]) span = spans[2] self._check_span(span, "INSERT") self.assertIn( diff --git a/tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mysql.py b/tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mysql.py index ceebe78ef5..6e4c85cd92 100644 --- a/tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mysql.py +++ b/tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mysql.py @@ -70,6 +70,7 @@ def test_engine_execute_errors(self): spans = self.memory_exporter.get_finished_spans() # one span for the connection and one for the query self.assertEqual(len(spans), 2) + self.check_meta(spans[0]) span = spans[1] # span fields self.assertEqual(span.name, "SELECT opentelemetry-tests")