Support enable_commenter for instrument_connection by psycopg(2

CHANGELOG.md

@@ -98,6 +98,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#2942](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2942))
 - `opentelemetry-instrumentation-click`: new instrumentation to trace click commands
   ([#2994](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2994))
+- `opentelemetry-instrumentation-psycopg2`, `opentelemetry-instrumentation-psycopg`: Add sqlcommenter support for `instrument_connection`
+  ([#3071](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3071))
 
 ### Fixed
 

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

docs/conf.py

@@ -122,6 +122,8 @@
         "https://opentelemetry-python.readthedocs.io/en/latest/",
         None,
     ),
+    "psycopg": ("https://www.psycopg.org/psycopg3/docs/", None),
+    "psycopg2": ("https://www.psycopg.org/docs/", None),
 }
 
 # http://www.sphinx-doc.org/en/master/config.html#confval-nitpicky

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

@@ -16,7 +16,7 @@
 The integration with PostgreSQL supports the `Psycopg`_ library, it can be enabled by
 using ``PsycopgInstrumentor``.
 
-.. _Psycopg: http://initd.org/psycopg/
+.. _Psycopg: https://www.psycopg.org/psycopg3/docs/
 
 SQLCOMMENTER
 *****************************************
@@ -140,11 +140,12 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Callable, Collection, TypeVar
+from typing import Any, Callable, Collection, Optional, TypeVar
 
 import psycopg  # pylint: disable=import-self
 from psycopg.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.psycopg.package import _instruments
@@ -175,7 +176,7 @@ def instrumentation_dependencies(self) -> Collection[str]:
 
     def _instrument(self, **kwargs: Any):
         """Integrate with PostgreSQL Psycopg library.
-        Psycopg: http://initd.org/psycopg/
+        Psycopg: https://www.psycopg.org/psycopg3/docs/
         """
         tracer_provider = kwargs.get("tracer_provider")
         enable_sqlcommenter = kwargs.get("enable_commenter", False)
@@ -239,16 +240,26 @@ def _uninstrument(self, **kwargs: Any):
     # TODO(owais): check if core dbapi can do this for all dbapi implementations e.g, pymysql and mysql
     @staticmethod
     def instrument_connection(
-        connection: ConnectionT, tracer_provider: TracerProvider | None = None
-    ) -> ConnectionT:
-        """Enable instrumentation in a psycopg connection.
+        connection: psycopg.Connection,
+        tracer_provider: Optional[trace_api.TracerProvider] = None,
+        enable_commenter: bool = False,
+        commenter_options: dict = None,
+        enable_attribute_commenter: bool = False,
+    ):
+        """Enable instrumentation of a Psycopg connection.
 
         Args:
             connection: psycopg.Connection
                 The psycopg connection object to be instrumented.
             tracer_provider: opentelemetry.trace.TracerProvider, optional
                 The TracerProvider to use for instrumentation. If not provided,
                 the global TracerProvider will be used.
+            enable_commenter: bool, optional
+                Optional flag to enable/disable sqlcommenter (default False).
+            commenter_options: dict, optional
+                Optional configurations for tags to be appended at the sql query.
+            enable_attribute_commenter:
+                Optional flag to enable/disable addition of sqlcomment to span attribute (default False). Requires enable_commenter=True.
 
         Returns:
             An instrumented psycopg connection object.
@@ -261,7 +272,10 @@ def instrument_connection(
                 connection, _OTEL_CURSOR_FACTORY_KEY, connection.cursor_factory
             )
             connection.cursor_factory = _new_cursor_factory(
-                tracer_provider=tracer_provider
+                tracer_provider=tracer_provider,
+                enable_commenter=enable_commenter,
+                commenter_options=commenter_options,
+                enable_attribute_commenter=enable_attribute_commenter,
             )
             connection._is_instrumented_by_opentelemetry = True
         else:
@@ -346,9 +360,12 @@ def get_statement(self, cursor: CursorT, args: list[Any]) -> str:
 
 
 def _new_cursor_factory(
-    db_api: DatabaseApiIntegration | None = None,
-    base_factory: type[psycopg.Cursor] | None = None,
-    tracer_provider: TracerProvider | None = None,
+    db_api: dbapi.DatabaseApiIntegration = None,
+    base_factory: psycopg.Cursor = None,
+    tracer_provider: Optional[trace_api.TracerProvider] = None,
+    enable_commenter: bool = False,
+    commenter_options: dict = None,
+    enable_attribute_commenter: bool = False,
 ):
     if not db_api:
         db_api = DatabaseApiIntegration(
@@ -357,6 +374,10 @@ def _new_cursor_factory(
             connection_attributes=PsycopgInstrumentor._CONNECTION_ATTRIBUTES,
             version=__version__,
             tracer_provider=tracer_provider,
+            enable_commenter=enable_commenter,
+            commenter_options=commenter_options,
+            connect_module=psycopg,
+            enable_attribute_commenter=enable_attribute_commenter,
         )
 
     base_factory = base_factory or psycopg.Cursor

instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py

@@ -20,6 +20,7 @@
 import opentelemetry.instrumentation.psycopg
 from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor
 from opentelemetry.sdk import resources
+from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.test.test_base import TestBase
 
 
@@ -378,6 +379,121 @@ def test_sqlcommenter_enabled(self, event_mocked):
         kwargs = event_mocked.call_args[1]
         self.assertEqual(kwargs["enable_commenter"], True)
 
+    def test_sqlcommenter_enabled_instrument_connection_defaults(self):
+        with mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.__version__",
+            "foobar",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.pq.__build_version__",
+            "foobaz",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.threadsafety",
+            "123",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.apilevel",
+            "123",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.paramstyle",
+            "test",
+        ):
+            cnx = psycopg.connect(database="test")
+            cnx = PsycopgInstrumentor().instrument_connection(
+                cnx,
+                enable_commenter=True,
+            )
+            query = "Select 1"
+            cursor = cnx.cursor()
+            cursor.execute(query)
+            spans_list = self.memory_exporter.get_finished_spans()
+            span = spans_list[0]
+            span_id = format(span.get_span_context().span_id, "016x")
+            trace_id = format(span.get_span_context().trace_id, "032x")
+            self.assertEqual(
+                MockCursor.execute.call_args[0][0],
+                f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/",
+            )
+            self.assertEqual(
+                span.attributes[SpanAttributes.DB_STATEMENT],
+                "Select 1",
+            )
+
+    def test_sqlcommenter_enabled_instrument_connection_stmt_enabled(self):
+        with mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.__version__",
+            "foobar",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.pq.__build_version__",
+            "foobaz",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.threadsafety",
+            "123",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.apilevel",
+            "123",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.paramstyle",
+            "test",
+        ):
+            cnx = psycopg.connect(database="test")
+            cnx = PsycopgInstrumentor().instrument_connection(
+                cnx,
+                enable_commenter=True,
+                enable_attribute_commenter=True,
+            )
+            query = "Select 1"
+            cursor = cnx.cursor()
+            cursor.execute(query)
+            spans_list = self.memory_exporter.get_finished_spans()
+            span = spans_list[0]
+            span_id = format(span.get_span_context().span_id, "016x")
+            trace_id = format(span.get_span_context().trace_id, "032x")
+            self.assertEqual(
+                MockCursor.execute.call_args[0][0],
+                f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/",
+            )
+            self.assertEqual(
+                span.attributes[SpanAttributes.DB_STATEMENT],
+                f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/",
+            )
+
+    def test_sqlcommenter_enabled_instrument_connection_with_options(self):
+        with mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.__version__",
+            "foobar",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.pq.__build_version__",
+            "foobaz",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.threadsafety",
+            "123",
+        ):
+            cnx = psycopg.connect(database="test")
+            cnx = PsycopgInstrumentor().instrument_connection(
+                cnx,
+                enable_commenter=True,
+                commenter_options={
+                    "dbapi_level": False,
+                    "dbapi_threadsafety": True,
+                    "driver_paramstyle": False,
+                    "foo": "ignored",
+                },
+            )
+            query = "Select 1"
+            cursor = cnx.cursor()
+            cursor.execute(query)
+            spans_list = self.memory_exporter.get_finished_spans()
+            span = spans_list[0]
+            span_id = format(span.get_span_context().span_id, "016x")
+            trace_id = format(span.get_span_context().trace_id, "032x")
+            self.assertEqual(
+                MockCursor.execute.call_args[0][0],
+                f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_threadsafety='123',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/",
+            )
+            self.assertEqual(
+                span.attributes[SpanAttributes.DB_STATEMENT],
+                "Select 1",
+            )
+
     @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
     def test_sqlcommenter_disabled(self, event_mocked):
         cnx = psycopg.connect(database="test")
@@ -388,6 +504,45 @@ def test_sqlcommenter_disabled(self, event_mocked):
         kwargs = event_mocked.call_args[1]
         self.assertEqual(kwargs["enable_commenter"], False)
 
+    def test_sqlcommenter_disabled_default_instrument_connection(self):
+        cnx = psycopg.connect(database="test")
+        cnx = PsycopgInstrumentor().instrument_connection(
+            cnx,
+        )
+        query = "Select 1"
+        cursor = cnx.cursor()
+        cursor.execute(query)
+        self.assertEqual(
+            MockCursor.execute.call_args[0][0],
+            "Select 1",
+        )
+        spans_list = self.memory_exporter.get_finished_spans()
+        span = spans_list[0]
+        self.assertEqual(
+            span.attributes[SpanAttributes.DB_STATEMENT],
+            "Select 1",
+        )
+
+    def test_sqlcommenter_disabled_explicit_instrument_connection(self):
+        cnx = psycopg.connect(database="test")
+        cnx = PsycopgInstrumentor().instrument_connection(
+            cnx,
+            enable_commenter=False,
+        )
+        query = "Select 1"
+        cursor = cnx.cursor()
+        cursor.execute(query)
+        self.assertEqual(
+            MockCursor.execute.call_args[0][0],
+            "Select 1",
+        )
+        spans_list = self.memory_exporter.get_finished_spans()
+        span = spans_list[0]
+        self.assertEqual(
+            span.attributes[SpanAttributes.DB_STATEMENT],
+            "Select 1",
+        )
+
 
 class TestPostgresqlIntegrationAsync(
     PostgresqlIntegrationTestMixin, TestBase, IsolatedAsyncioTestCase