Support enable_commenter for instrument_connection by psycopg(2

CHANGELOG.md

@@ -88,6 +88,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
 *****************************************
@@ -150,6 +150,7 @@
 )
 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
@@ -174,7 +175,7 @@ def instrumentation_dependencies(self) -> Collection[str]:
 
     def _instrument(self, **kwargs):
         """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)
@@ -237,15 +238,24 @@ def _uninstrument(self, **kwargs):
 
     # 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):
-        """Enable instrumentation in a psycopg connection.
+    def instrument_connection(
+        connection: psycopg.Connection,
+        tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
+        enable_commenter: bool = False,
+        commenter_options: dict = None,
+    ):
+        """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.
 
         Returns:
             An instrumented psycopg connection object.
@@ -258,7 +268,9 @@ def instrument_connection(connection, tracer_provider=None):
                 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,
             )
             connection._is_instrumented_by_opentelemetry = True
         else:
@@ -342,14 +354,23 @@ def get_statement(self, cursor, args):
         return statement
 
 
-def _new_cursor_factory(db_api=None, base_factory=None, tracer_provider=None):
+def _new_cursor_factory(
+    db_api: dbapi.DatabaseApiIntegration = None,
+    base_factory: pg_cursor = None,
+    tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
+    enable_commenter: bool = False,
+    commenter_options: dict = None,
+):
     if not db_api:
         db_api = DatabaseApiIntegration(
             __name__,
             PsycopgInstrumentor._DATABASE_SYSTEM,
             connection_attributes=PsycopgInstrumentor._CONNECTION_ATTRIBUTES,
             version=__version__,
             tracer_provider=tracer_provider,
+            enable_commenter=enable_commenter,
+            commenter_options=commenter_options,
+            connect_module=psycopg,
         )
 
     base_factory = base_factory or pg_cursor

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

@@ -378,6 +378,74 @@ 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'*/",
+            )
+
+    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'*/",
+            )
+
     @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
     def test_sqlcommenter_disabled(self, event_mocked):
         cnx = psycopg.connect(database="test")
@@ -388,6 +456,33 @@ 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",
+        )
+
+    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",
+        )
+
 
 class TestPostgresqlIntegrationAsync(
     PostgresqlIntegrationTestMixin, TestBase, IsolatedAsyncioTestCase

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

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 """
-The integration with PostgreSQL supports the `Psycopg`_ library, it can be enabled by
+The integration with PostgreSQL supports the `Psycopg2`_ library, it can be enabled by
 using ``Psycopg2Instrumentor``.
 
-.. _Psycopg: http://initd.org/psycopg/
+.. _Psycopg2: https://www.psycopg.org/docs/
 
 SQLCOMMENTER
 *****************************************
@@ -142,11 +142,15 @@
 from typing import Collection
 
 import psycopg2
+from psycopg2.extensions import (
+    connection as pg_connection,  # pylint: disable=no-name-in-module
+)
 from psycopg2.extensions import (
     cursor as pg_cursor,  # pylint: disable=no-name-in-module
 )
 from psycopg2.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.psycopg2.package import _instruments
@@ -170,8 +174,8 @@ def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
     def _instrument(self, **kwargs):
-        """Integrate with PostgreSQL Psycopg library.
-        Psycopg: http://initd.org/psycopg/
+        """Integrate with PostgreSQL Psycopg2 library.
+        Psycopg2: https://www.psycopg.org/docs/
         """
         tracer_provider = kwargs.get("tracer_provider")
         enable_sqlcommenter = kwargs.get("enable_commenter", False)
@@ -199,20 +203,28 @@ def _uninstrument(self, **kwargs):
 
     # 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):
-        """Enable instrumentation in a psycopg2 connection.
+    def instrument_connection(
+        connection: pg_connection,
+        tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
+        enable_commenter: bool = False,
+        commenter_options: dict = None,
+    ):
+        """Enable instrumentation of a Psycopg2 connection.
 
         Args:
             connection: psycopg2.extensions.connection
                 The psycopg2 connection object to be instrumented.
             tracer_provider: opentelemetry.trace.TracerProvider, optional
-                The TracerProvider to use for instrumentation. If not specified,
+                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.
 
         Returns:
             An instrumented psycopg2 connection object.
         """
-
         if not hasattr(connection, "_is_instrumented_by_opentelemetry"):
             connection._is_instrumented_by_opentelemetry = False
 
@@ -221,7 +233,9 @@ def instrument_connection(connection, tracer_provider=None):
                 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,
             )
             connection._is_instrumented_by_opentelemetry = True
         else:
@@ -284,14 +298,23 @@ def get_statement(self, cursor, args):
         return statement
 
 
-def _new_cursor_factory(db_api=None, base_factory=None, tracer_provider=None):
+def _new_cursor_factory(
+    db_api: dbapi.DatabaseApiIntegration = None,
+    base_factory: pg_cursor = None,
+    tracer_provider: typing.Optional[trace_api.TracerProvider] = None,
+    enable_commenter: bool = False,
+    commenter_options: dict = None,
+):
     if not db_api:
         db_api = DatabaseApiIntegration(
             __name__,
             Psycopg2Instrumentor._DATABASE_SYSTEM,
             connection_attributes=Psycopg2Instrumentor._CONNECTION_ATTRIBUTES,
             version=__version__,
             tracer_provider=tracer_provider,
+            enable_commenter=enable_commenter,
+            commenter_options=commenter_options,
+            connect_module=psycopg2,
         )
 
     base_factory = base_factory or pg_cursor