feat(span-first): Support before_send_span (#6239)

### Description
Add support for `before_send_span` in span streaming mode.

`before_send_span` is different from `before_send_metric` and
`before_send_log` in that:
- it doesn't allow users to drop a span (i.e., return `None`)
- it only allows to modify specific parts of the span

To that end, we're now serializing the span earlier, and exposing the
serialized dictionary in the `before_send` callback. This is consistent
with metrics and logs. It also means we're now queuing dictionaries
instead of `StreamedSpan` instances in the span batcher, which should
also decrease our memory footprint.

This aligns our implementation with JS.

See https://develop.sentry.dev/sdk/telemetry/spans/scrubbing-data/ for
spec.


#### Issues

* resolves: #5388
* resolves:
https://linear.app/getsentry/issue/PY-2057/add-before-send-span

Author: sentrivana

sentry_sdk/_span_batcher.py

@@ -14,10 +14,10 @@
 if TYPE_CHECKING:
     from typing import Any, Callable, Optional
 
-    from sentry_sdk.traces import StreamedSpan
+    from sentry_sdk._types import SpanJSON
 
 
-class SpanBatcher(Batcher["StreamedSpan"]):
+class SpanBatcher(Batcher["SpanJSON"]):
     # MAX_BEFORE_FLUSH should be lower than MAX_BEFORE_DROP, so that there is
     # a bit of a buffer for spans that appear between the trigger to flush
     # and actually flushing the buffer.
@@ -43,7 +43,7 @@ def __init__(
         # by trace_id, so that we can then send the buckets each in its own
         # envelope.
         # trace_id -> span buffer
-        self._span_buffer: dict[str, list["StreamedSpan"]] = defaultdict(list)
+        self._span_buffer: dict[str, list["SpanJSON"]] = defaultdict(list)
         self._running_size: dict[str, int] = defaultdict(lambda: 0)
         self._capture_func = capture_func
         self._record_lost_func = record_lost_func
@@ -100,7 +100,7 @@ def _flush_loop(self) -> None:
                 self._flush()
                 self._last_full_flush = time.monotonic()
 
-    def add(self, span: "StreamedSpan") -> None:
+    def add(self, span: "SpanJSON") -> None:
         # Bail out if the current thread is already executing batcher code.
         # This prevents deadlocks when code running inside the batcher (e.g.
         # _add_to_envelope during flush, or _flush_event.wait/set) triggers
@@ -116,7 +116,7 @@ def add(self, span: "StreamedSpan") -> None:
                 return None
 
             with self._lock:
-                size = len(self._span_buffer[span.trace_id])
+                size = len(self._span_buffer[span["trace_id"]])
                 if size >= self.MAX_BEFORE_DROP:
                     self._record_lost_func(
                         reason="queue_overflow",
@@ -125,14 +125,15 @@ def add(self, span: "StreamedSpan") -> None:
                     )
                     return None
 
-                self._span_buffer[span.trace_id].append(span)
-                self._running_size[span.trace_id] += self._estimate_size(span)
+                self._span_buffer[span["trace_id"]].append(span)
+                self._running_size[span["trace_id"]] += self._estimate_size(span)
 
                 if (
                     size + 1 >= self.MAX_BEFORE_FLUSH
-                    or self._running_size[span.trace_id] >= self.MAX_BYTES_BEFORE_FLUSH
+                    or self._running_size[span["trace_id"]]
+                    >= self.MAX_BYTES_BEFORE_FLUSH
                 ):
-                    self._pending_flush.add(span.trace_id)
+                    self._pending_flush.add(span["trace_id"])
                     notify = True
                 else:
                     notify = False
@@ -143,12 +144,12 @@ def add(self, span: "StreamedSpan") -> None:
             self._active.flag = False
 
     @staticmethod
-    def _estimate_size(item: "StreamedSpan") -> int:
+    def _estimate_size(item: "SpanJSON") -> int:
         # Rough estimate of serialized span size that's quick to compute.
         # 210 is the rough size of the payload without attributes, and then we
         # estimate the attributes separately.
         estimate = 210
-        for value in item._attributes.values():
+        for value in (item.get("attributes") or {}).values():
             estimate += 50
 
             if isinstance(value, str):
@@ -159,26 +160,15 @@ def _estimate_size(item: "StreamedSpan") -> int:
         return estimate
 
     @staticmethod
-    def _to_transport_format(item: "StreamedSpan") -> "Any":
-        res: "dict[str, Any]" = {
-            "trace_id": item.trace_id,
-            "span_id": item.span_id,
-            "name": item._name if item._name is not None else "<unlabeled span>",
-            "status": item._status,
-            "is_segment": item._is_segment(),
-            "start_timestamp": item._start_timestamp.timestamp(),
-        }
-
-        if item._end_timestamp:
-            res["end_timestamp"] = item._end_timestamp.timestamp()
-
-        if item._parent_span_id:
-            res["parent_span_id"] = item._parent_span_id
-
-        if item._attributes:
+    def _to_transport_format(item: "SpanJSON") -> "Any":
+        res = {k: v for k, v in item.items() if k not in ("_segment_span",)}
+
+        if item.get("attributes"):
             res["attributes"] = {
-                k: serialize_attribute(v) for (k, v) in item._attributes.items()
+                k: serialize_attribute(v) for (k, v) in item["attributes"].items()
             }
+        else:
+            del res["attributes"]
 
         return res
 
@@ -202,7 +192,7 @@ def _flush(self, only_pending: bool = False) -> None:
                 if not spans:
                     continue
 
-                dsc = spans[0]._dynamic_sampling_context()
+                dsc = spans[0]["_segment_span"]._dynamic_sampling_context()
 
                 # Max per envelope is 1000, so if we happen to have more than
                 # 1000 spans in one bucket, we'll need to separate them.

sentry_sdk/_types.py

@@ -145,6 +145,8 @@ def substituted_because_contains_sensitive_data(cls) -> "AnnotatedValue":
 
     from typing_extensions import Literal, TypedDict
 
+    import sentry_sdk
+
     class SDKInfo(TypedDict):
         name: str
         version: str
@@ -318,6 +320,22 @@ class SDKInfo(TypedDict):
 
     MetricProcessor = Callable[[Metric, Hint], Optional[Metric]]
 
+    SpanJSON = TypedDict(
+        "SpanJSON",
+        {
+            "trace_id": str,
+            "span_id": str,
+            "parent_span_id": NotRequired[str],
+            "name": str,
+            "status": str,
+            "is_segment": bool,
+            "start_timestamp": float,
+            "end_timestamp": NotRequired[float],
+            "attributes": NotRequired[Attributes],
+            "_segment_span": NotRequired["sentry_sdk.traces.StreamedSpan"],
+        },
+    )
+
     # TODO: Make a proper type definition for this (PRs welcome!)
     Breadcrumb = Dict[str, Any]
 

sentry_sdk/client.py

@@ -34,7 +34,7 @@
 from sentry_sdk.scrubber import EventScrubber
 from sentry_sdk.serializer import serialize
 from sentry_sdk.sessions import SessionFlusher
-from sentry_sdk.traces import SpanStatus
+from sentry_sdk.traces import SpanStatus, StreamedSpan
 from sentry_sdk.tracing import trace
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.transport import (
@@ -52,6 +52,7 @@
     format_timestamp,
     get_before_send_log,
     get_before_send_metric,
+    get_before_send_span,
     get_default_release,
     get_sdk_name,
     get_type_name,
@@ -1169,34 +1170,72 @@ 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)
+            serialized = telemetry
+
         elif ty == "metric":
             before_send = get_before_send_metric(self.options)
+            serialized = telemetry
+
+        elif ty == "span":
+            before_send = get_before_send_span(self.options)
+            serialized = telemetry._to_json()  # type: ignore[union-attr]
 
         if before_send is not None:
-            telemetry = before_send(telemetry, {})  # type: ignore
+            serialized = before_send(serialized, {})  # type: ignore[arg-type]
+
+            if ty in ("log", "metric"):
+                # Logs and metrics can be dropped in their respective
+                # before_send, so if we get None, don't queue them for sending.
+                if serialized is None:
+                    return
+
+            elif ty == "span" and isinstance(telemetry, StreamedSpan):
+                # Spans can't be dropped in before_send_span by design. They can
+                # be altered though (e.g. to sanitize). Only allow changes to
+                # name and attributes.
+                if isinstance(serialized, dict) and "name" in serialized:
+                    telemetry.name = serialized["name"]  # type: ignore[typeddict-item]
+                    telemetry._attributes = {}
+                    for k, v in (serialized.get("attributes") or {}).items():
+                        telemetry.set_attribute(k, v)
+
+                else:
+                    logger.debug(
+                        "[Tracing] Invalid return value from before_send_span. Keeping original span."
+                    )
 
-        if telemetry is None:
-            return
+                serialized = telemetry._to_json()
 
         batcher = None
         if ty == "log":
             batcher = self.log_batcher
+
         elif ty == "metric":
             batcher = self.metrics_batcher
+
         elif ty == "span":
+            # We need a reference to the segment span in the batcher to populate
+            # the dynamic sampling context (DSC)
+            serialized["_segment_span"] = telemetry._segment  # type: ignore
             batcher = self.span_batcher
 
         if batcher is not None:
-            batcher.add(telemetry)  # type: ignore
+            batcher.add(serialized)  # type: ignore
 
     def _capture_log(self, log: "Optional[Log]", scope: "Scope") -> None:
         self._capture_telemetry(log, "log", scope)

sentry_sdk/consts.py

@@ -56,6 +56,7 @@ class CompressionAlgo(Enum):
         Log,
         Metric,
         ProfilerMode,
+        SpanJSON,
         TracesSampler,
         TransactionProcessor,
     )
@@ -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[[SpanJSON, Hint], Optional[SpanJSON]]
+            ],
             "suppress_asgi_chained_exceptions": Optional[bool],
         },
         total=False,

