OTel log scrubbing (#903)

Author: alexmojaki

logfire/_internal/config.py

@@ -83,7 +83,7 @@
     ShowParentsConsoleSpanExporter,
     SimpleConsoleSpanExporter,
 )
-from .exporters.logs import CheckSuppressInstrumentationLogProcessorWrapper
+from .exporters.logs import CheckSuppressInstrumentationLogProcessorWrapper, MainLogProcessorWrapper
 from .exporters.otlp import OTLPExporterHttpSession, QuietLogExporter, QuietSpanExporter, RetryFewerSpansSpanExporter
 from .exporters.processor_wrapper import CheckSuppressInstrumentationProcessorWrapper, MainSpanProcessorWrapper
 from .exporters.quiet_metrics import QuietMetricExporter
@@ -1015,7 +1015,9 @@ def fix_pid():  # pragma: no cover
             multi_log_processor = SynchronousMultiLogRecordProcessor()
             for processor in log_record_processors:
                 multi_log_processor.add_log_record_processor(processor)
-            root_log_processor = CheckSuppressInstrumentationLogProcessorWrapper(multi_log_processor)
+            root_log_processor = CheckSuppressInstrumentationLogProcessorWrapper(
+                MainLogProcessorWrapper(multi_log_processor, self.scrubber)
+            )
             logger_provider = SDKLoggerProvider(resource)
             logger_provider.add_log_record_processor(root_log_processor)
 

logfire/_internal/exporters/logs.py

@@ -1,7 +1,10 @@
+from dataclasses import dataclass
+
 from opentelemetry.sdk._logs import LogData
 
 import logfire
 from logfire._internal.exporters.wrapper import WrapperLogProcessor
+from logfire._internal.scrubbing import BaseScrubber
 from logfire._internal.utils import is_instrumentation_suppressed
 
 
@@ -16,3 +19,12 @@ def emit(self, log_data: LogData):
             return
         with logfire.suppress_instrumentation():
             return super().emit(log_data)
+
+
+@dataclass
+class MainLogProcessorWrapper(WrapperLogProcessor):
+    scrubber: BaseScrubber
+
+    def emit(self, log_data: LogData):
+        log_data.log_record = self.scrubber.scrub_log(log_data.log_record)
+        return super().emit(log_data)

logfire/_internal/exporters/test.py

@@ -171,6 +171,9 @@ def build_attributes(
 class TestLogExporter(InMemoryLogExporter):
     """A LogExporter that stores exported logs in a list for asserting in tests."""
 
+    # NOTE: Avoid test discovery by pytest.
+    __test__ = False
+
     def __init__(self, ns_timestamp_generator: typing.Callable[[], int]) -> None:
         super().__init__()
         self.ns_timestamp_generator = ns_timestamp_generator
@@ -222,3 +225,6 @@ def build_log(log_data: LogData) -> dict[str, Any]:
             return res
 
         return [build_log(log) for log in self.get_finished_logs()]
+
+    def shutdown(self) -> None:
+        self.clear()

logfire/_internal/scrubbing.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import json
 import re
 from abc import ABC, abstractmethod
@@ -8,6 +9,7 @@
 
 import typing_extensions
 from opentelemetry.attributes import BoundedAttributes
+from opentelemetry.sdk._logs import LogRecord
 from opentelemetry.sdk.trace import Event
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import Link
@@ -129,11 +131,15 @@ class BaseScrubber(ABC):
         SpanAttributes.URL_FULL,
         SpanAttributes.URL_PATH,
         SpanAttributes.URL_QUERY,
+        'event.name',
     }
 
     @abstractmethod
     def scrub_span(self, span: ReadableSpanDict): ...  # pragma: no cover
 
+    @abstractmethod
+    def scrub_log(self, log: LogRecord) -> LogRecord: ...  # pragma: no cover
+
     @abstractmethod
     def scrub_value(self, path: JsonPath, value: Any) -> tuple[Any, list[ScrubbedNote]]: ...  # pragma: no cover
 
@@ -142,6 +148,9 @@ class NoopScrubber(BaseScrubber):
     def scrub_span(self, span: ReadableSpanDict):
         pass
 
+    def scrub_log(self, log: LogRecord) -> LogRecord:
+        return log
+
     def scrub_value(self, path: JsonPath, value: Any) -> tuple[Any, list[ScrubbedNote]]:  # pragma: no cover
         return value, []
 
