Add span for connection phase

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

@@ -55,6 +55,7 @@
 from typing import Collection
 
 import sqlalchemy
+from sqlalchemy.engine.base import Engine
 from packaging.version import parse as parse_version
 from wrapt import wrap_function_wrapper as _w
 
@@ -64,6 +65,7 @@
     _get_tracer,
     _wrap_create_async_engine,
     _wrap_create_engine,
+    _wrap_connect,
 )
 from opentelemetry.instrumentation.sqlalchemy.package import _instruments
 from opentelemetry.instrumentation.utils import unwrap
@@ -97,13 +99,17 @@ def _instrument(self, **kwargs):
             "create_engine",
             _wrap_create_engine(tracer_provider),
         )
+        _w(
+            "sqlalchemy.engine.base",
+            "Engine.connect",
+            _wrap_connect(tracer_provider)
+        )
         if parse_version(sqlalchemy.__version__).release >= (1, 4):
             _w(
                 "sqlalchemy.ext.asyncio",
                 "create_async_engine",
                 _wrap_create_async_engine(tracer_provider),
             )
-
         if kwargs.get("engine") is not None:
             return EngineTracer(
                 _get_tracer(tracer_provider),
@@ -127,5 +133,7 @@ def _instrument(self, **kwargs):
     def _uninstrument(self, **kwargs):
         unwrap(sqlalchemy, "create_engine")
         unwrap(sqlalchemy.engine, "create_engine")
+        unwrap(Engine, "connect")
         if parse_version(sqlalchemy.__version__).release >= (1, 4):
             unwrap(sqlalchemy.ext.asyncio, "create_async_engine")
+

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

@@ -77,6 +77,17 @@ def _wrap_create_engine_internal(func, module, args, kwargs):
     return _wrap_create_engine_internal
 
 
+def _wrap_connect(tracer_provider=None):
+    def _wrap_connect_internal(func, module, args, kwargs):
+        tracer = trace.get_tracer(
+            module.name,
+            __version__,
+            tracer_provider=tracer_provider,
+        )
+        with tracer.start_as_current_span("database-connect", kind=trace.SpanKind.CLIENT) as span:
+            return func(*args, **kwargs)
+    return _wrap_connect_internal
+
 class EngineTracer:
     def __init__(self, tracer, engine, enable_commenter=False):
         self.tracer = tracer

instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py

@@ -46,9 +46,13 @@ def test_trace_integration(self):
         cnx.execute("SELECT	1 + 1;").fetchall()
         spans = self.memory_exporter.get_finished_spans()
 
-        self.assertEqual(len(spans), 1)
-        self.assertEqual(spans[0].name, "SELECT :memory:")
+        self.assertEqual(len(spans), 2)
+        # first span - the connection to the db
+        self.assertEqual(spans[0].name, "database-connect")
         self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
+        # second span - the query it self
+        self.assertEqual(spans[1].name, "SELECT :memory:")
+        self.assertEqual(spans[1].kind, trace.SpanKind.CLIENT)
 
     def test_instrument_two_engines(self):
         engine_1 = create_engine("sqlite:///:memory:")
@@ -65,8 +69,20 @@ def test_instrument_two_engines(self):
         cnx_2.execute("SELECT	1 + 1;").fetchall()
 
         spans = self.memory_exporter.get_finished_spans()
+        # 2 queries + 2 engine connect
+        self.assertEqual(len(spans), 4)
 
-        self.assertEqual(len(spans), 2)
+    def test_instrument_engine_connect(self):
+        engine = create_engine("sqlite:///:memory:")
+
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+        )
+
+        engine.connect()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
 
     @pytest.mark.skipif(
         not sqlalchemy.__version__.startswith("1.4"),
@@ -85,11 +101,15 @@ async def run():
             async with engine.connect() as cnx:
                 await cnx.execute(sqlalchemy.text("SELECT	1 + 1;"))
             spans = self.memory_exporter.get_finished_spans()
-            self.assertEqual(len(spans), 1)
-            self.assertEqual(spans[0].name, "SELECT :memory:")
+            self.assertEqual(len(spans), 2)
+            # first span - the connection to the db
+            self.assertEqual(spans[0].name, "database-connect")
             self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
+            # second span - the query
+            self.assertEqual(spans[1].name, "SELECT :memory:")
+            self.assertEqual(spans[1].kind, trace.SpanKind.CLIENT)
             self.assertEqual(
-                spans[0].instrumentation_scope.name,
+                spans[1].instrumentation_scope.name,
                 "opentelemetry.instrumentation.sqlalchemy",
             )
 
@@ -99,7 +119,10 @@ def test_not_recording(self):
         mock_tracer = mock.Mock()
         mock_span = 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
         with mock.patch("opentelemetry.trace.get_tracer") as tracer:
             tracer.return_value = mock_tracer
             engine = create_engine("sqlite:///:memory:")
@@ -123,11 +146,15 @@ def test_create_engine_wrapper(self):
         cnx.execute("SELECT	1 + 1;").fetchall()
         spans = self.memory_exporter.get_finished_spans()
 
-        self.assertEqual(len(spans), 1)
-        self.assertEqual(spans[0].name, "SELECT :memory:")
+        self.assertEqual(len(spans), 2)
+        # first span - the connection to the db
+        self.assertEqual(spans[0].name, "database-connect")
         self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
+        # second span - the query
+        self.assertEqual(spans[1].name, "SELECT :memory:")
+        self.assertEqual(spans[1].kind, trace.SpanKind.CLIENT)
         self.assertEqual(
-            spans[0].instrumentation_scope.name,
+            spans[1].instrumentation_scope.name,
             "opentelemetry.instrumentation.sqlalchemy",
         )
 
@@ -153,7 +180,7 @@ def test_custom_tracer_provider(self):
         cnx.execute("SELECT	1 + 1;").fetchall()
         spans = self.memory_exporter.get_finished_spans()
 
-        self.assertEqual(len(spans), 1)
+        self.assertEqual(len(spans), 2)
         self.assertEqual(spans[0].resource.attributes["service.name"], "test")
         self.assertEqual(
             spans[0].resource.attributes["deployment.environment"], "env"
@@ -177,11 +204,15 @@ async def run():
             async with engine.connect() as cnx:
                 await cnx.execute(sqlalchemy.text("SELECT	1 + 1;"))
             spans = self.memory_exporter.get_finished_spans()
-            self.assertEqual(len(spans), 1)
-            self.assertEqual(spans[0].name, "SELECT :memory:")
+            self.assertEqual(len(spans), 2)
+            # first span - the connection to the db
+            self.assertEqual(spans[0].name, "database-connect")
             self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
+            # second span - the query
+            self.assertEqual(spans[1].name, "SELECT :memory:")
+            self.assertEqual(spans[1].kind, trace.SpanKind.CLIENT)
             self.assertEqual(
-                spans[0].instrumentation_scope.name,
+                spans[1].instrumentation_scope.name,
                 "opentelemetry.instrumentation.sqlalchemy",
             )
 
@@ -199,8 +230,8 @@ def test_generate_commenter(self):
         cnx = engine.connect()
         cnx.execute("SELECT	1 + 1;").fetchall()
         spans = self.memory_exporter.get_finished_spans()
-        self.assertEqual(len(spans), 1)
-        span = spans[0]
+        self.assertEqual(len(spans), 2)
+        span = spans[1]
         self.assertIn(
             EngineTracer._generate_comment(span),
             self.caplog.records[-2].getMessage(),