Fix psycopg2 instrument_connection AttributeError

CHANGELOG.md

@@ -50,6 +50,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#2816](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2816))
 - `opentelemetry-instrumentation-sqlalchemy`: Fix a remaining memory leak in EngineTracer
   ([#3053](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3053))
+- `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,4 +1,4 @@
-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
 *****************************************
@@ -111,13 +111,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):
@@ -130,12 +130,14 @@ class Psycopg2Instrumentor(BaseInstrumentor):
 
     _DATABASE_SYSTEM = "postgresql"
 
+    _otel_cursor_factory = None
+
     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)
@@ -158,32 +160,35 @@ 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):
-        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:
+    def instrument_connection(
+        self,
+        connection,
+        tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
+    ):
+        if self._is_instrumented_by_opentelemetry:
+            # _instrument (via BaseInstrumentor) or instrument_connection (this)
+            # was already called
             _logger.warning(
-                "Attempting to instrument Psycopg connection while already instrumented"
+                "Attempting to instrument Psycopg2 connection while already instrumented"
             )
+            self._is_instrumented_by_opentelemetry = True
+            return connection
+
+        # Save cursor_factory at instrumentor level because
+        # psycopg2 connection type does not allow arbitrary attrs
+        self._otel_cursor_factory = connection.cursor_factory
+        connection.cursor_factory = _new_cursor_factory(
+            base_factory=connection.cursor_factory,
+            tracer_provider=tracer_provider,
+        )
+        self._is_instrumented_by_opentelemetry = True
+
         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):
+        self._is_instrumented_by_opentelemetry = False
+        connection.cursor_factory = self._otel_cursor_factory
         return connection
 
 

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

@@ -16,10 +16,13 @@
 from unittest import mock
 
 import psycopg2
+import pytest
 
 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
 
@@ -66,6 +69,10 @@ 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(
@@ -190,7 +197,13 @@ 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.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
+        self.assertEqual(
+            instrumentor._otel_cursor_factory, original_cursor_factory
+        )
         cursor = cnx.cursor()
         cursor.execute(query)
 
@@ -207,8 +220,52 @@ 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()
+        self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
+
+        cnx = psycopg2.connect(database="test")
         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)
+        # 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)
+
+    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()
+        cnx = instrumentor.instrument_connection(cnx)
+        self.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
+        self.assertEqual(
+            instrumentor._otel_cursor_factory, 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)
+        self.assertEqual(
+            instrumentor._otel_cursor_factory, original_cursor_factory
+        )
         cursor = cnx.cursor()
         cursor.execute(query)
 
@@ -217,39 +274,50 @@ def test_instrument_connection_with_instrument(self):
 
     # 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()
         cursor.execute(query)
 
         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):
         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.assertTrue(instrumentor._is_instrumented_by_opentelemetry)
+        self.assertEqual(
+            instrumentor._otel_cursor_factory, original_cursor_factory
+        )
 
         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)
 
     @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
     def test_sqlcommenter_enabled(self, event_mocked):