respect supress_instrumentation

Author: marcm-ml

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

@@ -23,7 +23,10 @@
 
 from opentelemetry import trace
 from opentelemetry.instrumentation.sqlcommenter_utils import _add_sql_comment
-from opentelemetry.instrumentation.utils import _get_opentelemetry_values
+from opentelemetry.instrumentation.utils import (
+    _get_opentelemetry_values,
+    is_instrumentation_enabled,
+)
 from opentelemetry.semconv.trace import NetTransportValues, SpanAttributes
 from opentelemetry.trace.status import Status, StatusCode
 
@@ -54,6 +57,9 @@ def _wrap_create_async_engine_internal(func, module, args, kwargs):
         """Trace the SQLAlchemy engine, creating an `EngineTracer`
         object that will listen to SQLAlchemy events.
         """
+        if not is_instrumentation_enabled():
+            return func(*args, **kwargs)
+
         engine = func(*args, **kwargs)
         EngineTracer(
             tracer,
@@ -79,6 +85,9 @@ def _wrap_create_engine_internal(func, _module, args, kwargs):
         """Trace the SQLAlchemy engine, creating an `EngineTracer`
         object that will listen to SQLAlchemy events.
         """
+        if not is_instrumentation_enabled():
+            return func(*args, **kwargs)
+
         engine = func(*args, **kwargs)
         EngineTracer(
             tracer,
@@ -96,6 +105,9 @@ def _wrap_create_engine_internal(func, _module, args, kwargs):
 def _wrap_connect(tracer):
     # pylint: disable=unused-argument
     def _wrap_connect_internal(func, module, args, kwargs):
+        if not is_instrumentation_enabled():
+            return func(*args, **kwargs)
+
         with tracer.start_as_current_span(
             "connect", kind=trace.SpanKind.CLIENT
         ) as span:
@@ -144,6 +156,9 @@ def __init__(
         self._register_event_listener(engine, "checkout", self._pool_checkout)
 
     def _add_idle_to_connection_usage(self, value):
+        if not is_instrumentation_enabled():
+            return
+
         self.connections_usage.add(
             value,
             attributes={
@@ -153,6 +168,9 @@ def _add_idle_to_connection_usage(self, value):
         )
 
     def _add_used_to_connection_usage(self, value):
+        if not is_instrumentation_enabled():
+            return
+
         self.connections_usage.add(
             value,
             attributes={
@@ -259,6 +277,9 @@ def _set_db_client_span_attributes(self, span, statement, attrs) -> None:
     def _before_cur_exec(
         self, conn, cursor, statement, params, context, _executemany
     ):
+        if not is_instrumentation_enabled():
+            return statement, params
+
         attrs, found = _get_attributes_from_url(conn.engine.url)
         if not found:
             attrs = _get_attributes_from_cursor(self.vendor, cursor, attrs)

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

@@ -23,7 +23,11 @@
 )
 
 from opentelemetry import trace
-from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+from opentelemetry.instrumentation.sqlalchemy import (
+    EngineTracer,
+    SQLAlchemyInstrumentor,
+)
+from opentelemetry.instrumentation.utils import suppress_instrumentation
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider, export
 from opentelemetry.semconv.trace import SpanAttributes
@@ -630,3 +634,63 @@ def make_shortlived_engine():
         gc.collect()
         assert callback.call_count == 5
         assert len(EngineTracer._remove_event_listener_params) == 0
+
+    def test_suppress_instrumentation_create_engine(self):
+        SQLAlchemyInstrumentor().instrument()
+
+        from sqlalchemy import create_engine
+
+        with suppress_instrumentation():
+            engine = create_engine("sqlite:///:memory:")
+
+        self.assertTrue(not isinstance(engine, EngineTracer))
+
+    @pytest.mark.skipif(
+        not sqlalchemy.__version__.startswith("1.4"),
+        reason="only run async tests for 1.4",
+    )
+    def test_suppress_instrumentation_create_async_engine(self):
+        async def run():
+            SQLAlchemyInstrumentor().instrument()
+            from sqlalchemy.ext.asyncio import (  # pylint: disable-all
+                create_async_engine,
+            )
+
+            with suppress_instrumentation():
+                engine = create_async_engine("sqlite+aiosqlite:///:memory:")
+
+            self.assertTrue(not isinstance(engine, EngineTracer))
+
+        asyncio.get_event_loop().run_until_complete(run())
+
+    def test_suppress_instrumentation_connect(self):
+        engine = create_engine("sqlite:///:memory:")
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+        )
+
+        with suppress_instrumentation():
+            with engine.connect():
+                pass
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 0)
+
+    def test_suppress_instrumentation_cursor_and_metric(self):
+        engine = create_engine("sqlite:///:memory:")
+        SQLAlchemyInstrumentor().instrument(
+            engine=engine,
+            tracer_provider=self.tracer_provider,
+            enable_commenter=True,
+        )
+
+        with suppress_instrumentation():
+            with engine.connect() as conn:
+                conn.execute(text("SELECT 1 + 1;")).fetchall()
+
+        spans_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans_list), 0)
+
+        metric_list = self.get_sorted_metrics()
+        self.assertEqual(len(metric_list), 0)