fix: Introduce _get_current_streamed_span() to keep types backwards compatible (#6177)

Create streaming-lifecycle variants of the `tracing_utils.get_current_span()` function and the `Scope.span` property. The existing functions are narrowed to be backwards compatible for users of the transaction-based span lifecycle, allowing users to upgrade their SDK version without their type checkers raising errors.

Note that this breaks all current uses of these functions when the streaming lifecycle option is enabled (which is experimental). All call sites internal to the SDK are updated.

The `Scope.span` property returns `None` if a `StreamedSpan` is held, and `Scope.streamed_span` returns `None` if a `Span` is held on the scope.

Author: alexander-alderman-webb

sentry_sdk/api.py

@@ -6,7 +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.traces import StreamedSpan, _get_current_streamed_span
 from sentry_sdk.tracing import NoOpSpan, Transaction, trace
 from sentry_sdk.crons import monitor
 
@@ -422,7 +422,7 @@ def set_measurement(name: str, value: float, unit: "MeasurementUnit" = "") -> No
 
 def get_current_span(
     scope: "Optional[Scope]" = None,
-) -> "Optional[Union[Span, StreamedSpan]]":
+) -> "Optional[Span]":
     """
     Returns the currently active span if there is one running, otherwise `None`
     """
@@ -533,12 +533,7 @@ def update_current_span(
             attributes={"user_id": 123, "batch_size": 50}
         )
     """
-    current_span = get_current_span()
-
-    if current_span is None:
-        return
-
-    if isinstance(current_span, StreamedSpan):
+    if isinstance(_get_current_streamed_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 "
@@ -548,6 +543,11 @@ def update_current_span(
         )
         return
 
+    current_span = get_current_span()
+
+    if current_span is None:
+        return
+
     if op is not None:
         current_span.op = op
 

sentry_sdk/integrations/celery/__init__.py

@@ -15,7 +15,7 @@
 from sentry_sdk.integrations.celery.utils import _now_seconds_since_epoch
 from sentry_sdk.integrations.logging import ignore_logger
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.traces import StreamedSpan
+from sentry_sdk.traces import StreamedSpan, _get_current_streamed_span
 from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span, TransactionSource
 from sentry_sdk.tracing_utils import Baggage, has_span_streaming_enabled
 from sentry_sdk.utils import (
@@ -98,13 +98,16 @@ def setup_once() -> None:
 
 
 def _set_status(status: str) -> None:
+    client = sentry_sdk.get_client()
+    span_streaming = has_span_streaming_enabled(client.options)
+
     with capture_internal_exceptions():
         scope = sentry_sdk.get_current_scope()
-        if scope.span is not None:
-            if isinstance(scope.span, Span):
-                scope.span.set_status(status)
-            else:
-                scope.span.status = "ok" if status == "ok" else "error"
+
+        if span_streaming and scope.streamed_span is not None:
+            scope.streamed_span.status = "ok" if status == "ok" else "error"
+        elif not span_streaming and scope.span is not None:
+            scope.span.set_status(status)
 
 
 def _capture_exception(task: "Any", exc_info: "ExcInfo") -> None:
@@ -289,7 +292,7 @@ def apply_async(*args: "Any", **kwargs: "Any") -> "Any":
 
         span_mgr: "Union[StreamedSpan, Span, NoOpMgr]" = NoOpMgr()
         if span_streaming:
-            if not task_started_from_beat and sentry_sdk.get_current_span() is not None:
+            if not task_started_from_beat and _get_current_streamed_span() is not None:
                 span_mgr = sentry_sdk.traces.start_span(
                     name=task_name,
                     attributes={
@@ -570,7 +573,7 @@ def sentry_publish(self: "Producer", *args: "Any", **kwargs: "Any") -> "Any":
 
         span: "Union[StreamedSpan, Span, None]" = None
         if span_streaming:
-            if sentry_sdk.get_current_span() is not None:
+            if _get_current_streamed_span() is not None:
                 span = sentry_sdk.traces.start_span(
                     name=task_name,
                     attributes={

sentry_sdk/integrations/fastapi.py

@@ -90,13 +90,17 @@ def _sentry_get_request_handler(*args: "Any", **kwargs: "Any") -> "Any":
             @wraps(old_call)
             def _sentry_call(*args: "Any", **kwargs: "Any") -> "Any":
                 current_scope = sentry_sdk.get_current_scope()
-                current_span = current_scope.span
 
-                if isinstance(current_span, StreamedSpan) and not isinstance(
-                    current_span, NoOpStreamedSpan
-                ):
-                    segment = current_span._segment
-                    segment._update_active_thread()
+                client = sentry_sdk.get_client()
+                if has_span_streaming_enabled(client.options):
+                    current_span = current_scope.streamed_span
+
+                    if isinstance(current_span, StreamedSpan) and not isinstance(
+                        current_span, NoOpStreamedSpan
+                    ):
+                        segment = current_span._segment
+                        segment._update_active_thread()
+
                 elif current_scope.transaction is not None:
                     current_scope.transaction.update_active_thread()
 

sentry_sdk/integrations/starlette.py

@@ -21,7 +21,7 @@
 )
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan
+from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan, _get_current_streamed_span
 from sentry_sdk.tracing import (
     SOURCE_FOR_STYLE,
     TransactionSource,
@@ -254,7 +254,7 @@ def _default(value: "Any") -> "Any":
 def _set_request_body_data_on_streaming_segment(
     info: "Optional[Dict[str, Any]]",
 ) -> None:
-    current_span = sentry_sdk.get_current_span()
+    current_span = _get_current_streamed_span()
     if (
         info
         and "data" in info
@@ -545,19 +545,22 @@ def event_processor(
 
             @functools.wraps(old_func)
             def _sentry_sync_func(*args: "Any", **kwargs: "Any") -> "Any":
-                integration = sentry_sdk.get_client().get_integration(
-                    StarletteIntegration
-                )
+                client = sentry_sdk.get_client()
+
+                integration = client.get_integration(StarletteIntegration)
                 if integration is None:
                     return old_func(*args, **kwargs)
 
                 current_scope = sentry_sdk.get_current_scope()
-                current_span = current_scope.span
 
-                if isinstance(current_span, StreamedSpan) and not isinstance(
-                    current_span, NoOpStreamedSpan
-                ):
-                    current_span._segment._update_active_thread()
+                span_streaming = has_span_streaming_enabled(client.options)
+                if span_streaming:
+                    current_span = current_scope.streamed_span
+
+                    if isinstance(current_span, StreamedSpan) and not isinstance(
+                        current_span, NoOpStreamedSpan
+                    ):
+                        current_span._segment._update_active_thread()
                 elif current_scope.transaction is not None:
                     current_scope.transaction.update_active_thread()
 

sentry_sdk/scope.py

@@ -580,12 +580,18 @@ def get_traceparent(self, *args: "Any", **kwargs: "Any") -> "Optional[str]":
         """
         client = self.get_client()
 
+        if not has_tracing_enabled(client.options):
+            return self.get_active_propagation_context().to_traceparent()
+
+        span_streaming = has_span_streaming_enabled(client.options)
         # If we have an active span, return traceparent from there
         if (
-            has_tracing_enabled(client.options)
-            and self.span is not None
-            and not isinstance(self.span, NoOpStreamedSpan)
+            span_streaming
+            and self.streamed_span is not None
+            and not isinstance(self.streamed_span, NoOpStreamedSpan)
         ):
+            return self.streamed_span._to_traceparent()
+        elif not span_streaming and self.span is not None:
             return self.span._to_traceparent()
 
         # else return traceparent from the propagation context
@@ -598,12 +604,18 @@ def get_baggage(self, *args: "Any", **kwargs: "Any") -> "Optional[Baggage]":
         """
         client = self.get_client()
 
+        if not has_tracing_enabled(client.options):
+            return self.get_active_propagation_context().get_baggage()
+
+        span_streaming = has_span_streaming_enabled(client.options)
         # If we have an active span, return baggage from there
         if (
-            has_tracing_enabled(client.options)
-            and self.span is not None
-            and not isinstance(self.span, NoOpStreamedSpan)
+            span_streaming
+            and self.streamed_span is not None
+            and not isinstance(self.streamed_span, NoOpStreamedSpan)
         ):
+            return self.streamed_span._to_baggage()
+        elif not span_streaming and self.span is not None:
             return self.span._to_baggage()
 
         # else return baggage from the propagation context
@@ -680,7 +692,9 @@ def iter_trace_propagation_headers(
             return
 
         span = kwargs.pop("span", None)
-        span = span or self.span
+        if not span:
+            span_streaming = has_span_streaming_enabled(client.options)
+            span = self.streamed_span if span_streaming else self.span
 
         if (
             has_tracing_enabled(client.options)
@@ -877,12 +891,12 @@ def set_user(self, value: "Optional[Dict[str, Any]]") -> None:
             session.update(user=value)
 
     @property
-    def span(self) -> "Optional[Union[Span, StreamedSpan]]":
+    def span(self) -> "Optional[Span]":
         """Get/set current tracing span or transaction."""
-        return self._span
+        return self._span if isinstance(self._span, Span) else None
 
     @span.setter
-    def span(self, span: "Optional[Union[Span, StreamedSpan]]") -> None:
+    def span(self, span: "Optional[Span]") -> None:
         self._span = span
         # XXX: this differs from the implementation in JS, there Scope.setSpan
         # does not set Scope._transactionName.
@@ -893,6 +907,15 @@ def span(self, span: "Optional[Union[Span, StreamedSpan]]") -> None:
                 if transaction.source:
                     self._transaction_info["source"] = transaction.source
 
+    @property
+    def streamed_span(self) -> "Optional[StreamedSpan]":
+        """Get/set current tracing span."""
+        return self._span if isinstance(self._span, StreamedSpan) else None
+
+    @streamed_span.setter
+    def streamed_span(self, span: "Optional[StreamedSpan]") -> None:
+        self._span = span
+
         # Also set _transaction and _transaction_info in streaming mode as this
         # is used for populating events and linking them to segments
         if (
@@ -1267,7 +1290,7 @@ def start_streamed_span(
         if parent_span is _DEFAULT_PARENT_SPAN or isinstance(
             parent_span, NoOpStreamedSpan
         ):
-            parent_span = self.span  # type: ignore
+            parent_span = self.streamed_span
 
         # If no eligible parent_span was provided and there is no currently
         # active span, this is a segment

sentry_sdk/traces.py

@@ -341,8 +341,8 @@ 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
+            old_span = self._scope.streamed_span
+            self._scope.streamed_span = self
             self._previous_span_on_scope = old_span
 
     def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None:
@@ -360,7 +360,7 @@ def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None
             with capture_internal_exceptions():
                 old_span = self._previous_span_on_scope
                 del self._previous_span_on_scope
-                self._scope.span = old_span
+                self._scope.streamed_span = old_span
 
         # Set attributes from the segment. These are set on span end on purpose
         # so that we have the best chance to capture the segment's final name
@@ -586,8 +586,8 @@ def _start(self) -> None:
         if self._scope is None:
             return
 
-        old_span = self._scope.span
-        self._scope.span = self
+        old_span = self._scope.streamed_span
+        self._scope.streamed_span = self
         self._previous_span_on_scope = old_span
 
     def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None:
@@ -610,7 +610,7 @@ def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None
             with capture_internal_exceptions():
                 old_span = self._previous_span_on_scope
                 del self._previous_span_on_scope
-                self._scope.span = old_span
+                self._scope.streamed_span = old_span
 
         self._finished = True
 
@@ -755,3 +755,17 @@ def make_db_query(sql):
         return decorator(func)
     else:
         return decorator
+
+
+def _get_current_streamed_span(
+    scope: "Optional[sentry_sdk.Scope]" = None,
+) -> "Optional[StreamedSpan]":
+    """
+    Returns the currently active span on the scope if the span is a `StreamedSpan`, otherwise `None`.
+
+    This function will only return a non-`None` value when the streaming trace lifecycle is enabled.
+    To enable the lifecycle, pass `_experiments={"trace_lifecycle": "stream"}` to `sentry.init()`.
+    """
+    scope = scope or sentry_sdk.get_current_scope()
+    current_span = scope.streamed_span
+    return current_span

sentry_sdk/tracing_utils.py

@@ -1194,7 +1194,7 @@ def sync_wrapper(*args: "Any", **kwargs: "Any") -> "Any":
 
 def get_current_span(
     scope: "Optional[sentry_sdk.Scope]" = None,
-) -> "Optional[Union[Span, StreamedSpan]]":
+) -> "Optional[Span]":
     """
     Returns the currently active span if there is one running, otherwise `None`
     """
@@ -1208,9 +1208,16 @@ def set_span_errored(span: "Optional[Union[Span, StreamedSpan]]" = None) -> None
     Set the status of the current or given span to INTERNAL_ERROR.
     Also sets the status of the transaction (root span) to INTERNAL_ERROR.
     """
-    from sentry_sdk.traces import StreamedSpan, SpanStatus
+    from sentry_sdk.traces import StreamedSpan, SpanStatus, _get_current_streamed_span
 
-    span = span or get_current_span()
+    client = sentry_sdk.get_client()
+
+    if not span:
+        span = (
+            _get_current_streamed_span()
+            if has_span_streaming_enabled(client.options)
+            else sentry_sdk.get_current_span()
+        )
 
     if span is not None:
         if isinstance(span, Span):

tests/integrations/celery/test_celery.py

@@ -8,6 +8,7 @@
 
 import sentry_sdk
 from sentry_sdk import start_transaction, get_current_span
+from sentry_sdk.traces import _get_current_streamed_span
 from sentry_sdk.integrations.celery import (
     CeleryIntegration,
     _wrap_task_run,
@@ -663,7 +664,11 @@ def test_sentry_propagate_traces_override(span_streaming, init_celery):
 
     @celery.task(name="dummy_task", bind=True)
     def dummy_task(self, message):
-        trace_id = get_current_span().trace_id
+        trace_id = (
+            _get_current_streamed_span().trace_id
+            if span_streaming
+            else get_current_span().trace_id
+        )
         return trace_id
 
     if span_streaming:

tests/test_api.py

@@ -21,6 +21,8 @@
     get_isolation_scope,
 )
 
+from sentry_sdk.tracing import Span
+
 from sentry_sdk.client import Client, NonRecordingClient
 from tests.conftest import TestTransportWithOptions
 
@@ -40,7 +42,7 @@ def test_get_current_span_default_hub(sentry_init):
     assert get_current_span() is None
 
     scope = get_current_scope()
-    fake_span = mock.MagicMock()
+    fake_span = Span()
     scope.span = fake_span
 
     assert get_current_span() == fake_span