sentry_sdk/traces.py

@@ -43,7 +43,7 @@
         overload,
     )
 
-    from sentry_sdk._types import Attributes, AttributeValue
+    from sentry_sdk._types import Attributes, AttributeValue, SpanJSON
     from sentry_sdk.profiler.continuous_profiler import ContinuousProfile
 
     P = ParamSpec("P")
@@ -574,6 +574,26 @@ def _set_segment_attributes(self) -> None:
 
         self.set_attribute("process.command_args", sys.argv)
 
+    def _to_json(self) -> "SpanJSON":
+        res: "SpanJSON" = {
+            "trace_id": self.trace_id,
+            "span_id": self.span_id,
+            "name": self._name if self._name is not None else "<unlabeled span>",
+            "status": self._status,
+            "is_segment": self._is_segment(),
+            "start_timestamp": self._start_timestamp.timestamp(),
+        }
+
+        if self._end_timestamp:
+            res["end_timestamp"] = self._end_timestamp.timestamp()
+
+        if self._parent_span_id:
+            res["parent_span_id"] = self._parent_span_id
+
+        res["attributes"] = {k: v for k, v in self._attributes.items()}
+
+        return res
+
 
 class NoOpStreamedSpan(StreamedSpan):
     __slots__ = (

sentry_sdk/utils.py

@@ -76,6 +76,7 @@
         Log,
         Metric,
         SerializedAttributeValue,
+        SpanJSON,
     )
 
     P = ParamSpec("P")
