(instrumentation-sqlite3): trace connections made with `dbapi2.connec…

…t` (#873) * Change wrap_connect method * Wrap both connect defs * Add tests * Add CHANGELOG entry
open-telemetry · Jan 26, 2022 · 571c378 · 571c378
1 parent ec5b73c
commit 571c378
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -34,6 +34,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `opentelemetry-instrumentation-django` Django: fix issue preventing detection of MIDDLEWARE_CLASSES
 
+- `opentelemetry-instrumentation-sqlite3` Instrumentation now works with `dbapi2.connect`
+
 ## [1.8.0-0.27b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.8.0-0.27b0) - 2021-12-17
 
 ### Added

diff --git a/...entelemetry-instrumentation-sqlite3/src/opentelemetry/instrumentation/sqlite3/__init__.py b/...entelemetry-instrumentation-sqlite3/src/opentelemetry/instrumentation/sqlite3/__init__.py
@@ -40,6 +40,7 @@
 """
 
 import sqlite3
+from sqlite3 import dbapi2
 from typing import Collection
 
 from opentelemetry.instrumentation import dbapi
@@ -54,6 +55,8 @@
 
 
 class SQLite3Instrumentor(BaseInstrumentor):
+    _TO_WRAP = [sqlite3, dbapi2]
+
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
@@ -63,19 +66,21 @@ def _instrument(self, **kwargs):
         """
         tracer_provider = kwargs.get("tracer_provider")
 
-        dbapi.wrap_connect(
-            __name__,
-            sqlite3,
-            "connect",
-            _DATABASE_SYSTEM,
-            _CONNECTION_ATTRIBUTES,
-            version=__version__,
-            tracer_provider=tracer_provider,
-        )
+        for module in self._TO_WRAP:
+            dbapi.wrap_connect(
+                __name__,
+                module,
+                "connect",
+                _DATABASE_SYSTEM,
+                _CONNECTION_ATTRIBUTES,
+                version=__version__,
+                tracer_provider=tracer_provider,
+            )
 
     def _uninstrument(self, **kwargs):
         """ "Disable SQLite3 instrumentation"""
-        dbapi.unwrap_connect(sqlite3, "connect")
+        for module in self._TO_WRAP:
+            dbapi.unwrap_connect(module, "connect")
 
     @staticmethod
     def instrument_connection(connection, tracer_provider=None):

diff --git a/instrumentation/opentelemetry-instrumentation-sqlite3/tests/test_sqlite3.py b/instrumentation/opentelemetry-instrumentation-sqlite3/tests/test_sqlite3.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import sqlite3
+from sqlite3 import dbapi2
 
 from opentelemetry import trace as trace_api
 from opentelemetry.instrumentation.sqlite3 import SQLite3Instrumentor
@@ -25,20 +26,29 @@ def setUpClass(cls):
         super().setUpClass()
         cls._connection = None
         cls._cursor = None
+        cls._connection2 = None
+        cls._cursor2 = None
         cls._tracer = cls.tracer_provider.get_tracer(__name__)
         SQLite3Instrumentor().instrument(tracer_provider=cls.tracer_provider)
         cls._connection = sqlite3.connect(":memory:")
         cls._cursor = cls._connection.cursor()
+        cls._connection2 = dbapi2.connect(":memory:")
+        cls._cursor2 = cls._connection2.cursor()
 
     @classmethod
     def tearDownClass(cls):
         if cls._cursor:
             cls._cursor.close()
         if cls._connection:
             cls._connection.close()
+        if cls._cursor2:
+            cls._cursor2.close()
+        if cls._connection2:
+            cls._connection2.close()
 
     def validate_spans(self, span_name):
         spans = self.memory_exporter.get_finished_spans()
+        self.memory_exporter.clear()
         self.assertEqual(len(spans), 2)
         for span in spans:
             if span.name == "rootSpan":
@@ -62,14 +72,22 @@ def test_execute(self):
             self._cursor.execute(stmt)
         self.validate_spans("CREATE")
 
+        with self._tracer.start_as_current_span("rootSpan"):
+            self._cursor2.execute(stmt)
+        self.validate_spans("CREATE")
+
     def test_executemany(self):
         """Should create a child span for executemany"""
         stmt = "INSERT INTO test (id) VALUES (?)"
+        data = [("1",), ("2",), ("3",)]
         with self._tracer.start_as_current_span("rootSpan"):
-            data = [("1",), ("2",), ("3",)]
             self._cursor.executemany(stmt, data)
         self.validate_spans("INSERT")
 
+        with self._tracer.start_as_current_span("rootSpan"):
+            self._cursor2.executemany(stmt, data)
+        self.validate_spans("INSERT")
+
     def test_callproc(self):
         """Should create a child span for callproc"""
         with self._tracer.start_as_current_span("rootSpan"), self.assertRaises(