db.statement inclusion of sqlcomment as opt-in

Author: tammy-baylis-swi

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

@@ -65,6 +65,28 @@
 ::
 Enabling this flag will add traceparent values /*traceparent='00-03afa25236b8cd948fa853d67038ac79-405ff022e8247c46-01'*/
 
+SQLComment in span attribute
+****************************
+If sqlcommenter is enabled, you can optionally configure SQLAlchemy instrumentation to append sqlcomment to query span attribute for convenience of your platform.
+
+.. code:: python
+
+    from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+
+    SQLAlchemyInstrumentor().instrument(
+        enable_commenter=True,
+        commenter_options={},
+        enable_attribute_commenter=True,
+    )
+
+
+For example,
+::
+
+    Invoking engine.execute("select * from auth_users") will lead to sql query "select * from auth_users" but when SQLCommenter and attribute_commenter is enabled
+    the query will get appended with some configurable tags like "select * from auth_users /*tag=value*/;" for both server query and `db.statement` span attribute.
+
+
 Usage
 -----
 .. code:: python
@@ -138,6 +160,7 @@ def _instrument(self, **kwargs):
                 ``meter_provider``: a MeterProvider, defaults to global
                 ``enable_commenter``: bool to enable sqlcommenter, defaults to False
                 ``commenter_options``: dict of sqlcommenter config, defaults to {}
+                ``enable_attribute_commenter``: bool to enable sqlcomment addition to span attribute, defaults to False. Must also set `enable_commenter`.
 
         Returns:
             An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise.
@@ -166,19 +189,30 @@ def _instrument(self, **kwargs):
 
         enable_commenter = kwargs.get("enable_commenter", False)
         commenter_options = kwargs.get("commenter_options", {})
+        enable_attribute_commenter = kwargs.get(
+            "enable_attribute_commenter", False
+        )
 
         _w(
             "sqlalchemy",
             "create_engine",
             _wrap_create_engine(
-                tracer, connections_usage, enable_commenter, commenter_options
+                tracer,
+                connections_usage,
+                enable_commenter,
+                commenter_options,
+                enable_attribute_commenter,
             ),
         )
         _w(
             "sqlalchemy.engine",
             "create_engine",
             _wrap_create_engine(
-                tracer, connections_usage, enable_commenter, commenter_options
+                tracer,
+                connections_usage,
+                enable_commenter,
+                commenter_options,
+                enable_attribute_commenter,
             ),
         )
         # sqlalchemy.engine.create is not present in earlier versions of sqlalchemy (which we support)
@@ -191,6 +225,7 @@ def _instrument(self, **kwargs):
                     connections_usage,
                     enable_commenter,
                     commenter_options,
+                    enable_attribute_commenter,
                 ),
             )
         _w(
@@ -207,6 +242,7 @@ def _instrument(self, **kwargs):
                     connections_usage,
                     enable_commenter,
                     commenter_options,
+                    enable_attribute_commenter,
                 ),
             )
         if kwargs.get("engine") is not None:
@@ -216,6 +252,7 @@ def _instrument(self, **kwargs):
                 connections_usage,
                 kwargs.get("enable_commenter", False),
                 kwargs.get("commenter_options", {}),
+                kwargs.get("enable_attribute_commenter", False),
             )
         if kwargs.get("engines") is not None and isinstance(
             kwargs.get("engines"), Sequence
@@ -227,6 +264,7 @@ def _instrument(self, **kwargs):
                     connections_usage,
                     kwargs.get("enable_commenter", False),
                     kwargs.get("commenter_options", {}),
+                    kwargs.get("enable_attribute_commenter", False),
                 )
                 for engine in kwargs.get("engines")
             ]

instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py

@@ -43,7 +43,11 @@ def _normalize_vendor(vendor):
 
 
 def _wrap_create_async_engine(
-    tracer, connections_usage, enable_commenter=False, commenter_options=None
+    tracer,
+    connections_usage,
+    enable_commenter=False,
+    commenter_options=None,
+    enable_attribute_commenter=False,
 ):
     # pylint: disable=unused-argument
     def _wrap_create_async_engine_internal(func, module, args, kwargs):
@@ -57,14 +61,19 @@ def _wrap_create_async_engine_internal(func, module, args, kwargs):
             connections_usage,
             enable_commenter,
             commenter_options,
+            enable_attribute_commenter,
         )
         return engine
 
     return _wrap_create_async_engine_internal
 
 
 def _wrap_create_engine(
-    tracer, connections_usage, enable_commenter=False, commenter_options=None
+    tracer,
+    connections_usage,
+    enable_commenter=False,
+    commenter_options=None,
+    enable_attribute_commenter=False,
 ):
     def _wrap_create_engine_internal(func, _module, args, kwargs):
         """Trace the SQLAlchemy engine, creating an `EngineTracer`
@@ -77,6 +86,7 @@ def _wrap_create_engine_internal(func, _module, args, kwargs):
             connections_usage,
             enable_commenter,
             commenter_options,
+            enable_attribute_commenter,
         )
         return engine
 
@@ -110,12 +120,14 @@ def __init__(
         connections_usage,
         enable_commenter=False,
         commenter_options=None,
+        enable_attribute_commenter=False,
     ):
         self.tracer = tracer
         self.connections_usage = connections_usage
         self.vendor = _normalize_vendor(engine.name)
         self.enable_commenter = enable_commenter
         self.commenter_options = commenter_options if commenter_options else {}
+        self.enable_attribute_commenter = enable_attribute_commenter
         self._engine_attrs = _get_attributes_from_engine(engine)
         self._leading_comment_remover = re.compile(r"^/\*.*?\*/")
 
@@ -251,13 +263,22 @@ def _before_cur_exec(
                         if self.commenter_options.get(k, True)
                     }
 
-                    statement = _add_sql_comment(statement, **commenter_data)
+                    if self.enable_attribute_commenter:
+                        statement = _add_sql_comment(
+                            statement, **commenter_data
+                        )
 
                 span.set_attribute(SpanAttributes.DB_STATEMENT, statement)
                 span.set_attribute(SpanAttributes.DB_SYSTEM, self.vendor)
                 for key, value in attrs.items():
                     span.set_attribute(key, value)
 
+                if (
+                    self.enable_commenter
+                    and not self.enable_attribute_commenter
+                ):
+                    statement = _add_sql_comment(statement, **commenter_data)
+
         context._otel_span = span
 
         return statement, params

instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlcommenter.py

@@ -61,13 +61,45 @@ def test_sqlcommenter_enabled(self):
             r"SELECT  1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
         )
 
-    def test_sqlcommenter_enabled_matches_db_statement_attribute(self):
+    def test_sqlcommenter_enabled_stmt_disabled_default_matches_db_statement_attribute(
+        self,
+    ):
         engine = create_engine("sqlite:///:memory:")
         SQLAlchemyInstrumentor().instrument(
             engine=engine,
             tracer_provider=self.tracer_provider,
             enable_commenter=True,
             commenter_options={"db_framework": False},
+            # enable_attribute_commenter not set
+        )
+        cnx = engine.connect()
+        cnx.execute(text("SELECT  1;")).fetchall()
+        query_log = self.caplog.records[-2].getMessage()
+        self.assertRegex(
+            query_log,
+            r"SELECT  1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+        )
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 2)
+        # first span is connection to db
+        self.assertEqual(spans[0].name, "connect")
+        # second span is query itself
+        query_span = spans[1]
+        self.assertEqual(
+            query_span.attributes[SpanAttributes.DB_STATEMENT],
+            "SELECT  1;",
+        )
+
+    def test_sqlcommenter_enabled_stmt_enabled_matches_db_statement_attribute(
+        self,
+    ):
+        engine = create_engine("sqlite:///:memory:")
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+            enable_commenter=True,
+            commenter_options={"db_framework": False},
+            enable_attribute_commenter=True,
         )
         cnx = engine.connect()
         cnx.execute(text("SELECT  1;")).fetchall()