Sqlcommenter integration into SQLAlchemy (#924)

* Integrating sqlcommenter into psycopg2

* Integrating sqlcommenter into psycopg2 - Converted public local variable into private

* Added test cases for sqlcommenter & PR Changes

* Code refactoring for generate sqlcommenter

* Added testcase for sqlcommenter integration into sqlalchemy

* updated change log

* updated to accept latest logs

* Updated lint changes

* Fixed errors due to linting

* Fixed linting errors

* Fixed linting errors

* Fixed linting errors

* Update CHANGELOG.md

* Update CHANGELOG.md

Co-authored-by: Srikanth Chekuri <[email protected]>

Co-authored-by: Srikanth Chekuri <[email protected]>

Author: Thiyagu55

CHANGELOG.md

@@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- `opentelemetry-instrumentation-sqlalchemy` added experimental sql commenter capability
+   ([#924](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/924))
 - `opentelemetry-instrumentation-dbapi` add experimental sql commenter capability
   ([#908](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/908))
 - `opentelemetry-instrumentation-requests` make span attribute available to samplers

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

@@ -106,6 +106,7 @@ def _instrument(self, **kwargs):
             return EngineTracer(
                 _get_tracer(kwargs.get("engine"), tracer_provider),
                 kwargs.get("engine"),
+                kwargs.get("enable_commenter", False),
             )
         return None
 

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

@@ -17,7 +17,12 @@
 
 from opentelemetry import trace
 from opentelemetry.instrumentation.sqlalchemy.version import __version__
+from opentelemetry.instrumentation.utils import (
+    _generate_opentelemetry_traceparent,
+    _generate_sql_comment,
+)
 from opentelemetry.semconv.trace import NetTransportValues, SpanAttributes
+from opentelemetry.trace import Span
 from opentelemetry.trace.status import Status, StatusCode
 
 
@@ -70,12 +75,15 @@ def _wrap_create_engine_internal(func, module, args, kwargs):
 
 
 class EngineTracer:
-    def __init__(self, tracer, engine):
+    def __init__(self, tracer, engine, enable_commenter=False):
         self.tracer = tracer
         self.engine = engine
         self.vendor = _normalize_vendor(engine.name)
+        self.enable_commenter = enable_commenter
 
-        listen(engine, "before_cursor_execute", self._before_cur_exec)
+        listen(
+            engine, "before_cursor_execute", self._before_cur_exec, retval=True
+        )
         listen(engine, "after_cursor_execute", _after_cur_exec)
         listen(engine, "handle_error", _handle_error)
 
@@ -115,6 +123,18 @@ def _before_cur_exec(
                     span.set_attribute(key, value)
 
         context._otel_span = span
+        if self.enable_commenter:
+            statement = statement + EngineTracer._generate_comment(span=span)
+
+        return statement, params
+
+    @staticmethod
+    def _generate_comment(span: Span) -> str:
+        span_context = span.get_span_context()
+        meta = {}
+        if span_context.is_valid:
+            meta.update(_generate_opentelemetry_traceparent(span))
+        return _generate_sql_comment(**meta)
 
 
 # pylint: disable=unused-argument

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

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import asyncio
+import logging
 from unittest import mock
 
 import pytest
@@ -20,12 +21,17 @@
 
 from opentelemetry import trace
 from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+from opentelemetry.instrumentation.sqlalchemy.engine import EngineTracer
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider, export
 from opentelemetry.test.test_base import TestBase
 
 
 class TestSqlalchemyInstrumentation(TestBase):
+    @pytest.fixture(autouse=True)
+    def inject_fixtures(self, caplog):
+        self.caplog = caplog  # pylint: disable=attribute-defined-outside-init
+
     def tearDown(self):
         super().tearDown()
         SQLAlchemyInstrumentor().uninstrument()
@@ -150,3 +156,22 @@ async def run():
             self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
 
         asyncio.get_event_loop().run_until_complete(run())
+
+    def test_generate_commenter(self):
+        logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+        engine = create_engine("sqlite:///:memory:")
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+            enable_commenter=True,
+        )
+
+        cnx = engine.connect()
+        cnx.execute("SELECT	1 + 1;").fetchall()
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertIn(
+            EngineTracer._generate_comment(span),
+            self.caplog.records[-2].getMessage(),
+        )

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

@@ -0,0 +1,53 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from sqlalchemy import create_engine
+
+from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+from opentelemetry.test.test_base import TestBase
+
+
+class TestSqlalchemyInstrumentationWithSQLCommenter(TestBase):
+    @pytest.fixture(autouse=True)
+    def inject_fixtures(self, caplog):
+        self.caplog = caplog  # pylint: disable=attribute-defined-outside-init
+
+    def tearDown(self):
+        super().tearDown()
+        SQLAlchemyInstrumentor().uninstrument()
+
+    def test_sqlcommenter_enabled(self):
+        engine = create_engine("sqlite:///:memory:")
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+            enable_commenter=True,
+        )
+        cnx = engine.connect()
+        cnx.execute("SELECT  1;").fetchall()
+        self.assertRegex(
+            self.caplog.records[-2].getMessage(),
+            r"SELECT  1; /\*traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/",
+        )
+
+    def test_sqlcommenter_disabled(self):
+        engine = create_engine("sqlite:///:memory:", echo=True)
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine, tracer_provider=self.tracer_provider
+        )
+        cnx = engine.connect()
+        cnx.execute("SELECT 1;").fetchall()
+
+        self.assertEqual(self.caplog.records[-2].getMessage(), "SELECT 1;")

opentelemetry-instrumentation/src/opentelemetry/instrumentation/utils.py

@@ -23,7 +23,7 @@
 # pylint: disable=E0611
 from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY  # noqa: F401
 from opentelemetry.propagate import extract
-from opentelemetry.trace import StatusCode
+from opentelemetry.trace import Span, StatusCode
 
 
 def extract_attributes_from_object(
@@ -152,3 +152,14 @@ def _url_quote(s):  # pylint: disable=invalid-name
     # thus in our quoting, we need to escape it too to finally give
     #      foo,bar --> foo%%2Cbar
     return quoted.replace("%", "%%")
+
+
+def _generate_opentelemetry_traceparent(span: Span) -> str:
+    meta = {}
+    _version = "00"
+    _span_id = trace.format_span_id(span.context.span_id)
+    _trace_id = trace.format_trace_id(span.context.trace_id)
+    _flags = str(trace.TraceFlags.SAMPLED)
+    _traceparent = _version + "-" + _trace_id + "-" + _span_id + "-" + _flags
+    meta.update({"traceparent": _traceparent})
+    return meta