Fix psycopg2 intrusment issue

Signed-off-by: Qiu, David <[email protected]>
fidelity-contributions · Jul 31, 2024 · bc77e8f · bc77e8f
1 parent dbfa681
commit bc77e8f
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 22 deletions.
diff --git a/...telemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py b/...telemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py
@@ -158,33 +158,40 @@ 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:
-            _logger.warning(
-                "Attempting to instrument Psycopg connection while already instrumented"
-            )
-        return connection
+    # TODO: comment out below codes to fix issue #2522, will try to set the properties later.
+    def instrument_connection(self, 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:
+        #     _logger.warning(
+        #         "Attempting to instrument Psycopg connection while already instrumented"
+        #     )
+        # return connection
+        return dbapi.instrument_connection(
+            __name__,
+            connection,
+            self._DATABASE_SYSTEM,
+            self._CONNECTION_ATTRIBUTES,
+            version=__version__,
+            tracer_provider=tracer_provider,
+        )
 
     # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
-    @staticmethod
-    def uninstrument_connection(connection):
+    def uninstrument_connection(self, connection):
         connection.cursor_factory = getattr(
             connection, _OTEL_CURSOR_FACTORY_KEY, None
         )
 
-        return connection
+        return dbapi.uninstrument_connection(connection)
 
 
 # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql

diff --git a/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py b/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py
@@ -235,7 +235,7 @@ 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)
+        cnx = Psycopg2Instrumentor().instrument_connection(cnx)
         query = "SELECT * FROM test"
         cursor = cnx.cursor()
         cursor.execute(query)