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

Couple things going on in this PR. Bear with me, this is probably the
most all over the place span first PR because the outgoing trace
propagation changes make mypy complain about things elsewhere in the
sdk.

### 1. Outgoing trace propagation

Support getting trace propagation information from the span with
`_get_traceparent`, `_get_baggage`, `_iter_headers`, etc. These mirror
the old `Span` class to make integrating `StreamedSpan`s with the rest
of the SDK easier (since they're used throughout), with one difference:
they're explicitly private, while the corresponding `Span` methods were
public. Added aliases to them so that we can use the private methods
everywhere.

There is definite clean up potential here once we get rid of the old
spans and we no longer have to make the streaming span interface work
with the existing helper scope methods.

### 2. Addressing cascading mypy issues

Now that we're officially allowing `StreamedSpan`s to be set on the
scope, a LOT of type hints need updating all over the SDK. In many
places, I've added explicit guards against functionality that doesn't
exist in span first mode. This should prevent folks from using the wrong
APIs in the wrong SDK mode (tracing vs. static) as well as make mypy
happy.

---------

Co-authored-by: Erica Pisani <pisani.erica@gmail.com>

Author: sentrivana

sentry_sdk/_span_batcher.py

@@ -114,9 +114,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

@@ -12,6 +12,8 @@
 
 import sentry_sdk
 from sentry_sdk.utils import logger
+from sentry_sdk.traces import StreamedSpan
+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
@@ -523,7 +525,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 isinstance(current_span, StreamedSpan):
+        # 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/integrations/celery/__init__.py

@@ -14,7 +14,7 @@
 )
 from sentry_sdk.integrations.celery.utils import _now_seconds_since_epoch
 from sentry_sdk.integrations.logging import ignore_logger
-from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, TransactionSource
+from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span, TransactionSource
 from sentry_sdk.tracing_utils import Baggage
 from sentry_sdk.utils import (
     capture_internal_exceptions,
@@ -34,7 +34,6 @@
     from typing import Union
 
     from sentry_sdk._types import EventProcessor, Event, Hint, ExcInfo
-    from sentry_sdk.tracing import Span
 
     F = TypeVar("F", bound=Callable[..., Any])
 
@@ -100,7 +99,10 @@ def _set_status(status: str) -> None:
     with capture_internal_exceptions():
         scope = sentry_sdk.get_current_scope()
         if scope.span is not None:
-            scope.span.set_status(status)
+            if isinstance(scope.span, Span):
+                scope.span.set_status(status)
+            else:
+                scope.span.status = "ok" if status == "ok" else "error"
 
 
 def _capture_exception(task: "Any", exc_info: "ExcInfo") -> None:

sentry_sdk/integrations/openai_agents/utils.py

@@ -11,6 +11,7 @@
 from sentry_sdk.consts import SPANDATA, SPANSTATUS, OP
 from sentry_sdk.integrations import DidNotEnable
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing import Span
 from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import event_from_exception, safe_serialize
 from sentry_sdk.ai._openai_completions_api import _transform_system_instructions
@@ -22,10 +23,10 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Union
     from agents import Usage, TResponseInputItem
 
-    from sentry_sdk.tracing import Span
+    from sentry_sdk.traces import StreamedSpan
     from sentry_sdk._types import TextPart
 
 try:
@@ -46,8 +47,15 @@ def _capture_exception(exc: "Any") -> None:
     sentry_sdk.capture_event(event, hint=hint)
 
 
-def _record_exception_on_span(span: "Span", error: Exception) -> "Any":
+def _record_exception_on_span(
+    span: "Union[Span, StreamedSpan]", error: Exception
+) -> "Any":
     set_span_errored(span)
+
+    if not isinstance(span, Span):
+        # TODO[span-first]: make this work with streamedspans
+        return
+
     span.set_data("span.status", "error")
 
     # Optionally capture the error details if we have them

sentry_sdk/scope.py

@@ -581,8 +581,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()
@@ -595,8 +599,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()
@@ -605,8 +613,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()
@@ -670,8 +682,12 @@ def iter_trace_propagation_headers(
         span = kwargs.pop("span", None)
         span = span or self.span
 
-        if has_tracing_enabled(client.options) and span is not None:
-            for header in span.iter_headers():
+        if (
+            has_tracing_enabled(client.options)
+            and span is not None
+            and not isinstance(span, NoOpStreamedSpan)
+        ):
+            for header in span._iter_headers():
                 yield header
         elif has_external_propagation_context():
             # when we have an external_propagation_context (otlp)
@@ -718,7 +734,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
 
@@ -772,6 +788,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
@@ -801,17 +825,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
@@ -834,12 +877,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.
@@ -1148,6 +1191,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.",
@@ -1167,6 +1219,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