sqlalchemy wrap_create_engine now accepts sqlcommenter options (#1873)

* sqlalchemy wrap_create_engine accepts sqlcommenter options * Changelog * Lint * Fix default val * Add sqlalchemy tests * Change a default in _instrument get * Lint * More lint * Update default Co-authored-by: Shalev Roda <[email protected]> * Update args doc * lintttt --------- Co-authored-by: Shalev Roda <[email protected]>
open-telemetry · Jun 27, 2023 · 79d62b3 · 79d62b3
1 parent 2e49ba1
commit 79d62b3
Show file tree

Hide file tree

Showing 5 changed files with 146 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#1870](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1870))
 - Update falcon instrumentation to follow semantic conventions
   ([#1824](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1824))
+- Fix sqlalchemy instrumentation wrap methods to accept sqlcommenter options([#1873](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1873))
 
 ### Added
 

diff --git a/...metry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py b/...metry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py
@@ -134,6 +134,9 @@ def _instrument(self, **kwargs):
                 ``engine``: a SQLAlchemy engine instance
                 ``engines``: a list of SQLAlchemy engine instances
                 ``tracer_provider``: a TracerProvider, defaults to global
+                ``meter_provider``: a MeterProvider, defaults to global
+                ``enable_commenter``: bool to enable sqlcommenter, defaults to False
+                ``commenter_options``: dict of sqlcommenter config, defaults to None
 
         Returns:
             An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise.
@@ -151,16 +154,21 @@ def _instrument(self, **kwargs):
         )
 
         enable_commenter = kwargs.get("enable_commenter", False)
+        commenter_options = kwargs.get("commenter_options", {})
 
         _w(
             "sqlalchemy",
             "create_engine",
-            _wrap_create_engine(tracer, connections_usage, enable_commenter),
+            _wrap_create_engine(
+                tracer, connections_usage, enable_commenter, commenter_options
+            ),
         )
         _w(
             "sqlalchemy.engine",
             "create_engine",
-            _wrap_create_engine(tracer, connections_usage, enable_commenter),
+            _wrap_create_engine(
+                tracer, connections_usage, enable_commenter, commenter_options
+            ),
         )
         _w(
             "sqlalchemy.engine.base",
@@ -172,7 +180,10 @@ def _instrument(self, **kwargs):
                 "sqlalchemy.ext.asyncio",
                 "create_async_engine",
                 _wrap_create_async_engine(
-                    tracer, connections_usage, enable_commenter
+                    tracer,
+                    connections_usage,
+                    enable_commenter,
+                    commenter_options,
                 ),
             )
         if kwargs.get("engine") is not None:

diff --git a/...lemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py b/...lemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py
@@ -43,7 +43,7 @@ def _normalize_vendor(vendor):
 
 
 def _wrap_create_async_engine(
-    tracer, connections_usage, enable_commenter=False
+    tracer, connections_usage, enable_commenter=False, commenter_options=None
 ):
     # pylint: disable=unused-argument
     def _wrap_create_async_engine_internal(func, module, args, kwargs):
@@ -52,20 +52,32 @@ def _wrap_create_async_engine_internal(func, module, args, kwargs):
         """
         engine = func(*args, **kwargs)
         EngineTracer(
-            tracer, engine.sync_engine, connections_usage, enable_commenter
+            tracer,
+            engine.sync_engine,
+            connections_usage,
+            enable_commenter,
+            commenter_options,
         )
         return engine
 
     return _wrap_create_async_engine_internal
 
 
-def _wrap_create_engine(tracer, connections_usage, enable_commenter=False):
+def _wrap_create_engine(
+    tracer, connections_usage, enable_commenter=False, commenter_options=None
+):
     def _wrap_create_engine_internal(func, _module, args, kwargs):
         """Trace the SQLAlchemy engine, creating an `EngineTracer`
         object that will listen to SQLAlchemy events.
         """
         engine = func(*args, **kwargs)
-        EngineTracer(tracer, engine, connections_usage, enable_commenter)
+        EngineTracer(
+            tracer,
+            engine,
+            connections_usage,
+            enable_commenter,
+            commenter_options,
+        )
         return engine
 
     return _wrap_create_engine_internal

diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py b/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
@@ -176,6 +177,43 @@ def test_create_engine_wrapper(self):
             "opentelemetry.instrumentation.sqlalchemy",
         )
 
+    def test_create_engine_wrapper_enable_commenter(self):
+        logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+        SQLAlchemyInstrumentor().instrument(
+            enable_commenter=True,
+            commenter_options={"db_framework": False},
+        )
+        from sqlalchemy import create_engine  # pylint: disable-all
+
+        engine = create_engine("sqlite:///:memory:")
+        cnx = engine.connect()
+        cnx.execute("SELECT  1;").fetchall()
+        # sqlcommenter
+        self.assertRegex(
+            self.caplog.records[-2].getMessage(),
+            r"SELECT  1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+        )
+
+    def test_create_engine_wrapper_enable_commenter_otel_values_false(self):
+        logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+        SQLAlchemyInstrumentor().instrument(
+            enable_commenter=True,
+            commenter_options={
+                "db_framework": False,
+                "opentelemetry_values": False,
+            },
+        )
+        from sqlalchemy import create_engine  # pylint: disable-all
+
+        engine = create_engine("sqlite:///:memory:")
+        cnx = engine.connect()
+        cnx.execute("SELECT  1;").fetchall()
+        # sqlcommenter
+        self.assertRegex(
+            self.caplog.records[-2].getMessage(),
+            r"SELECT  1 /\*db_driver='(.*)'\*/;",
+        )
+
     def test_custom_tracer_provider(self):
         provider = TracerProvider(
             resource=Resource.create(
@@ -242,6 +280,65 @@ async def run():
 
         asyncio.get_event_loop().run_until_complete(run())
 
+    @pytest.mark.skipif(
+        not sqlalchemy.__version__.startswith("1.4"),
+        reason="only run async tests for 1.4",
+    )
+    def test_create_async_engine_wrapper_enable_commenter(self):
+        async def run():
+            logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+            SQLAlchemyInstrumentor().instrument(
+                enable_commenter=True,
+                commenter_options={
+                    "db_framework": False,
+                },
+            )
+            from sqlalchemy.ext.asyncio import (  # pylint: disable-all
+                create_async_engine,
+            )
+
+            engine = create_async_engine("sqlite+aiosqlite:///:memory:")
+            async with engine.connect() as cnx:
+                await cnx.execute(sqlalchemy.text("SELECT  1;"))
+            # sqlcommenter
+            self.assertRegex(
+                self.caplog.records[1].getMessage(),
+                r"SELECT  1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
+            )
+
+        asyncio.get_event_loop().run_until_complete(run())
+
+    @pytest.mark.skipif(
+        not sqlalchemy.__version__.startswith("1.4"),
+        reason="only run async tests for 1.4",
+    )
+    def test_create_async_engine_wrapper_enable_commenter_otel_values_false(
+        self,
+    ):
+        async def run():
+            logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+            SQLAlchemyInstrumentor().instrument(
+                enable_commenter=True,
+                commenter_options={
+                    "db_framework": False,
+                    "opentelemetry_values": False,
+                },
+            )
+            from sqlalchemy.ext.asyncio import (  # pylint: disable-all
+                create_async_engine,
+            )
+
+            engine = create_async_engine("sqlite+aiosqlite:///:memory:")
+            async with engine.connect() as cnx:
+                await cnx.execute(sqlalchemy.text("SELECT  1;"))
+            # sqlcommenter
+            self.assertRegex(
+                self.caplog.records[1].getMessage(),
+                r"SELECT  1 /\*db_driver='(.*)'\*/;",
+            )
+
+        asyncio.get_event_loop().run_until_complete(run())
+
     def test_uninstrument(self):
         engine = create_engine("sqlite:///:memory:")
         SQLAlchemyInstrumentor().instrument(

diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlcommenter.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlcommenter.py
@@ -56,6 +56,24 @@ 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_otel_values_false(self):
+        engine = create_engine("sqlite:///:memory:")
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+            enable_commenter=True,
+            commenter_options={
+                "db_framework": False,
+                "opentelemetry_values": False,
+            },
+        )
+        cnx = engine.connect()
+        cnx.execute("SELECT  1;").fetchall()
+        self.assertRegex(
+            self.caplog.records[-2].getMessage(),
+            r"SELECT  1 /\*db_driver='(.*)'\*/;",
+        )
+
     def test_sqlcommenter_flask_integration(self):
         engine = create_engine("sqlite:///:memory:")
         SQLAlchemyInstrumentor().instrument(