@@ -2111,6 +2112,15 @@ def get_before_send_metric(
     )
 
 
+def get_before_send_span(
+    options: "Optional[dict[str, Any]]",
+) -> "Optional[Callable[[SpanJSON, Hint], Optional[SpanJSON]]]":
+    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/integrations/sqlalchemy/test_sqlalchemy.py

@@ -1079,19 +1079,19 @@ class Person(Base):
 
             class fake_record_sql_queries:  # noqa: N801
                 def __init__(self, *args, **kwargs):
-                    with record_sql_queries_supporting_streaming(
+                    self._ctx_mgr = record_sql_queries_supporting_streaming(
                         *args, **kwargs
-                    ) as span:
-                        self.span = span
+                    )
 
+                def __enter__(self):
+                    self.span = self._ctx_mgr.__enter__()
                     self.span._start_timestamp = datetime(2024, 1, 1, microsecond=0)
                     self.span._end_timestamp = datetime(2024, 1, 1, microsecond=101000)
-
-                def __enter__(self):
                     return self.span
 
                 def __exit__(self, type, value, traceback):
-                    pass
+                    self.span._end_timestamp = None
+                    self._ctx_mgr.__exit__(type, value, traceback)
 
             with mock.patch(
                 "sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming",

tests/tracing/test_span_batcher.py

@@ -236,7 +236,7 @@ def test_weight_based_flushing_by_attribute_size(
     with sentry_sdk.traces.start_span(name="small span") as bare_span:
         pass
 
-    bare_span_size = SpanBatcher._estimate_size(bare_span)
+    bare_span_size = SpanBatcher._estimate_size(bare_span._to_json())
     big_attr = "x" * bare_span_size
 
     monkeypatch.setattr(SpanBatcher, "MAX_BYTES_BEFORE_FLUSH", bare_span_size * 3)