Support enable_commenter for instrument_connection by psycopg(2

CHANGELOG.md

@@ -55,6 +55,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#4141](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4141))
 - `opentelemetry-instrumentation-pyramid`: pass request attributes at span creation
   ([#4139](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4139))
+- `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

@@ -39,6 +39,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

@@ -129,6 +129,8 @@
         "https://opentelemetry-python.readthedocs.io/en/latest/",
         None,
     ),
+    "psycopg": ("https://www.psycopg.org/psycopg3/docs/", None),
+    "psycopg2": ("https://www.psycopg.org/docs/", None),
     "redis": ("https://redis.readthedocs.io/en/latest/", None),
     "fsspec": ("https://filesystem-spec.readthedocs.io/en/latest/", None),
 }

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

@@ -72,6 +72,13 @@
     from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor
 
     PsycopgInstrumentor().instrument(enable_commenter=True)
+    # OR with specific connection
+    cnx = psycopg.connect(database='Database')
+    instrumented_cnx = PsycopgInstrumentor().instrument_connection(
+        cnx,
+        enable_commenter=True,
+    )
+
 
 
 SQLCommenter with commenter_options
@@ -144,11 +151,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 Composable  # 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
@@ -179,7 +187,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)
@@ -247,16 +255,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.
@@ -270,11 +288,17 @@ def instrument_connection(
             )
             if isinstance(connection, psycopg.AsyncConnection):
                 connection.cursor_factory = _new_cursor_async_factory(
-                    tracer_provider=tracer_provider
+                    tracer_provider=tracer_provider,
+                    enable_commenter=enable_commenter,
+                    commenter_options=commenter_options,
+                    enable_attribute_commenter=enable_attribute_commenter,
                 )
             else:
                 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:
@@ -359,9 +383,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(
@@ -370,6 +397,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
@@ -398,6 +429,9 @@ def _new_cursor_async_factory(
     db_api: DatabaseApiAsyncIntegration | None = None,
     base_factory: type[psycopg.AsyncCursor] | None = None,
     tracer_provider: TracerProvider | None = None,
+    enable_commenter: bool = False,
+    commenter_options: dict = None,
+    enable_attribute_commenter: bool = False,
 ):
     if not db_api:
         db_api = DatabaseApiAsyncIntegration(
@@ -406,6 +440,10 @@ def _new_cursor_async_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.AsyncCursor
     _cursor_tracer = CursorTracer(db_api)

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

@@ -22,6 +22,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
 
 
@@ -451,6 +452,149 @@ 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.pq.version",
+                new=lambda: "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.pq.version",
+                new=lambda: "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.pq.version",
+                new=lambda: "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")
@@ -461,6 +605,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