@@ -158,6 +167,10 @@ def __init__(self, patterns: Sequence[str] | None, callback: ScrubCallback | Non
         self._pattern = re.compile('|'.join(patterns), re.IGNORECASE | re.DOTALL)
         self._callback = callback
 
+    def scrub_log(self, log: LogRecord) -> LogRecord:
+        span_scrubber = SpanScrubber(self)
+        return span_scrubber.scrub_log(log)
+
     def scrub_span(self, span: ReadableSpanDict):
         scope = span['instrumentation_scope']
         if scope and scope.name in ['logfire.openai', 'logfire.anthropic']:
@@ -194,15 +207,19 @@ def __init__(self, parent: Scrubber):
         self._pattern = parent._pattern  # type: ignore
         self._callback = parent._callback  # type: ignore
         self.scrubbed: list[ScrubbedNote] = []
+        self.did_scrub = False
 
     def scrub_span(self, span: ReadableSpanDict):
         # We need to use BoundedAttributes because:
         # 1. For events and links, we get an error otherwise:
         #      https://github.com/open-telemetry/opentelemetry-python/issues/3761
         # 2. The callback might return a value that isn't of the type required by OTEL,
         #      in which case BoundAttributes will discard it to prevent an error.
-        # TODO silently throwing away the result is bad, and BoundedAttributes might be bad for performance.
-        span['attributes'] = BoundedAttributes(attributes=self.scrub(('attributes',), span['attributes']))
+        # TODO silently throwing away the result is bad, and BoundedAttributes is bad for performance.
+        new_attributes = self.scrub(('attributes',), span['attributes'])
+        if self.did_scrub:
+            span['attributes'] = BoundedAttributes(attributes=new_attributes)
+
         span['events'] = [
             Event(
                 # We don't scrub the event name because in theory it should be a low-cardinality general description,
@@ -221,6 +238,22 @@ def scrub_span(self, span: ReadableSpanDict):
             for i, link in enumerate(span['links'])
         ]
 
+    def scrub_log(self, log: LogRecord) -> LogRecord:
+        new_attributes: dict[str, Any] | None = self.scrub(('attributes',), log.attributes)
+        new_body = self.scrub(('log_body',), log.body)
+
+        if not self.did_scrub:
+            return log
+
+        if self.scrubbed:
+            new_attributes = new_attributes or {}
+            new_attributes[ATTRIBUTES_SCRUBBED_KEY] = json.dumps(self.scrubbed)
+
+        result = copy.copy(log)
+        result.attributes = BoundedAttributes(attributes=new_attributes)
+        result.body = new_body
+        return result
+
     def scrub_event_attributes(self, event: Event, index: int):
         attributes = event.attributes or {}
         path = ('otel_events', index, 'attributes')
@@ -265,7 +298,9 @@ def scrub(self, path: JsonPath, value: Any) -> Any:
 
     def _redact(self, match: ScrubMatch) -> Any:
         if self._callback and (result := self._callback(match)) is not None:
+            self.did_scrub = self.did_scrub or result is not match.value
             return result
+        self.did_scrub = True
         matched_substring = match.pattern_match.group(0)
         self.scrubbed.append(ScrubbedNote(path=match.path, matched_substring=matched_substring))
         return f'[Scrubbed due to {matched_substring!r}]'

tests/test_configure.py

@@ -53,7 +53,7 @@
     sanitize_project_name,
 )
 from logfire._internal.exporters.console import ConsoleLogExporter, ShowParentsConsoleSpanExporter
-from logfire._internal.exporters.logs import CheckSuppressInstrumentationLogProcessorWrapper
+from logfire._internal.exporters.logs import CheckSuppressInstrumentationLogProcessorWrapper, MainLogProcessorWrapper
 from logfire._internal.exporters.otlp import QuietLogExporter, QuietSpanExporter
 from logfire._internal.exporters.processor_wrapper import (
     CheckSuppressInstrumentationProcessorWrapper,
@@ -1630,11 +1630,14 @@ def get_metric_readers() -> Iterable[SpanProcessor]:
 
 
 def get_log_record_processors() -> Iterable[LogRecordProcessor]:
-    [root] = get_logger_provider().provider._multi_log_record_processor._log_record_processors  # type: ignore
-    assert isinstance(root, CheckSuppressInstrumentationLogProcessorWrapper)
-    assert isinstance(root.processor, SynchronousMultiLogRecordProcessor)
-
-    return root.processor._log_record_processors  # type: ignore
+    [processor] = get_logger_provider().provider._multi_log_record_processor._log_record_processors  # type: ignore
+    assert isinstance(processor, CheckSuppressInstrumentationLogProcessorWrapper)
+    processor = processor.processor
+    assert isinstance(processor, MainLogProcessorWrapper)
+    processor = processor.processor
+    assert isinstance(processor, SynchronousMultiLogRecordProcessor)
+
+    return processor._log_record_processors  # type: ignore
 
 
 def test_dynamic_module_ignored_in_ensure_flush_after_aws_lambda(

tests/test_otel_logs.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import Any, Sequence
 from unittest import mock
 
 import pytest
@@ -16,9 +16,7 @@
 import logfire
 from logfire import suppress_instrumentation
 from logfire._internal.exporters.otlp import QuietLogExporter
-
-if TYPE_CHECKING:
-    from logfire._internal.exporters.test import TestLogExporter
+from logfire.testing import TestLogExporter
 
 
 def test_otel_logs_supress_scopes(logs_exporter: InMemoryLogExporter, config_kwargs: dict[str, Any]) -> None: