Add connection attributes to sqlalchemy connect span

Author: bschoenmaeckers

CHANGELOG.md

@@ -26,6 +26,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#1592](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1592))
 - `opentelemetry-instrumentation-django` Allow explicit `excluded_urls` configuration through `instrument()`
   ([#1618](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1618))
+- Add connection attributes to sqlalchemy connect span
+  ([#1608](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1608))
 
 ### Fixed
 

instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py

@@ -91,7 +91,11 @@ 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

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,8 @@ 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 +221,8 @@ 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:")

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(

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")

tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_postgres.py

@@ -64,6 +64,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")