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
 

docs-requirements.txt

@@ -32,6 +32,7 @@ mysqlclient~=2.1.1
 openai >= 1.26.0
 psutil>=5
 psycopg~=3.1.17
+psycopg2~=2.9.9
 pika>=0.12.0
 pymongo~=4.6.3
 PyMySQL~=1.1.1

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,39 @@ 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,
+        enable_commenter: bool = False,
+        commenter_options: dict = 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,
+            enable_commenter=enable_commenter,
+            commenter_options=commenter_options,
+        )
+        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
 
 
@@ -231,14 +240,23 @@ def get_statement(self, cursor, args):
         return statement
 
 
-def _new_cursor_factory(db_api=None, base_factory=None, tracer_provider=None):
+def _new_cursor_factory(
+    db_api: dbapi.DatabaseApiIntegration = None,
+    base_factory: pg_cursor = None,
+    tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
+    enable_commenter: bool = False,
+    commenter_options: dict = None,
+):
     if not db_api:
         db_api = DatabaseApiIntegration(
             __name__,
             Psycopg2Instrumentor._DATABASE_SYSTEM,
             connection_attributes=Psycopg2Instrumentor._CONNECTION_ATTRIBUTES,
             version=__version__,
             tracer_provider=tracer_provider,
+            enable_commenter=enable_commenter,
+            commenter_options=commenter_options,
+            connect_module=psycopg2,
         )
 
     base_factory = base_factory or pg_cursor

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

@@ -208,6 +208,7 @@ def test_instrument_connection_with_instrument(self):
         self.assertEqual(len(spans_list), 0)
 
         Psycopg2Instrumentor().instrument()
+        cnx = psycopg2.connect(database="test")
         cnx = Psycopg2Instrumentor().instrument_connection(cnx)
         cursor = cnx.cursor()
         cursor.execute(query)