Rm instrument_connection is-instrumented check; use dict to store cnx

Author: tammy-baylis-swi

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

@@ -146,7 +146,7 @@ class Psycopg2Instrumentor(BaseInstrumentor):
 
     _DATABASE_SYSTEM = "postgresql"
 
-    _otel_cursor_factory = None
+    _otel_cursor_factories = {}
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
@@ -193,30 +193,25 @@ def instrument_connection(
         Returns:
             An instrumented psycopg2 connection object.
         """
-        if self._is_instrumented_by_opentelemetry:
-            # _instrument (via BaseInstrumentor) or instrument_connection (this)
-            # was already called
-            _logger.warning(
-                "Attempting to instrument Psycopg2 connection while already instrumented"
-            )
-            self._is_instrumented_by_opentelemetry = True
-            return connection
+        # TODO Add check for attempt to instrument a connection when already instrumented
+        #      https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3138
 
-        # Save cursor_factory at instrumentor level because
+        # Save cursor_factory in instrumentor map because
         # psycopg2 connection type does not allow arbitrary attrs
-        self._otel_cursor_factory = connection.cursor_factory
+        self._otel_cursor_factories[connection] = connection.cursor_factory
         connection.cursor_factory = _new_cursor_factory(
             base_factory=connection.cursor_factory,
             tracer_provider=tracer_provider,
         )
-        self._is_instrumented_by_opentelemetry = True
+        _logger.warning("factories: %s", self._otel_cursor_factories)
 
         return connection
 
     # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
     def uninstrument_connection(self, connection):
-        self._is_instrumented_by_opentelemetry = False
-        connection.cursor_factory = self._otel_cursor_factory
+        connection.cursor_factory = self._otel_cursor_factories.pop(
+            connection, None
+        )
         return connection
 
 

instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py

@@ -16,7 +16,6 @@
 from unittest import mock
 
 import psycopg2
-import pytest
 
 import opentelemetry.instrumentation.psycopg2
 from opentelemetry import trace
@@ -69,10 +68,6 @@ def get_dsn_parameters(self):  # pylint: disable=no-self-use
 
 
 class TestPostgresqlIntegration(TestBase):
-    @pytest.fixture(autouse=True)
-    def inject_fixtures(self, caplog):
-        self.caplog = caplog  # pylint: disable=attribute-defined-outside-init
-
     def setUp(self):
         super().setUp()
         self.cursor_mock = mock.patch(
@@ -200,9 +195,8 @@ def test_instrument_connection(self):
         instrumentor = Psycopg2Instrumentor()
         original_cursor_factory = cnx.cursor_factory
         cnx = instrumentor.instrument_connection(cnx)
-        self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
         self.assertEqual(
-            instrumentor._otel_cursor_factory, original_cursor_factory
+            instrumentor._otel_cursor_factories[cnx], original_cursor_factory
         )
         cursor = cnx.cursor()
         cursor.execute(query)
@@ -222,22 +216,27 @@ def test_instrument_connection_with_instrument(self):
 
         instrumentor = Psycopg2Instrumentor()
         instrumentor.instrument()
-        self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
 
         cnx = psycopg2.connect(database="test")
+        original_cursor_factory = cnx.cursor_factory
+        orig_cnx = cnx
         cnx = Psycopg2Instrumentor().instrument_connection(cnx)
+
+        # TODO Should not set cursor_factory if already instrumented
+        #      https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3138
+        # self.assertEqual(instrumentor._otel_cursor_factories[cnx], None)
         self.assertEqual(
-            self.caplog.records[0].getMessage(),
-            "Attempting to instrument Psycopg2 connection while already instrumented",
+            instrumentor._otel_cursor_factories.get(orig_cnx),
+            original_cursor_factory,
         )
-        self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
-        # Will not set cursor_factory if already instrumented
-        self.assertEqual(instrumentor._otel_cursor_factory, None)
         cursor = cnx.cursor()
         cursor.execute(query)
 
         spans_list = self.memory_exporter.get_finished_spans()
-        self.assertEqual(len(spans_list), 1)
+        # TODO Add check for attempt to instrument a connection when already instrumented
+        #      https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3138
+        # self.assertEqual(len(spans_list), 1)
+        self.assertEqual(len(spans_list), 2)
 
     def test_instrument_connection_with_instrument_connection(self):
         cnx = psycopg2.connect(database="test")
@@ -251,26 +250,29 @@ def test_instrument_connection_with_instrument_connection(self):
         cnx = psycopg2.connect(database="test")
         original_cursor_factory = cnx.cursor_factory
         instrumentor = Psycopg2Instrumentor()
+        orig_cnx = cnx
         cnx = instrumentor.instrument_connection(cnx)
-        self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
         self.assertEqual(
-            instrumentor._otel_cursor_factory, original_cursor_factory
+            instrumentor._otel_cursor_factories.get(orig_cnx),
+            original_cursor_factory,
         )
 
-        cnx = Psycopg2Instrumentor().instrument_connection(cnx)
-        self.assertEqual(
-            self.caplog.records[0].getMessage(),
-            "Attempting to instrument Psycopg2 connection while already instrumented",
-        )
-        self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
+        orig_cnx = cnx
+        original_cursor_factory = cnx.cursor_factory
+        instrumentor = Psycopg2Instrumentor()
+        cnx = instrumentor.instrument_connection(cnx)
         self.assertEqual(
-            instrumentor._otel_cursor_factory, original_cursor_factory
+            instrumentor._otel_cursor_factories.get(orig_cnx),
+            original_cursor_factory,
         )
         cursor = cnx.cursor()
         cursor.execute(query)
 
         spans_list = self.memory_exporter.get_finished_spans()
-        self.assertEqual(len(spans_list), 1)
+        # TODO Add check for attempt to instrument a connection when already instrumented
+        #      https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3138
+        # self.assertEqual(len(spans_list), 1)
+        self.assertEqual(len(spans_list), 2)
 
     # pylint: disable=unused-argument
     def test_uninstrument_connection_with_instrument(self):
@@ -283,15 +285,13 @@ def test_uninstrument_connection_with_instrument(self):
 
         spans_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans_list), 1)
-        self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
 
         cnx = Psycopg2Instrumentor().uninstrument_connection(cnx)
         cursor = cnx.cursor()
         cursor.execute(query)
 
         spans_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans_list), 1)
-        self.assertFalse(instrumentor._is_instrumented_by_opentelemetry)
 
     # pylint: disable=unused-argument
     def test_uninstrument_connection_with_instrument_connection(self):
@@ -305,19 +305,20 @@ def test_uninstrument_connection_with_instrument_connection(self):
 
         spans_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans_list), 1)
-        self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
         self.assertEqual(
-            instrumentor._otel_cursor_factory, original_cursor_factory
+            instrumentor._otel_cursor_factories[cnx], original_cursor_factory
         )
 
+        orig_cnx = cnx
         cnx = Psycopg2Instrumentor().uninstrument_connection(cnx)
         cursor = cnx.cursor()
         cursor.execute(query)
 
         spans_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans_list), 1)
-        self.assertFalse(instrumentor._is_instrumented_by_opentelemetry)
-        self.assertEqual(instrumentor._otel_cursor_factory, None)
+        self.assertEqual(
+            instrumentor._otel_cursor_factories.get(orig_cnx), None
+        )
 
     @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
     def test_sqlcommenter_enabled(self, event_mocked):