Add logfire.suppress_instrumentation context manager, silence urlli…

…b3 debug logs from exporting (#197)

Co-authored-by: Alex Hall <[email protected]>

logfire/__init__.py

@@ -16,6 +16,7 @@
 from ._internal.exporters.file import load_file as load_spans_from_file
 from ._internal.main import Logfire, LogfireSpan
 from ._internal.scrubbing import ScrubMatch
+from ._internal.utils import suppress_instrumentation
 from .integrations.logging import LogfireLoggingHandler
 from .integrations.structlog import LogfireProcessor as StructlogProcessor
 from .version import VERSION
@@ -101,6 +102,7 @@ def loguru_handler() -> dict[str, Any]:
     'METRICS_PREFERRED_TEMPORALITY',
     'ScrubMatch',
     'VERSION',
+    'suppress_instrumentation',
     'StructlogProcessor',
     'LogfireLoggingHandler',
 )

logfire/_internal/config.py

@@ -17,7 +17,6 @@
 
 import requests
 from opentelemetry import metrics, trace
-from opentelemetry.context import attach, detach, set_value
 from opentelemetry.environment_variables import OTEL_TRACES_EXPORTER
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
@@ -57,7 +56,6 @@
     DEFAULT_FALLBACK_FILE_NAME,
     OTLP_MAX_BODY_SIZE,
     RESOURCE_ATTRIBUTES_PACKAGE_VERSIONS,
