Add psycopg2 native tags to sqlcommenter (#1203)

Author: Thiyagu55

CHANGELOG.md

@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#1187](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1187))
 - SQLCommenter semicolon bug fix
   ([#1200](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1200/files))
+- Add psycopg2 native tags to sqlcommenter
+  ([#1203](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1203))
 
 ### Added
 - `opentelemetry-instrumentation-redis` add support to instrument RedisCluster clients

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

@@ -105,6 +105,7 @@ def wrap_connect(
     capture_parameters: bool = False,
     enable_commenter: bool = False,
     db_api_integration_factory=None,
+    commenter_options: dict = None,
 ):
     """Integrate with DB API library.
     https://www.python.org/dev/peps/pep-0249/
@@ -119,6 +120,8 @@ def wrap_connect(
         tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to
             use. If omitted the current configured one is used.
         capture_parameters: Configure if db.statement.parameters should be captured.
+        enable_commenter: Flag to enable/disable sqlcommenter.
+        commenter_options: Configurations for tags to be appended at the sql query.
 
     """
     db_api_integration_factory = (
@@ -140,6 +143,8 @@ def wrap_connect_(
             tracer_provider=tracer_provider,
             capture_parameters=capture_parameters,
             enable_commenter=enable_commenter,
+            commenter_options=commenter_options,
+            connect_module=connect_module,
         )
         return db_integration.wrapped_connection(wrapped, args, kwargs)
 
@@ -173,6 +178,7 @@ def instrument_connection(
     tracer_provider: typing.Optional[TracerProvider] = None,
     capture_parameters: bool = False,
     enable_commenter: bool = False,
+    commenter_options: dict = None,
 ):
     """Enable instrumentation in a database connection.
 
@@ -185,6 +191,9 @@ def instrument_connection(
         tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to
             use. If omitted the current configured one is used.
         capture_parameters: Configure if db.statement.parameters should be captured.
+        enable_commenter: Flag to enable/disable sqlcommenter.
+        commenter_options: Configurations for tags to be appended at the sql query.
+
     Returns:
         An instrumented connection.
     """
@@ -200,6 +209,7 @@ def instrument_connection(
         tracer_provider=tracer_provider,
         capture_parameters=capture_parameters,
         enable_commenter=enable_commenter,
+        commenter_options=commenter_options,
     )
     db_integration.get_connection_attributes(connection)
     return get_traced_connection_proxy(connection, db_integration)
@@ -231,6 +241,8 @@ def __init__(
         tracer_provider: typing.Optional[TracerProvider] = None,
         capture_parameters: bool = False,
         enable_commenter: bool = False,
+        commenter_options: dict = None,
+        connect_module: typing.Callable[..., typing.Any] = None,
     ):
         self.connection_attributes = connection_attributes
         if self.connection_attributes is None:
@@ -249,11 +261,13 @@ def __init__(
         )
         self.capture_parameters = capture_parameters
         self.enable_commenter = enable_commenter
+        self.commenter_options = commenter_options
         self.database_system = database_system
         self.connection_props = {}
         self.span_attributes = {}
         self.name = ""
         self.database = ""
+        self.connect_module = connect_module
 
     def wrapped_connection(
         self,
@@ -335,12 +349,18 @@ class CursorTracer:
     def __init__(self, db_api_integration: DatabaseApiIntegration) -> None:
         self._db_api_integration = db_api_integration
         self._commenter_enabled = self._db_api_integration.enable_commenter
+        self._commenter_options = (
+            self._db_api_integration.commenter_options
+            if self._db_api_integration.commenter_options
+            else {}
+        )
+        self._connect_module = self._db_api_integration.connect_module
 
     def _populate_span(
         self,
         span: trace_api.Span,
         cursor,
-        *args: typing.Tuple[typing.Any, typing.Any]
+        *args: typing.Tuple[typing.Any, typing.Any],
     ):
         if not span.is_recording():
             return
@@ -380,7 +400,7 @@ def traced_execution(
         cursor,
         query_method: typing.Callable[..., typing.Any],
         *args: typing.Tuple[typing.Any, typing.Any],
-        **kwargs: typing.Dict[typing.Any, typing.Any]
+        **kwargs: typing.Dict[typing.Any, typing.Any],
     ):
         name = self.get_operation_name(cursor, args)
         if not name:
@@ -397,14 +417,32 @@ def traced_execution(
             if args and self._commenter_enabled:
                 try:
                     args_list = list(args)
-                    commenter_data = {}
-                    commenter_data.update(_get_opentelemetry_values())
+                    commenter_data = dict(
+                        # Psycopg2/framework information
+                        db_driver=f"psycopg2:{self._connect_module.__version__.split(' ')[0]}",
+                        dbapi_threadsafety=self._connect_module.threadsafety,
+                        dbapi_level=self._connect_module.apilevel,
+                        libpq_version=self._connect_module.__libpq_version__,
+                        driver_paramstyle=self._connect_module.paramstyle,
+                    )
+                    if self._commenter_options.get(
+                        "opentelemetry_values", True
+                    ):
+                        commenter_data.update(**_get_opentelemetry_values())
+
+                    # Filter down to just the requested attributes.
+                    commenter_data = {
+                        k: v
+                        for k, v in commenter_data.items()
+                        if self._commenter_options.get(k, True)
+                    }
                     statement = _add_sql_comment(
                         args_list[0], **commenter_data
                     )
 
                     args_list[0] = statement
                     args = tuple(args_list)
+
                 except Exception as exc:  # pylint: disable=broad-except
                     _logger.exception(
                         "Exception while generating sql comment: %s", exc

instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py

@@ -229,8 +229,20 @@ def test_executemany(self):
         )
 
     def test_executemany_comment(self):
+
+        connect_module = mock.MagicMock()
+        connect_module.__version__ = mock.MagicMock()
+        connect_module.__libpq_version__ = 123
+        connect_module.apilevel = 123
+        connect_module.threadsafety = 123
+        connect_module.paramstyle = "test"
+
         db_integration = dbapi.DatabaseApiIntegration(
-            "testname", "testcomponent", enable_commenter=True
+            "testname",
+            "testcomponent",
+            enable_commenter=True,
+            commenter_options={"db_driver": False, "dbapi_level": False},
+            connect_module=connect_module,
         )
         mock_connection = db_integration.wrapped_connection(
             mock_connect, {}, {}
@@ -239,7 +251,7 @@ def test_executemany_comment(self):
         cursor.executemany("Select 1;")
         self.assertRegex(
             cursor.query,
-            r"Select 1 /\*traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+            r"Select 1 /\*dbapi_threadsafety=123,driver_paramstyle='test',libpq_version=123,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
         )
 
     def test_callproc(self):

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

@@ -18,6 +18,68 @@
 
 .. _Psycopg: http://initd.org/psycopg/
 
+SQLCOMMENTER
+*****************************************
+You can optionally configure Psycopg2 instrumentation to enable sqlcommenter which enriches
+the query with contextual information.
+
+Usage
+-----
+
+.. code:: python
+
+    from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
+
+    Psycopg2Instrumentor().instrument(enable_commenter=True, commenter_options={})
+
+
+For example,
+::
+
+   Invoking cursor.execute("select * from auth_users") will lead to sql query "select * from auth_users" but when SQLCommenter is enabled
+   the query will get appended with some configurable tags like "select * from auth_users /*tag=value*/;"
+
+
+SQLCommenter Configurations
+***************************
+We can configure the tags to be appended to the sqlquery log by adding configuration inside commenter_options(default:{}) keyword
+
+db_driver = True(Default) or False
+
+For example,
+::
+Enabling this flag will add psycopg2 and it's version which is /*psycopg2%%3A2.9.3*/
+
+dbapi_threadsafety = True(Default) or False
+
+For example,
+::
+Enabling this flag will add threadsafety /*dbapi_threadsafety=2*/
+
+dbapi_level = True(Default) or False
+
+For example,
+::
+Enabling this flag will add dbapi_level /*dbapi_level='2.0'*/
+
+libpq_version = True(Default) or False
+
+For example,
+::
+Enabling this flag will add libpq_version /*libpq_version=140001*/
+
+driver_paramstyle = True(Default) or False
+
+For example,
+::
+Enabling this flag will add driver_paramstyle /*driver_paramstyle='pyformat'*/
+
+opentelemetry_values = True(Default) or False
+
+For example,
+::
+Enabling this flag will add traceparent values /*traceparent='00-03afa25236b8cd948fa853d67038ac79-405ff022e8247c46-01'*/
+
 Usage
 -----
 
@@ -77,6 +139,7 @@ def _instrument(self, **kwargs):
         """
         tracer_provider = kwargs.get("tracer_provider")
         enable_sqlcommenter = kwargs.get("enable_commenter", False)
+        commenter_options = kwargs.get("commenter_options", {})
         dbapi.wrap_connect(
             __name__,
             psycopg2,
@@ -87,6 +150,7 @@ def _instrument(self, **kwargs):
             tracer_provider=tracer_provider,
             db_api_integration_factory=DatabaseApiIntegration,
             enable_commenter=enable_sqlcommenter,
+            commenter_options=commenter_options,
         )
 
     def _uninstrument(self, **kwargs):

tests/opentelemetry-docker-tests/tests/postgres/test_psycopg_sqlcommenter.py

@@ -50,5 +50,5 @@ def test_commenter_enabled(self):
         self._cursor.execute("SELECT  1;")
         self.assertRegex(
             self._cursor.query.decode("ascii"),
-            r"SELECT  1 /\*traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+            r"SELECT  1 /\*db_driver='psycopg2(.*)',dbapi_level='\d.\d',dbapi_threadsafety=\d,driver_paramstyle=(.*),libpq_version=\d*,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
         )