ref: Support outgoing trace propagation in span first (18)

sentry_sdk/_span_batcher.py

@@ -116,9 +116,7 @@ def _flush(self) -> None:
             envelopes = []
             for trace_id, spans in self._span_buffer.items():
                 if spans:
-                    # TODO[span-first]
-                    # dsc = spans[0].dynamic_sampling_context()
-                    dsc = None
+                    dsc = spans[0]._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/ai/utils.py

@@ -14,6 +14,7 @@
 
 import sentry_sdk
 from sentry_sdk.utils import logger
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 MAX_GEN_AI_MESSAGE_BYTES = 20_000  # 20KB
 # Maximum characters when only a single message is left after bytes truncation
@@ -525,7 +526,14 @@ def normalize_message_roles(messages: "list[dict[str, Any]]") -> "list[dict[str,
 
 
 def get_start_span_function() -> "Callable[..., Any]":
+    if has_span_streaming_enabled(sentry_sdk.get_client().options):
+        return sentry_sdk.traces.start_span
+
     current_span = sentry_sdk.get_current_span()
+    if not isinstance(current_span, Span):
+        # mypy
+        return sentry_sdk.traces.start_span
+
     transaction_exists = (
         current_span is not None and current_span.containing_transaction is not None
     )

sentry_sdk/api.py

@@ -6,6 +6,7 @@
 from sentry_sdk._init_implementation import init
 from sentry_sdk.consts import INSTRUMENTER
 from sentry_sdk.scope import Scope, _ScopeManager, new_scope, isolation_scope
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing import NoOpSpan, Transaction, trace
 from sentry_sdk.crons import monitor
 
@@ -37,6 +38,7 @@
         LogLevelStr,
         SamplingContext,
     )
+    from sentry_sdk.traces import StreamedSpan
     from sentry_sdk.tracing import Span, TransactionKwargs
 
     T = TypeVar("T")
@@ -409,7 +411,9 @@ def set_measurement(name: str, value: float, unit: "MeasurementUnit" = "") -> No
         transaction.set_measurement(name, value, unit)
 
 
-def get_current_span(scope: "Optional[Scope]" = None) -> "Optional[Span]":
+def get_current_span(
+    scope: "Optional[Scope]" = None,
+) -> "Optional[Union[Span, StreamedSpan]]":
     """
     Returns the currently active span if there is one running, otherwise `None`
     """