-    SUPPRESS_INSTRUMENTATION_CONTEXT_KEY,
     LevelName,
 )
 from .exporters.console import (
@@ -75,7 +73,7 @@
 from .metrics import ProxyMeterProvider, configure_metrics
 from .scrubbing import Scrubber, ScrubCallback
 from .tracer import PendingSpanProcessor, ProxyTracerProvider
-from .utils import UnexpectedResponse, ensure_data_dir_exists, get_version, read_toml_file
+from .utils import UnexpectedResponse, ensure_data_dir_exists, get_version, read_toml_file, suppress_instrumentation
 
 CREDENTIALS_FILENAME = 'logfire_credentials.json'
 """Default base URL for the Logfire API."""
@@ -551,8 +549,7 @@ def _initialize(self) -> ProxyTracerProvider:
         if self._initialized:  # pragma: no cover
             return self._tracer_provider
 
-        backup_context = attach(set_value(SUPPRESS_INSTRUMENTATION_CONTEXT_KEY, True))
-        try:
+        with suppress_instrumentation():
             otel_resource_attributes: dict[str, Any] = {
                 ResourceAttributes.SERVICE_NAME: self.service_name,
                 RESOURCE_ATTRIBUTES_PACKAGE_VERSIONS: json.dumps(collect_package_info(), separators=(',', ':')),
@@ -703,8 +700,6 @@ def add_span_processor(span_processor: SpanProcessor) -> None:
             # set up context propagation for ThreadPoolExecutor and ProcessPoolExecutor
             instrument_executors()
             return self._tracer_provider
-        finally:
-            detach(backup_context)
 
     def get_tracer_provider(self) -> ProxyTracerProvider:
         """Get a tracer provider from this `LogfireConfig`.

logfire/_internal/constants.py

@@ -148,10 +148,6 @@ def log_level_attributes(level: LevelName | int) -> dict[str, otel_types.Attribu
 DEFAULT_FALLBACK_FILE_NAME = 'logfire_spans.bin'
 """The default name of the fallback file, used when the API is unreachable."""
 
-# see https://github.com/open-telemetry/opentelemetry-python/blob/d054dff47d2da663a39b9656d106c3d15f344269/opentelemetry-api/src/opentelemetry/context/__init__.py#L171
-SUPPRESS_INSTRUMENTATION_CONTEXT_KEY = 'suppress_instrumentation'
-"""Key in OTEL context that indicates whether instrumentation should be suppressed."""
-
 ATTRIBUTES_SAMPLE_RATE_KEY = 'logfire.sample_rate'
 """Key in attributes that indicates the sample rate for this span."""
 

logfire/_internal/exporters/processor_wrapper.py

@@ -15,7 +15,7 @@
     log_level_attributes,
 )
 from ..scrubbing import Scrubber
-from ..utils import ReadableSpanDict, span_to_dict, truncate_string
+from ..utils import ReadableSpanDict, is_instrumentation_suppressed, span_to_dict, truncate_string
 
 
 class SpanProcessorWrapper(SpanProcessor):
@@ -34,13 +34,13 @@ def on_start(
         span: Span,
         parent_context: context.Context | None = None,
     ) -> None:
-        if context.get_value('suppress_instrumentation'):  # pragma: no cover
+        if is_instrumentation_suppressed():
             return
         _set_log_level_on_asgi_send_receive_spans(span)
         self.processor.on_start(span, parent_context)
 
     def on_end(self, span: ReadableSpan) -> None:
-        if context.get_value('suppress_instrumentation'):  # pragma: no cover
+        if is_instrumentation_suppressed():
             return
         span_dict = span_to_dict(span)
         _tweak_asgi_send_receive_spans(span_dict)

logfire/_internal/integrations/llm_providers/llm_provider.py

@@ -4,9 +4,8 @@
 from contextlib import ExitStack, contextmanager, nullcontext
 from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, ContextManager, Iterator, cast
 
-from opentelemetry import context
-
 from ...constants import ONE_SECOND_IN_NANOSECONDS
+from ...utils import is_instrumentation_suppressed, suppress_instrumentation
 
 if TYPE_CHECKING:
     from ...main import Logfire, LogfireSpan
@@ -74,7 +73,7 @@ def uninstrument_context():
     is_async = is_async_client_fn(client if isinstance(client, type) else type(client))
 
     def _instrumentation_setup(**kwargs: Any) -> Any:
-        if context.get_value('suppress_instrumentation'):
+        if is_instrumentation_suppressed():
             return None, None, kwargs
 
         options = kwargs['options']
@@ -172,12 +171,8 @@ def uninstrument_context():
 @contextmanager
 def maybe_suppress_instrumentation(suppress: bool) -> Iterator[None]:
     if suppress:
-        new_context = context.set_value('suppress_instrumentation', True)
-        token = context.attach(new_context)
-        try:
+        with suppress_instrumentation():
             yield
-        finally:
-            context.detach(token)
     else:
         yield
 

logfire/_internal/utils.py

@@ -2,10 +2,11 @@
 
 import json
 import sys
+from contextlib import contextmanager
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Sequence, Tuple, TypedDict, TypeVar, Union
 
-from opentelemetry import trace as trace_api
+from opentelemetry import context, trace as trace_api
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import Event, ReadableSpan
 from opentelemetry.sdk.util.instrumentation import InstrumentationScope
@@ -184,3 +185,37 @@ def get_version(version: str) -> Version:
     except ImportError:  # pragma: no cover
         from setuptools._vendor.packaging.version import Version
     return Version(version)  # type: ignore
+
+
+# OTEL uses two different keys to supress instrumentation. We need to check both.
+SUPPRESS_INSTRUMENTATION_CONTEXT_KEYS = [
+    # This is still used in some places in OTEL, and probably more in older versions.
+    'suppress_instrumentation',
+]
+
+try:
+    # This is the 'main' key used by OTEL in recent versions
+    SUPPRESS_INSTRUMENTATION_CONTEXT_KEYS.append(context._SUPPRESS_INSTRUMENTATION_KEY)  # type: ignore
+except AttributeError:  # pragma: no cover
+    pass
+
+
+def is_instrumentation_suppressed() -> bool:
+    """Return True if the `suppress_instrumentation` context manager is currently active.
+
+    This means that any logs/spans generated by logfire or OpenTelemetry will not be logged in any way.
+    """
+    return any(context.get_value(key) for key in SUPPRESS_INSTRUMENTATION_CONTEXT_KEYS)
+
+
+@contextmanager
+def suppress_instrumentation():
+    """Context manager to suppress all logs/spans generated by logfire or OpenTelemetry."""
+    new_context = context.get_current()
+    for key in SUPPRESS_INSTRUMENTATION_CONTEXT_KEYS:
+        new_context = context.set_value(key, True, new_context)
+    token = context.attach(new_context)
+    try:
+        yield
+    finally:
+        context.detach(token)

logfire/integrations/logging.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import inspect
 from logging import NOTSET, Handler as LoggingHandler, LogRecord, StreamHandler
 from typing import Any, ClassVar, Mapping, cast
 
@@ -14,6 +13,7 @@
     ATTRIBUTES_MESSAGE_TEMPLATE_KEY,
     LOGGING_TO_OTEL_LEVEL_NUMBERS,
 )
+from .._internal.utils import is_instrumentation_suppressed
 
 # skip natural LogRecord attributes
 # http://docs.python.org/library/logging.html#logrecord-attributes
@@ -61,20 +61,9 @@ def emit(self, record: LogRecord) -> None:
         Args:
             record: The log record to send.
         """
-        if record.name.startswith('opentelemetry.'):
-            # This method can lead to OTEL calling logging methods which recursively calls this again.
-            # If we detect recursion, use the fallback handler instead.
-            # TODO find a better way to handle this,
-            #  or document the fallback clearly and nudge the user to configure it.
-            frame = inspect.currentframe()
-            assert frame is not None
-            code_here = frame.f_code
-            frame = frame.f_back
-            while frame:
-                if frame.f_code is code_here:
-                    self.fallback.emit(record)
-                    return
-                frame = frame.f_back
+        if is_instrumentation_suppressed():
+            self.fallback.handle(record)
+            return
 
         attributes = self.fill_attributes(record)
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "logfire"
-version = "0.36.1"
+version = "0.37.0"
 description = "The best Python observability tool! 🪵🔥"
 authors = [
     { name = "Pydantic Team", email = "[email protected]" },

tests/otel_integrations/test_openai.py

@@ -23,11 +23,11 @@
     images_response,
 )
 from openai.types.chat import chat_completion, chat_completion_chunk as cc_chunk, chat_completion_message
-from opentelemetry import context
 from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
 
 import logfire
 from logfire._internal.integrations.llm_providers.openai import get_endpoint_config
+from logfire._internal.utils import suppress_instrumentation
 from logfire.testing import TestExporter
 
 
@@ -845,25 +845,15 @@ def test_suppress_httpx(exporter: TestExporter) -> None:
 
 
 def test_openai_suppressed(instrumented_client: openai.Client, exporter: TestExporter) -> None:
-    new_context = context.set_value('suppress_instrumentation', True)
-    token = context.attach(new_context)
-    try:
+    with suppress_instrumentation():
         response = instrumented_client.completions.create(model='gpt-3.5-turbo-instruct', prompt='xxx')
-    finally:
-        context.detach(token)
-
     assert response.choices[0].text == 'Nine'
     assert exporter.exported_spans_as_dict() == []
 
 
 async def test_async_openai_suppressed(instrumented_async_client: openai.AsyncClient, exporter: TestExporter) -> None:
-    new_context = context.set_value('suppress_instrumentation', True)
-    token = context.attach(new_context)
-    try:
+    with suppress_instrumentation():
         response = await instrumented_async_client.completions.create(model='gpt-3.5-turbo-instruct', prompt='xxx')
-    finally:
-        context.detach(token)
-
     assert response.choices[0].text == 'Nine'
     assert exporter.exported_spans_as_dict() == []
 

tests/test_logfire.py

@@ -20,7 +20,7 @@
 from pydantic_core import ValidationError
 
 import logfire
-from logfire import Logfire
+from logfire import Logfire, suppress_instrumentation
 from logfire._internal.config import LogfireConfig, configure
 from logfire._internal.constants import (
     ATTRIBUTES_MESSAGE_KEY,
@@ -31,6 +31,7 @@
     NULL_ARGS_KEY,
 )
 from logfire._internal.formatter import InspectArgumentsFailedWarning
+from logfire._internal.utils import is_instrumentation_suppressed
 from logfire.integrations.logging import LogfireLoggingHandler
 from logfire.testing import IncrementalIdGenerator, TestExporter, TimeGenerator
 
@@ -2454,3 +2455,51 @@ def test_wrong_fstring_source_segment(exporter: TestExporter):
             }
         ]
     )
+
+
+def test_suppress_instrumentation(exporter: TestExporter):
+    logfire.info('log1')
+    assert not is_instrumentation_suppressed()
+    with suppress_instrumentation():
+        assert is_instrumentation_suppressed()
+        # Not included in the asserted spans below
+        logfire.info('log2')
+    assert not is_instrumentation_suppressed()
+    logfire.info('log3')
+
+    assert exporter.exported_spans_as_dict() == snapshot(
+        [
+            {
+                'name': 'log1',
+                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                'parent': None,
+                'start_time': 1000000000,
+                'end_time': 1000000000,
+                'attributes': {
+                    'logfire.span_type': 'log',
+                    'logfire.level_num': 9,
+                    'logfire.msg_template': 'log1',
+                    'logfire.msg': 'log1',
+                    'code.filepath': 'test_logfire.py',
+                    'code.function': 'test_suppress_instrumentation',
+                    'code.lineno': 123,
+                },
+            },
+            {
+                'name': 'log3',
+                'context': {'trace_id': 3, 'span_id': 3, 'is_remote': False},
+                'parent': None,
+                'start_time': 3000000000,
+                'end_time': 3000000000,
+                'attributes': {
+                    'logfire.span_type': 'log',
+                    'logfire.level_num': 9,
+                    'logfire.msg_template': 'log3',
+                    'logfire.msg': 'log3',
+                    'code.filepath': 'test_logfire.py',
+                    'code.function': 'test_suppress_instrumentation',
+                    'code.lineno': 123,
+                },
+            },
+        ]
+    )

tests/test_stdlib_logging.py

@@ -212,9 +212,10 @@ def __init__(self):
 
     try:
         # This calls ExceptionExporter.export which causes OTEL to log an exception.
-        # That log call goes to LogfireLoggingHandler.emit, which tries to emit another logfire log,
+        # That log call goes to LogfireLoggingHandler.emit, which usually tries to emit another logfire log,
         # causing another stdlib log from OTEL, potentially leading to infinite recursion.
-        # This tests that we handle that by using the fallback test_logging_handler when recursion is detected.
+        # Recursion is prevented by OTEL suppressing instrumentation, so the second logfire log isn't emitted.
+        # But when we detect this, we use the fallback handler instead, so this tests that.
         logfire.info('test')
     finally:
         # Don't mess with the root logger longer than needed.