PYD-1066: Remove default_span_processor parameter from configure (#400)

alexmojaki · web-flow · commit fdfe06ddd973 · 2024-08-26T13:04:55.000Z
diff --git a/logfire-api/logfire_api/_internal/config.pyi b/logfire-api/logfire_api/_internal/config.pyi
@@ -27,7 +27,6 @@ from logfire.version import VERSION as VERSION
 from opentelemetry import metrics
 from opentelemetry.sdk.metrics.export import MetricReader as MetricReader
 from opentelemetry.sdk.trace import SpanProcessor
-from opentelemetry.sdk.trace.export import SpanExporter
 from opentelemetry.sdk.trace.id_generator import IdGenerator
 from pathlib import Path
 from typing import Any, Callable, Literal, Sequence
@@ -56,7 +55,7 @@ class PydanticPlugin:
     include: set[str] = ...
     exclude: set[str] = ...
 
-def configure(*, send_to_logfire: bool | Literal['if-token-present'] | None = None, token: str | None = None, project_name: str | None = None, service_name: str | None = None, service_version: str | None = None, trace_sample_rate: float | None = None, console: ConsoleOptions | Literal[False] | None = None, show_summary: bool | None = None, config_dir: Path | str | None = None, data_dir: Path | str | None = None, base_url: str | None = None, collect_system_metrics: None = None, id_generator: IdGenerator | None = None, ns_timestamp_generator: Callable[[], int] | None = None, processors: None = None, additional_span_processors: Sequence[SpanProcessor] | None = None, default_span_processor: Callable[[SpanExporter], SpanProcessor] | None = None, metric_readers: None = None, additional_metric_readers: Sequence[MetricReader] | None = None, pydantic_plugin: PydanticPlugin | None = None, fast_shutdown: bool = False, scrubbing_patterns: Sequence[str] | None = None, scrubbing_callback: ScrubCallback | None = None, scrubbing: ScrubbingOptions | Literal[False] | None = None, inspect_arguments: bool | None = None, tail_sampling: TailSamplingOptions | None = None) -> None:
+def configure(*, send_to_logfire: bool | Literal['if-token-present'] | None = None, token: str | None = None, project_name: str | None = None, service_name: str | None = None, service_version: str | None = None, trace_sample_rate: float | None = None, console: ConsoleOptions | Literal[False] | None = None, show_summary: bool | None = None, config_dir: Path | str | None = None, data_dir: Path | str | None = None, base_url: str | None = None, collect_system_metrics: None = None, id_generator: IdGenerator | None = None, ns_timestamp_generator: Callable[[], int] | None = None, processors: None = None, additional_span_processors: Sequence[SpanProcessor] | None = None, metric_readers: None = None, additional_metric_readers: Sequence[MetricReader] | None = None, pydantic_plugin: PydanticPlugin | None = None, fast_shutdown: bool = False, scrubbing_patterns: Sequence[str] | None = None, scrubbing_callback: ScrubCallback | None = None, scrubbing: ScrubbingOptions | Literal[False] | None = None, inspect_arguments: bool | None = None, tail_sampling: TailSamplingOptions | None = None) -> None:
     """Configure the logfire SDK.
 
     Args:
@@ -89,9 +88,6 @@ def configure(*, send_to_logfire: bool | Literal['if-token-present'] | None = No
             Python standard library.
         processors: Legacy argument, use `additional_span_processors` instead.
         additional_span_processors: Span processors to use in addition to the default processor which exports spans to Logfire's API.
-        default_span_processor: A function to create the default span processor. Defaults to `BatchSpanProcessor` from the OpenTelemetry SDK. You can configure the export delay for
-            [`BatchSpanProcessor`](https://opentelemetry-python.readthedocs.io/en/latest/sdk/trace.export.html#opentelemetry.sdk.trace.export.BatchSpanProcessor)
-            by setting the `OTEL_BSP_SCHEDULE_DELAY_MILLIS` environment variable.
         metric_readers: Legacy argument, use `additional_metric_readers` instead.
         additional_metric_readers: Sequence of metric readers to be used in addition to the default reader
             which exports metrics to Logfire's API.
@@ -133,21 +129,20 @@ class _LogfireConfigData:
     ns_timestamp_generator: Callable[[], int]
     additional_span_processors: Sequence[SpanProcessor] | None
     pydantic_plugin: PydanticPlugin
-    default_span_processor: Callable[[SpanExporter], SpanProcessor]
     fast_shutdown: bool
     scrubbing: ScrubbingOptions | Literal[False]
     inspect_arguments: bool
     tail_sampling: TailSamplingOptions | None
 
 class LogfireConfig(_LogfireConfigData):
-    def __init__(self, base_url: str | None = None, send_to_logfire: bool | None = None, token: str | None = None, project_name: str | None = None, service_name: str | None = None, service_version: str | None = None, trace_sample_rate: float | None = None, console: ConsoleOptions | Literal[False] | None = None, show_summary: bool | None = None, config_dir: Path | None = None, data_dir: Path | None = None, id_generator: IdGenerator | None = None, ns_timestamp_generator: Callable[[], int] | None = None, additional_span_processors: Sequence[SpanProcessor] | None = None, default_span_processor: Callable[[SpanExporter], SpanProcessor] | None = None, additional_metric_readers: Sequence[MetricReader] | None = None, pydantic_plugin: PydanticPlugin | None = None, fast_shutdown: bool = False, scrubbing: ScrubbingOptions | Literal[False] | None = None, inspect_arguments: bool | None = None, tail_sampling: TailSamplingOptions | None = None) -> None:
+    def __init__(self, base_url: str | None = None, send_to_logfire: bool | None = None, token: str | None = None, project_name: str | None = None, service_name: str | None = None, service_version: str | None = None, trace_sample_rate: float | None = None, console: ConsoleOptions | Literal[False] | None = None, show_summary: bool | None = None, config_dir: Path | None = None, data_dir: Path | None = None, id_generator: IdGenerator | None = None, ns_timestamp_generator: Callable[[], int] | None = None, additional_span_processors: Sequence[SpanProcessor] | None = None, additional_metric_readers: Sequence[MetricReader] | None = None, pydantic_plugin: PydanticPlugin | None = None, fast_shutdown: bool = False, scrubbing: ScrubbingOptions | Literal[False] | None = None, inspect_arguments: bool | None = None, tail_sampling: TailSamplingOptions | None = None) -> None:
         """Create a new LogfireConfig.
 
         Users should never need to call this directly, instead use `logfire.configure`.
 
         See `_LogfireConfigData` for parameter documentation.
         """
-    def configure(self, base_url: str | None, send_to_logfire: bool | Literal['if-token-present'] | None, token: str | None, project_name: str | None, service_name: str | None, service_version: str | None, trace_sample_rate: float | None, console: ConsoleOptions | Literal[False] | None, show_summary: bool | None, config_dir: Path | None, data_dir: Path | None, id_generator: IdGenerator | None, ns_timestamp_generator: Callable[[], int] | None, additional_span_processors: Sequence[SpanProcessor] | None, default_span_processor: Callable[[SpanExporter], SpanProcessor] | None, additional_metric_readers: Sequence[MetricReader] | None, pydantic_plugin: PydanticPlugin | None, fast_shutdown: bool, scrubbing: ScrubbingOptions | Literal[False] | None, inspect_arguments: bool | None, tail_sampling: TailSamplingOptions | None) -> None: ...
+    def configure(self, base_url: str | None, send_to_logfire: bool | Literal['if-token-present'] | None, token: str | None, project_name: str | None, service_name: str | None, service_version: str | None, trace_sample_rate: float | None, console: ConsoleOptions | Literal[False] | None, show_summary: bool | None, config_dir: Path | None, data_dir: Path | None, id_generator: IdGenerator | None, ns_timestamp_generator: Callable[[], int] | None, additional_span_processors: Sequence[SpanProcessor] | None, additional_metric_readers: Sequence[MetricReader] | None, pydantic_plugin: PydanticPlugin | None, fast_shutdown: bool, scrubbing: ScrubbingOptions | Literal[False] | None, inspect_arguments: bool | None, tail_sampling: TailSamplingOptions | None) -> None: ...
     def initialize(self) -> ProxyTracerProvider:
         """Configure internals to start exporting traces and metrics."""
     def force_flush(self, timeout_millis: int = 30000) -> bool:
diff --git a/logfire/_internal/config.py b/logfire/_internal/config.py
@@ -44,7 +44,7 @@
 from opentelemetry.sdk.metrics.view import ExponentialBucketHistogramAggregation, View
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import SpanProcessor, TracerProvider as SDKTracerProvider
-from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor, SpanExporter
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
 from opentelemetry.sdk.trace.id_generator import IdGenerator, RandomIdGenerator
 from opentelemetry.sdk.trace.sampling import ParentBasedTraceIdRatio
 from opentelemetry.semconv.resource import ResourceAttributes
@@ -162,7 +162,6 @@ def configure(
     ns_timestamp_generator: Callable[[], int] | None = None,
     processors: None = None,
     additional_span_processors: Sequence[SpanProcessor] | None = None,
-    default_span_processor: Callable[[SpanExporter], SpanProcessor] | None = None,
     metric_readers: None = None,
     additional_metric_readers: Sequence[MetricReader] | None = None,
     pydantic_plugin: PydanticPlugin | None = None,
@@ -205,9 +204,6 @@ def configure(
             Python standard library.
         processors: Legacy argument, use `additional_span_processors` instead.
         additional_span_processors: Span processors to use in addition to the default processor which exports spans to Logfire's API.
-        default_span_processor: A function to create the default span processor. Defaults to `BatchSpanProcessor` from the OpenTelemetry SDK. You can configure the export delay for
-            [`BatchSpanProcessor`](https://opentelemetry-python.readthedocs.io/en/latest/sdk/trace.export.html#opentelemetry.sdk.trace.export.BatchSpanProcessor)
-            by setting the `OTEL_BSP_SCHEDULE_DELAY_MILLIS` environment variable.
         metric_readers: Legacy argument, use `additional_metric_readers` instead.
         additional_metric_readers: Sequence of metric readers to be used in addition to the default reader
             which exports metrics to Logfire's API.
@@ -275,7 +271,6 @@ def configure(
         id_generator=id_generator,
         ns_timestamp_generator=ns_timestamp_generator,
         additional_span_processors=additional_span_processors,
-        default_span_processor=default_span_processor,
         additional_metric_readers=additional_metric_readers,
         pydantic_plugin=pydantic_plugin,
         fast_shutdown=fast_shutdown,
@@ -346,9 +341,6 @@ class _LogfireConfigData:
     pydantic_plugin: PydanticPlugin
     """Options for the Pydantic plugin"""
 
-    default_span_processor: Callable[[SpanExporter], SpanProcessor]
-    """The span processor used for the logfire exporter and console exporter"""
-
     fast_shutdown: bool
     """Whether to shut down exporters and providers quickly, mostly used for tests"""
 
@@ -380,7 +372,6 @@ def _load_configuration(
         id_generator: IdGenerator | None,
         ns_timestamp_generator: Callable[[], int] | None,
         additional_span_processors: Sequence[SpanProcessor] | None,
-        default_span_processor: Callable[[SpanExporter], SpanProcessor] | None,
         additional_metric_readers: Sequence[MetricReader] | None,
         pydantic_plugin: PydanticPlugin | None,
         fast_shutdown: bool,
@@ -454,7 +445,6 @@ def _load_configuration(
         self.id_generator = id_generator or RandomIdGenerator()
         self.ns_timestamp_generator = ns_timestamp_generator or time.time_ns
         self.additional_span_processors = additional_span_processors
-        self.default_span_processor = default_span_processor or _get_default_span_processor
         self.additional_metric_readers = additional_metric_readers
         if self.service_version is None:
             try:
@@ -482,7 +472,6 @@ def __init__(
         id_generator: IdGenerator | None = None,
         ns_timestamp_generator: Callable[[], int] | None = None,
         additional_span_processors: Sequence[SpanProcessor] | None = None,
-        default_span_processor: Callable[[SpanExporter], SpanProcessor] | None = None,
         additional_metric_readers: Sequence[MetricReader] | None = None,
         pydantic_plugin: PydanticPlugin | None = None,
         fast_shutdown: bool = False,
@@ -513,7 +502,6 @@ def __init__(
             id_generator=id_generator,
             ns_timestamp_generator=ns_timestamp_generator,
             additional_span_processors=additional_span_processors,
-            default_span_processor=default_span_processor,
             additional_metric_readers=additional_metric_readers,
             pydantic_plugin=pydantic_plugin,
             fast_shutdown=fast_shutdown,
@@ -548,7 +536,6 @@ def configure(
         id_generator: IdGenerator | None,
         ns_timestamp_generator: Callable[[], int] | None,
         additional_span_processors: Sequence[SpanProcessor] | None,
-        default_span_processor: Callable[[SpanExporter], SpanProcessor] | None,
         additional_metric_readers: Sequence[MetricReader] | None,
         pydantic_plugin: PydanticPlugin | None,
         fast_shutdown: bool,
@@ -573,7 +560,6 @@ def configure(
                 id_generator,
                 ns_timestamp_generator,
                 additional_span_processors,
-                default_span_processor,
                 additional_metric_readers,
                 pydantic_plugin,
                 fast_shutdown,
@@ -719,7 +705,8 @@ def check_token():
                     span_exporter, FileSpanExporter(self.data_dir / DEFAULT_FALLBACK_FILE_NAME, warn=True)
                 )
                 span_exporter = RemovePendingSpansExporter(span_exporter)
-                add_span_processor(self.default_span_processor(span_exporter))
+                schedule_delay_millis = _get_int_from_env(OTEL_BSP_SCHEDULE_DELAY) or 500
+                add_span_processor(BatchSpanProcessor(span_exporter, schedule_delay_millis=schedule_delay_millis))
 
                 metric_readers += [
                     PeriodicExportingMetricReader(
@@ -905,11 +892,6 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
                 traceback.print_exception(e)
 
 
-def _get_default_span_processor(exporter: SpanExporter) -> SpanProcessor:
-    schedule_delay_millis = _get_int_from_env(OTEL_BSP_SCHEDULE_DELAY) or 500
-    return BatchSpanProcessor(exporter, schedule_delay_millis=schedule_delay_millis)
-
-
 # The global config is the single global object in logfire
 # It also does not initialize anything when it's created (right now)
 # but when `logfire.configure` aka `GLOBAL_CONFIG.configure` is called
diff --git a/tests/test_configure.py b/tests/test_configure.py
@@ -7,6 +7,7 @@
 import threading
 from contextlib import ExitStack
 from pathlib import Path
+from time import sleep, time
 from typing import Any, Iterable, Sequence
 from unittest import mock
 from unittest.mock import call, patch
@@ -534,35 +535,98 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
             # This should cause FallbackSpanExporter to call its own fallback file exporter.
             return SpanExportResult.FAILURE
 
-    def default_span_processor(exporter: SpanExporter) -> SimpleSpanProcessor:
-        # It's OK if these exporter types change.
-        # We just need access to the FallbackSpanExporter either way to swap out its underlying exporter.
-        assert isinstance(exporter, WrapperSpanExporter)
-        fallback_exporter = exporter.wrapped_exporter
-        assert isinstance(fallback_exporter, FallbackSpanExporter)
-        fallback_exporter.exporter = FailureExporter()
-
-        return SimpleSpanProcessor(exporter)
-
+    data_dir = Path(tmp_path) / 'logfire_data'
     with request_mocker:
-        data_dir = Path(tmp_path) / 'logfire_data'
         logfire.configure(
             send_to_logfire=True,
             data_dir=data_dir,
             token='abc1',
-            default_span_processor=default_span_processor,
-            additional_metric_readers=[InMemoryMetricReader()],
+            console=False,
         )
         wait_for_check_token_thread()
 
+    send_to_logfire_processor, *_ = get_span_processors()
+    # It's OK if these processor/exporter types change.
+    # We just need access to the FallbackSpanExporter either way to swap out its underlying exporter.
+    assert isinstance(send_to_logfire_processor, MainSpanProcessorWrapper)
+    batch_span_processor = send_to_logfire_processor.processor
+    assert isinstance(batch_span_processor, BatchSpanProcessor)
+    exporter = batch_span_processor.span_exporter
+    assert isinstance(exporter, WrapperSpanExporter)
+    fallback_exporter = exporter.wrapped_exporter
+    assert isinstance(fallback_exporter, FallbackSpanExporter)
+    fallback_exporter.exporter = FailureExporter()
+
+    with logfire.span('test'):
+        pass
+
     assert not data_dir.exists()
     path = data_dir / 'logfire_spans.bin'
 
     with pytest.warns(WritingFallbackWarning, match=f'Failed to export spans, writing to fallback file: {path}'):
+        logfire.force_flush()
+
+    assert path.exists()
+
+
+def test_configure_export_delay() -> None:
+    class TrackingExporter(SpanExporter):
+        def __init__(self) -> None:
+            self.last_export_timestamp: float | None = None
+            self.export_delays: list[float] = []
+
+        def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+            t = time()
+            if self.last_export_timestamp is not None:
+                self.export_delays.append(t - self.last_export_timestamp)
+            self.last_export_timestamp = t
+            return SpanExportResult.SUCCESS
+
+    def configure_tracking_exporter():
+        request_mocker = requests_mock.Mocker()
+        request_mocker.get(
+            'https://logfire-api.pydantic.dev/v1/info',
+            json={'project_name': 'myproject', 'project_url': 'fake_project_url'},
+        )
+
+        with request_mocker:
+            logfire.configure(
+                send_to_logfire=True,
+                token='abc1',
+                console=False,
+                fast_shutdown=True,
+            )
+            wait_for_check_token_thread()
+
+        send_to_logfire_processor, *_ = get_span_processors()
+        assert isinstance(send_to_logfire_processor, MainSpanProcessorWrapper)
+        batch_span_processor = send_to_logfire_processor.processor
+        assert isinstance(batch_span_processor, BatchSpanProcessor)
+
+        batch_span_processor.span_exporter = TrackingExporter()
+        return batch_span_processor.span_exporter
+
+    def check_delays(exp: TrackingExporter, min_delay: float, max_delay: float) -> None:
+        for delay in exp.export_delays:
+            assert min_delay < delay < max_delay, f'delay was {delay}, which is not between {min_delay} and {max_delay}'
+
+    # test the default value
+    exporter = configure_tracking_exporter()
+    while not exporter.export_delays:
         with logfire.span('test'):
             pass
+        sleep(0.1)
+    check_delays(exporter, 0.4, 1.0)  # our default is 500ms
 
-    assert path.exists()
+    # test a very small value
+    with patch.dict(os.environ, {'OTEL_BSP_SCHEDULE_DELAY': '1'}):
+        exporter = configure_tracking_exporter()
+
+    while not exporter.export_delays:
+        with logfire.span('test'):
+            pass
+        sleep(0.03)
+    check_delays(exporter, 0.0, 0.1)  # since we set 1ms it should be a very short delay
 
 
 def test_configure_service_version(tmp_path: str) -> None:
