From bc77e8f54613d4a869bf6c87966b6da25149a73b Mon Sep 17 00:00:00 2001 From: "Qiu, David" Date: Wed, 31 Jul 2024 20:29:56 +0800 Subject: [PATCH] Fix psycopg2 intrusment issue Signed-off-by: Qiu, David --- .../instrumentation/psycopg2/__init__.py | 49 +++++++++++-------- .../tests/test_psycopg2_integration.py | 2 +- 2 files changed, 29 insertions(+), 22 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py b/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py index de2e49f4c3..7fa549f22c 100644 --- a/instrumentation/opentelemetry-instrumentation-psycopg2/src/opentelemetry/instrumentation/psycopg2/__init__.py +++ b/instrumentation/opentelemetry-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 index 369d63d5cf..ee685a3b1a 100644 --- 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)