@@ -525,6 +529,16 @@ def update_current_span(
     if current_span is None:
         return
 
+    if isinstance(current_span, StreamedSpan):
+        warnings.warn(
+            "The `update_current_span` API isn't available in streaming mode. "
+            "Retrieve the current span with get_current_span() and use its API "
+            "directly.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return
+
     if op is not None:
         current_span.op = op
 

sentry_sdk/feature_flags.py

@@ -1,6 +1,7 @@
 import copy
 import sentry_sdk
 from sentry_sdk._lru_cache import LRUCache
+from sentry_sdk.tracing import Span
 from threading import Lock
 
 from typing import TYPE_CHECKING, Any
@@ -61,5 +62,5 @@ def add_feature_flag(flag: str, result: bool) -> None:
     flags.set(flag, result)
 
     span = sentry_sdk.get_current_span()
-    if span:
+    if span and isinstance(span, Span):
         span.set_flag(f"flag.evaluation.{flag}", result)

sentry_sdk/scope.py

@@ -583,8 +583,12 @@ def get_traceparent(self, *args: "Any", **kwargs: "Any") -> "Optional[str]":
         client = self.get_client()
 
         # If we have an active span, return traceparent from there
-        if has_tracing_enabled(client.options) and self.span is not None:
-            return self.span.to_traceparent()
+        if (
+            has_tracing_enabled(client.options)
+            and self.span is not None
+            and not isinstance(self.span, NoOpStreamedSpan)
+        ):
+            return self.span._to_traceparent()
 
         # else return traceparent from the propagation context
         return self.get_active_propagation_context().to_traceparent()
@@ -597,8 +601,12 @@ def get_baggage(self, *args: "Any", **kwargs: "Any") -> "Optional[Baggage]":
         client = self.get_client()
 
         # If we have an active span, return baggage from there
-        if has_tracing_enabled(client.options) and self.span is not None:
-            return self.span.to_baggage()
+        if (
+            has_tracing_enabled(client.options)
+            and self.span is not None
+            and not isinstance(self.span, NoOpStreamedSpan)
+        ):
+            return self.span._to_baggage()
 
         # else return baggage from the propagation context
         return self.get_active_propagation_context().get_baggage()
@@ -607,8 +615,12 @@ def get_trace_context(self) -> "Dict[str, Any]":
         """
         Returns the Sentry "trace" context from the Propagation Context.
         """
-        if has_tracing_enabled(self.get_client().options) and self._span is not None:
-            return self._span.get_trace_context()
+        if (
+            has_tracing_enabled(self.get_client().options)
+            and self._span is not None
+            and not isinstance(self.span, NoOpStreamedSpan)
+        ):
+            return self._span._get_trace_context()
 
         # if we are tracing externally (otel), those values take precedence
         external_propagation_context = get_external_propagation_context()
@@ -673,7 +685,7 @@ def iter_trace_propagation_headers(
         span = span or self.span
 
         if has_tracing_enabled(client.options) and span is not None:
-            for header in span.iter_headers():
+            for header in span._iter_headers():
                 yield header
         elif has_external_propagation_context():
             # when we have an external_propagation_context (otlp)
@@ -720,7 +732,7 @@ def clear(self) -> None:
         self.clear_breadcrumbs()
         self._should_capture: bool = True
 
-        self._span: "Optional[Span]" = None
+        self._span: "Optional[Union[Span, StreamedSpan]]" = None
         self._session: "Optional[Session]" = None
         self._force_auto_session_tracking: "Optional[bool]" = None
 
@@ -774,6 +786,14 @@ def transaction(self) -> "Any":
         if self._span is None:
             return None
 
+        if isinstance(self._span, StreamedSpan):
+            warnings.warn(
+                "Scope.transaction is not available in streaming mode.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return None
+
         # there is an orphan span on the scope
         if self._span.containing_transaction is None:
             return None
@@ -803,17 +823,36 @@ def transaction(self, value: "Any") -> None:
             "Assigning to scope.transaction directly is deprecated: use scope.set_transaction_name() instead."
         )
         self._transaction = value
-        if self._span and self._span.containing_transaction:
-            self._span.containing_transaction.name = value
+        if self._span:
+            if isinstance(self._span, StreamedSpan):
+                warnings.warn(
+                    "Scope.transaction is not available in streaming mode.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+                return None
+
+            if self._span.containing_transaction:
+                self._span.containing_transaction.name = value
 
     def set_transaction_name(self, name: str, source: "Optional[str]" = None) -> None:
         """Set the transaction name and optionally the transaction source."""
         self._transaction = name
-
-        if self._span and self._span.containing_transaction:
-            self._span.containing_transaction.name = name
-            if source:
-                self._span.containing_transaction.source = source
+        if self._span:
+            if isinstance(self._span, NoOpStreamedSpan):
+                return
+
+            elif isinstance(self._span, StreamedSpan):
+                self._span._segment.name = name
+                if source:
+                    self._span._segment.set_attribute(
+                        "sentry.span.source", getattr(source, "value", source)
+                    )
+
+            elif self._span.containing_transaction:
+                self._span.containing_transaction.name = name
+                if source:
+                    self._span.containing_transaction.source = source
 
         if source:
             self._transaction_info["source"] = source
@@ -836,12 +875,12 @@ def set_user(self, value: "Optional[Dict[str, Any]]") -> None:
             session.update(user=value)
 
     @property
-    def span(self) -> "Optional[Span]":
+    def span(self) -> "Optional[Union[Span, StreamedSpan]]":
         """Get/set current tracing span or transaction."""
         return self._span
 
     @span.setter
-    def span(self, span: "Optional[Span]") -> None:
+    def span(self, span: "Optional[Union[Span, StreamedSpan]]") -> None:
         self._span = span
         # XXX: this differs from the implementation in JS, there Scope.setSpan
         # does not set Scope._transactionName.
@@ -1150,6 +1189,15 @@ def start_span(
         be removed in the next major version. Going forward, it should only
         be used by the SDK itself.
         """
+        client = sentry_sdk.get_client()
+        if has_span_streaming_enabled(client.options):
+            warnings.warn(
+                "Scope.start_span is not available in streaming mode.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return NoOpSpan()
+
         if kwargs.get("description") is not None:
             warnings.warn(
                 "The `description` parameter is deprecated. Please use `name` instead.",
@@ -1169,6 +1217,9 @@ def start_span(
 
             # get current span or transaction
             span = self.span or self.get_isolation_scope().span
+            if isinstance(span, StreamedSpan):
+                # make mypy happy
+                return NoOpSpan()
 
             if span is None:
                 # New spans get the `trace_id` from the scope

sentry_sdk/traces.py

@@ -26,13 +26,17 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Optional, ParamSpec, TypeVar, Union
+    from typing import Any, Callable, Iterator, Optional, ParamSpec, TypeVar, Union
     from sentry_sdk._types import Attributes, AttributeValue
 
     P = ParamSpec("P")
     R = TypeVar("R")
 
 
+BAGGAGE_HEADER_NAME = "baggage"
+SENTRY_TRACE_HEADER_NAME = "sentry-trace"
+
+
 class SpanStatus(str, Enum):
     OK = "ok"
     ERROR = "error"
@@ -328,7 +332,7 @@ def finish(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> No
     def _start(self) -> None:
         if self._active:
             old_span = self._scope.span
-            self._scope.span = self  # type: ignore
+            self._scope.span = self
             self._previous_span_on_scope = old_span
 
     def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None:
@@ -463,6 +467,67 @@ def _update_active_thread(self) -> None:
             if thread_name is not None:
                 self.set_attribute(SPANDATA.THREAD_NAME, thread_name)
 
+    def _dynamic_sampling_context(self) -> "dict[str, str]":
+        return self._segment._get_baggage().dynamic_sampling_context()
+
+    def _to_traceparent(self) -> str:
+        if self.sampled is True:
+            sampled = "1"
+        elif self.sampled is False:
+            sampled = "0"
+        else:
+            sampled = None
+
+        traceparent = "%s-%s" % (self.trace_id, self.span_id)
+        if sampled is not None:
+            traceparent += "-%s" % (sampled,)
+
+        return traceparent
+
+    def _to_baggage(self) -> "Optional[Baggage]":
+        if self._segment:
+            return self._segment._get_baggage()
+        return None
+
+    def _get_baggage(self) -> "Baggage":
+        """
+        Return the :py:class:`~sentry_sdk.tracing_utils.Baggage` associated with
+        the segment.
+
+        The first time a new baggage with Sentry items is made, it will be frozen.
+        """
+        if not self._baggage or self._baggage.mutable:
+            self._baggage = Baggage.populate_from_segment(self)
+
+        return self._baggage
+
+    def _iter_headers(self) -> "Iterator[tuple[str, str]]":
+        if not self._segment:
+            return
+
+        yield SENTRY_TRACE_HEADER_NAME, self._to_traceparent()
+
+        baggage = self._segment._get_baggage().serialize()
+        if baggage:
+            yield BAGGAGE_HEADER_NAME, baggage
+
+    def _get_trace_context(self) -> "dict[str, Any]":
+        # Even if spans themselves are not event-based anymore, we need this
+        # to populate trace context on events
+        context: "dict[str, Any]" = {
+            "trace_id": self.trace_id,
+            "span_id": self.span_id,
+            "parent_span_id": self._parent_span_id,
+            "dynamic_sampling_context": self._dynamic_sampling_context(),
+        }
+
+        if "sentry.op" in self._attributes:
+            context["op"] = self._attributes["sentry.op"]
+        if "sentry.origin" in self._attributes:
+            context["origin"] = self._attributes["sentry.origin"]
+
+        return context
+
 
 class NoOpStreamedSpan(StreamedSpan):
     __slots__ = (
@@ -498,7 +563,7 @@ def _start(self) -> None:
             return
 
         old_span = self._scope.span
-        self._scope.span = self  # type: ignore
+        self._scope.span = self
         self._previous_span_on_scope = old_span
 
     def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None:

sentry_sdk/tracing.py

@@ -771,6 +771,12 @@ def update_active_thread(self) -> None:
         thread_id, thread_name = get_current_thread_meta()
         self.set_thread(thread_id, thread_name)
 
+    # Private aliases matching StreamedSpan's private API
+    _to_traceparent = to_traceparent
+    _to_baggage = to_baggage
+    _iter_headers = iter_headers
+    _get_trace_context = get_trace_context
+
 
 class Transaction(Span):
     """The Transaction is the root element that holds all the spans
@@ -1243,6 +1249,10 @@ def _set_initial_sampling_decision(
                 )
             )
 
+    # Private aliases matching StreamedSpan's private API
+    _get_baggage = get_baggage
+    _get_trace_context = get_trace_context
+
 
 class NoOpSpan(Span):
     def __repr__(self) -> str:
@@ -1324,6 +1334,13 @@ def _set_initial_sampling_decision(
     ) -> None:
         pass
 
+    # Private aliases matching StreamedSpan's private API
+    _to_traceparent = to_traceparent
+    _to_baggage = to_baggage
+    _get_baggage = get_baggage
+    _iter_headers = iter_headers
+    _get_trace_context = get_trace_context
+
 
 if TYPE_CHECKING: