Fix psycopg2 instrument issue

CHANGELOG.md

@@ -28,6 +28,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  ([#2823](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2823))
 - `opentelemetry-instrumentation-system-metrics` fix `process.runtime.cpu.utilization` values to be shown in range of 0 to 1
  ([#2812](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2812))
+- `opentelemetry-instrumentation-psycopg2` Fix psycopg2 instrument issue
+ ([#2795](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2795))
 - `opentelemetry-instrumentation-fastapi` fix `fastapi` auto-instrumentation by removing `fastapi-slim` support, `fastapi-slim` itself is discontinued from maintainers
  ([#2783](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2783))
 - `opentelemetry-instrumentation-aws-lambda` Avoid exception when a handler is not present.

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

@@ -110,6 +110,7 @@
  cursor as pg_cursor, # pylint: disable=no-name-in-module
 )
 from psycopg2.sql import Composed # pylint: disable=no-name-in-module
+from wrapt import ObjectProxy
 
 from opentelemetry.instrumentation import dbapi
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
@@ -158,24 +159,25 @@ 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=None):
+ if isinstance(connection, ObjectProxy):
  _logger.warning(
  "Attempting to instrument Psycopg connection while already instrumented"
  )
- return connection
+ return connection
+
+ db_integration = DatabaseApiIntegration(
+ __name__,
+ self._DATABASE_SYSTEM,
+ connection_attributes=self._CONNECTION_ATTRIBUTES,
+ version=__version__,
+ tracer_provider=tracer_provider,
+ )
+ db_integration.get_connection_attributes(connection)
+ db_integration.cursor_factory = _new_cursor_factory(
+ tracer_provider=tracer_provider
+ )
+ return dbapi.get_traced_connection_proxy(connection, db_integration)
 
  # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
  @staticmethod
@@ -184,7 +186,7 @@ def uninstrument_connection(connection):
  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

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

@@ -191,6 +191,7 @@ def test_instrument_connection(self):
  self.assertEqual(len(spans_list), 0)
 
  cnx = Psycopg2Instrumentor().instrument_connection(cnx)
+ self.assertFalse(hasattr(cnx, "_is_instrumented_by_opentelemetry"))
  cursor = cnx.cursor()
  cursor.execute(query)
 
@@ -209,6 +210,7 @@ def test_instrument_connection_with_instrument(self):
 
  Psycopg2Instrumentor().instrument()
  cnx = Psycopg2Instrumentor().instrument_connection(cnx)
+ self.assertFalse(hasattr(cnx, "_is_instrumented_by_opentelemetry"))
  cursor = cnx.cursor()
  cursor.execute(query)
 
@@ -236,7 +238,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)
@@ -281,3 +283,30 @@ def test_no_op_tracer_provider(self):
  cursor.execute(query)
  spans_list = self.memory_exporter.get_finished_spans()
  self.assertEqual(len(spans_list), 0)
+
+ # pylint: disable=unused-argument
+ def test_duplicated_instrumentation_works(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)
+
+ Psycopg2Instrumentor().instrument()
+ Psycopg2Instrumentor().instrument()
+ cnx = Psycopg2Instrumentor().instrument_connection(cnx)
+ self.assertFalse(hasattr(cnx, "_is_instrumented_by_opentelemetry"))
+ cursor = cnx.cursor()
+ cursor.execute(query)
+
+ spans_list = self.memory_exporter.get_finished_spans()
+ self.assertEqual(len(spans_list), 1)
+
+ 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)