Fix psycopg2 instrument_connection AttributeError

CHANGELOG.md

@@ -78,7 +78,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `opentelemetry-instrumentation-sqlalchemy`: Fix a remaining memory leak in EngineTracer
   ([#3053](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3053))
 - `opentelemetry-instrumentation-sqlite3`: Update documentation on explicit cursor support of tracing
-  ([#3088](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3088))
+  ([#3088](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3088))'
+- `opentelemetry-instrumentation-psycopg2`: fix AttributeError at `instrument_connection`
+  ([#3043](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3043))
 
 ### Breaking changes
 

instrumentation/opentelemetry-instrumentation-psycopg2/README.rst

@@ -1,5 +1,5 @@
-OpenTelemetry Psycopg Instrumentation
-=====================================
+OpenTelemetry Psycopg2 Instrumentation
+======================================
 
 |pypi|
 
@@ -16,6 +16,6 @@ Installation
 
 References
 ----------
-* `OpenTelemetry Psycopg Instrumentation <https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg2/psycopg2.html>`_
+* `OpenTelemetry Psycopg2 Instrumentation <https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/psycopg2/psycopg2.html>`_
 * `OpenTelemetry Project <https://opentelemetry.io/>`_
 * `OpenTelemetry Python Examples <https://github.com/open-telemetry/opentelemetry-python/tree/main/docs/examples>`_

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

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 """
-The integration with PostgreSQL supports the `Psycopg`_ library, it can be enabled by
+The integration with PostgreSQL supports the `Psycopg2`_ library, it can be enabled by
 using ``Psycopg2Instrumentor``.
 
-.. _Psycopg: http://initd.org/psycopg/
+.. _Psycopg2: https://www.psycopg.org/docs/
 
 SQLCOMMENTER
 *****************************************
@@ -127,13 +127,13 @@
 )
 from psycopg2.sql import Composed  # pylint: disable=no-name-in-module
 
+from opentelemetry import trace as trace_api
 from opentelemetry.instrumentation import dbapi
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.psycopg2.package import _instruments
 from opentelemetry.instrumentation.psycopg2.version import __version__
 
 _logger = logging.getLogger(__name__)
-_OTEL_CURSOR_FACTORY_KEY = "_otel_orig_cursor_factory"
 
 
 class Psycopg2Instrumentor(BaseInstrumentor):
@@ -146,12 +146,14 @@ class Psycopg2Instrumentor(BaseInstrumentor):
 
     _DATABASE_SYSTEM = "postgresql"
 
+    _otel_cursor_factories = {}
+
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
     def _instrument(self, **kwargs):
-        """Integrate with PostgreSQL Psycopg library.
-        Psycopg: http://initd.org/psycopg/
+        """Integrate with PostgreSQL Psycopg2 library.
+        Psycopg2: https://www.psycopg.org/docs/
         """
         tracer_provider = kwargs.get("tracer_provider")
         enable_sqlcommenter = kwargs.get("enable_commenter", False)
@@ -174,8 +176,11 @@ def _uninstrument(self, **kwargs):
         dbapi.unwrap_connect(psycopg2, "connect")
 
     # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
-    @staticmethod
-    def instrument_connection(connection, tracer_provider=None):
+    def instrument_connection(
+        self,
+        connection,
+        tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
+    ):
         """Enable instrumentation in a psycopg2 connection.
 
         Args:
@@ -188,31 +193,25 @@ def instrument_connection(connection, tracer_provider=None):
         Returns:
             An instrumented psycopg2 connection object.
         """
+        # 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 in instrumentor map because
+        # psycopg2 connection type does not allow arbitrary attrs
+        self._otel_cursor_factories[connection] = connection.cursor_factory
+        connection.cursor_factory = _new_cursor_factory(
+            base_factory=connection.cursor_factory,
+            tracer_provider=tracer_provider,
+        )
+        _logger.warning("factories: %s", self._otel_cursor_factories)
 
-        if not hasattr(connection, "_is_instrumented_by_opentelemetry"):
-            connection._is_instrumented_by_opentelemetry = False
-
-        if not connection._is_instrumented_by_opentelemetry:
-            setattr(
-                connection, _OTEL_CURSOR_FACTORY_KEY, connection.cursor_factory
-            )
-            connection.cursor_factory = _new_cursor_factory(
-                tracer_provider=tracer_provider
-            )
-            connection._is_instrumented_by_opentelemetry = True
-        else:
-            _logger.warning(
-                "Attempting to instrument Psycopg connection while already instrumented"
-            )
         return connection
 
     # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
-    @staticmethod
-    def uninstrument_connection(connection):
-        connection.cursor_factory = getattr(
-            connection, _OTEL_CURSOR_FACTORY_KEY, None
+    def uninstrument_connection(self, connection):
+        connection.cursor_factory = self._otel_cursor_factories.pop(
+            connection, None
         )
-
         return connection
 
 

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

@@ -19,7 +19,9 @@
 
 import opentelemetry.instrumentation.psycopg2
 from opentelemetry import trace
-from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
+from opentelemetry.instrumentation.psycopg2 import (
+    Psycopg2Instrumentor,
+)
 from opentelemetry.sdk import resources
 from opentelemetry.test.test_base import TestBase
 
@@ -190,7 +192,12 @@ def test_instrument_connection(self):
         spans_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans_list), 0)
 
-        cnx = Psycopg2Instrumentor().instrument_connection(cnx)
+        instrumentor = Psycopg2Instrumentor()
+        original_cursor_factory = cnx.cursor_factory
+        cnx = instrumentor.instrument_connection(cnx)
+        self.assertEqual(
+            instrumentor._otel_cursor_factories[cnx], original_cursor_factory
+        )
         cursor = cnx.cursor()
         cursor.execute(query)
 
@@ -207,17 +214,70 @@ def test_instrument_connection_with_instrument(self):
         spans_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans_list), 0)
 
-        Psycopg2Instrumentor().instrument()
+        instrumentor = Psycopg2Instrumentor()
+        instrumentor.instrument()
+
+        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(
+            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)
+
+    def test_instrument_connection_with_instrument_connection(self):
+        cnx = psycopg2.connect(database="test")
+        query = "SELECT * FROM test"
+        cursor = cnx.cursor()
+        cursor.execute(query)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 0)
+
+        cnx = psycopg2.connect(database="test")
+        original_cursor_factory = cnx.cursor_factory
+        instrumentor = Psycopg2Instrumentor()
+        orig_cnx = cnx
+        cnx = instrumentor.instrument_connection(cnx)
+        self.assertEqual(
+            instrumentor._otel_cursor_factories.get(orig_cnx),
+            original_cursor_factory,
+        )
+
+        orig_cnx = cnx
+        original_cursor_factory = cnx.cursor_factory
+        instrumentor = Psycopg2Instrumentor()
+        cnx = instrumentor.instrument_connection(cnx)
+        self.assertEqual(
+            instrumentor._otel_cursor_factories.get(orig_cnx),
+            original_cursor_factory,
+        )
+        cursor = cnx.cursor()
+        cursor.execute(query)
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        # 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):
-        Psycopg2Instrumentor().instrument()
+        instrumentor = Psycopg2Instrumentor()
+        instrumentor.instrument()
         cnx = psycopg2.connect(database="test")
         query = "SELECT * FROM test"
         cursor = cnx.cursor()
@@ -236,20 +296,29 @@ def test_uninstrument_connection_with_instrument(self):
     # pylint: disable=unused-argument
     def test_uninstrument_connection_with_instrument_connection(self):
         cnx = psycopg2.connect(database="test")
-        Psycopg2Instrumentor().instrument_connection(cnx)
+        original_cursor_factory = cnx.cursor_factory
+        instrumentor = Psycopg2Instrumentor()
+        instrumentor.instrument_connection(cnx)
         query = "SELECT * FROM test"
         cursor = cnx.cursor()
         cursor.execute(query)
 
         spans_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans_list), 1)
+        self.assertEqual(
+            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.assertEqual(
+            instrumentor._otel_cursor_factories.get(orig_cnx), None
+        )
 
     @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
     def test_sqlcommenter_enabled(self, event_mocked):