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 <srikanth.chekuri92@gmail.com> Co-authored-by: Shalev Roda <65566801+shalevr@users.noreply.github.com>
open-telemetry · Feb 20, 2023 · 3bcc043 · 3bcc043
1 parent d5cbd67
commit 3bcc043
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 5 deletions.
diff --git 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/...lemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py b/...lemetry-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
@@ -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
@@ -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
@@ -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")