feat(span-first): Support before_send_span

sentry_sdk/client.py

@@ -3,6 +3,7 @@
 import random
 import socket
 from collections.abc import Mapping
+from copy import deepcopy
 from datetime import datetime, timezone
 from importlib import import_module
 from typing import TYPE_CHECKING, List, Dict, cast, overload
@@ -25,10 +26,12 @@
     logger,
     get_before_send_log,
     get_before_send_metric,
+    get_before_send_span,
     has_logs_enabled,
     has_metrics_enabled,
 )
 from sentry_sdk.serializer import serialize
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing import trace
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.transport import (
@@ -71,7 +74,6 @@
     from sentry_sdk.scope import Scope
     from sentry_sdk.session import Session
     from sentry_sdk.spotlight import SpotlightClient
-    from sentry_sdk.traces import StreamedSpan
     from sentry_sdk.transport import Transport, Item
     from sentry_sdk._log_batcher import LogBatcher
     from sentry_sdk._metrics_batcher import MetricsBatcher
@@ -938,23 +940,54 @@ def _capture_telemetry(
         ty: str,
         scope: "Scope",
     ) -> None:
-        # Capture attributes-based telemetry (logs, metrics, spansV2)
+        """
+        Capture attributes-based telemetry (logs, metrics, streamed spans).
+
+        Apply any attributes set on the scope to it, and run the user's
+        before_send_{telemetry} on it, if applicable.
+        """
         if telemetry is None:
             return
 
         scope.apply_to_telemetry(telemetry)
 
         before_send = None
+
         if ty == "log":
             before_send = get_before_send_log(self.options)
+            snapshot = telemetry
+
         elif ty == "metric":
             before_send = get_before_send_metric(self.options)
+            snapshot = telemetry
 
-        if before_send is not None:
-            telemetry = before_send(telemetry, {})  # type: ignore
+        elif ty == "span":
+            before_send = get_before_send_span(self.options)
+            # We don't want to expose the actual underlying span in
+            # before_send_span to not allow arbitrary edits. Expose a copy
+            # instead.
+            snapshot = deepcopy(telemetry)
 
-        if telemetry is None:
-            return
+        if before_send is not None:
+            result = before_send(snapshot, {})
+
+            # Logs and metrics can be dropped in their respective
+            # before_send, so if we get None, don't queue them for sending.
+            if ty in ("log", "metric"):
+                if result is None:
+                    return
+
+            # Spans can't be dropped in before_send_span by design. They can
+            # be altered though (name and attributes can be changed, e.g. to
+            # sanitize).
+            #
+            # If we get anything but a StreamedSpan back from before_send_span,
+            # just ignore it. Otherwise, take the returned StreamedSpan and
+            # merge it with the original.
+            elif ty == "span":
+                if isinstance(result, StreamedSpan):
+                    telemetry._attributes = result._attributes
+                    telemetry._name = result._name
 
         batcher = None
         if ty == "log":

sentry_sdk/consts.py

@@ -46,6 +46,7 @@ class CompressionAlgo(Enum):
     from typing_extensions import Literal, TypedDict
 
     import sentry_sdk
+    from sentry_sdk.traces import StreamedSpan
     from sentry_sdk._types import (
         BreadcrumbProcessor,
         ContinuousProfilerMode,
@@ -85,6 +86,9 @@ class CompressionAlgo(Enum):
             "before_send_metric": Optional[Callable[[Metric, Hint], Optional[Metric]]],
             "trace_lifecycle": Optional[Literal["static", "stream"]],
             "ignore_spans": Optional[IgnoreSpansConfig],
+            "before_send_span": Optional[
+                Callable[[StreamedSpan, Hint], Optional[StreamedSpan]]
+            ],
             "suppress_asgi_chained_exceptions": Optional[bool],
         },
         total=False,

sentry_sdk/utils.py

@@ -77,6 +77,7 @@
         Metric,
         SerializedAttributeValue,
     )
+    from sentry_sdk.traces import StreamedSpan
 
     P = ParamSpec("P")
     R = TypeVar("R")
@@ -2111,6 +2112,15 @@ def get_before_send_metric(
     )
 
 
+def get_before_send_span(
+    options: "Optional[dict[str, Any]]",
+) -> "Optional[Callable[[StreamedSpan, Hint], Optional[StreamedSpan]]]":
+    if options is None:
+        return None
+
+    return options["_experiments"].get("before_send_span")
+
+
 def format_attribute(val: "Any") -> "AttributeValue":
     """
     Turn unsupported attribute value types into an AttributeValue.

tests/tracing/test_span_streaming.py

@@ -271,6 +271,135 @@ def traces_sampler(sampling_context):
         ...
 
 
+def test_before_send_span_basic(sentry_init, capture_items):
+    def before_send_span(span, hint):
+        assert isinstance(span, StreamedSpan)
+
+        span.name = "Better span name"
+        span.remove_attribute("drop")
+        span.set_attribute("sanitize", "[Removed]")
+        span.set_attribute("add", "new")
+
+        return span
+
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={
+            "before_send_span": before_send_span,
+            "trace_lifecycle": "stream",
+        },
+    )
+
+    items = capture_items("span")
+
+    with sentry_sdk.traces.start_span(
+        name="span",
+        attributes={
+            "drop": True,
+            "sanitize": "myamazingpassword",
+        },
+    ):
+        ...
+
+    sentry_sdk.get_client().flush()
+    spans = [item.payload for item in items]
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert span["name"] == "Better span name"
+    assert "drop" not in span["attributes"]
+    assert span["attributes"]["sanitize"] == "[Removed]"
+    assert span["attributes"]["add"] == "new"
+
+
+def test_before_send_span_invalid_return_value(sentry_init, capture_items):
+    def before_send_span(span, hint):
+        # Spans can't be dropped in before_send_span, so unsupported return
+        # values will be ignored
+        return None
+
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={
+            "before_send_span": before_send_span,
+            "trace_lifecycle": "stream",
+        },
+    )
+
+    items = capture_items("span")
+
+    with sentry_sdk.traces.start_span(name="span"):
+        ...
+
+    sentry_sdk.get_client().flush()
+    spans = [item.payload for item in items]
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert span["name"] == "span"
+
+
+def test_before_send_span_unsupported_edit(sentry_init, capture_items):
+    def before_send_span(span, hint):
+        # Anything beyond attribute and name changes will be ignored
+        span._trace_id = "my-trace-id"
+
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={
+            "before_send_span": before_send_span,
+            "trace_lifecycle": "stream",
+        },
+    )
+
+    items = capture_items("span")
+
+    with sentry_sdk.traces.start_span(name="span"):
+        ...
+
+    sentry_sdk.get_client().flush()
+    spans = [item.payload for item in items]
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert span["name"] == "span"
+    assert span["trace_id"] != "my-trace-id"
+
+
+def test_before_send_span_doesnt_receive_ignored_spans(sentry_init, capture_items):
+    before_send_span_called = False
+
+    def before_send_span(span, hint):
+        nonlocal before_send_span_called
+        before_send_span_called = True
+        return span
+
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={
+            "before_send_span": before_send_span,
+            "trace_lifecycle": "stream",
+            "ignore_spans": [
+                "ignored",
+            ],
+        },
+    )
+
+    items = capture_items("span")
+
+    with sentry_sdk.traces.start_span(name="ignored"):
+        ...
+
+    sentry_sdk.get_client().flush()
+    spans = [item.payload for item in items]
+
+    assert not spans
+    assert not before_send_span_called
+
+
 def test_span_attributes(sentry_init, capture_items):
     sentry_init(
         traces_sample_rate=1.0,