fix: Introduce _get_current_streamed_span() to keep types backwards compatible

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