compat, making mypy happy

Author: sentrivana

sentry_sdk/api.py

@@ -37,6 +37,7 @@
         LogLevelStr,
         SamplingContext,
     )
+    from sentry_sdk.traces import StreamedSpan
     from sentry_sdk.tracing import Span, TransactionKwargs
 
     T = TypeVar("T")
@@ -409,7 +410,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 +528,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

@@ -584,7 +584,7 @@ def get_traceparent(self, *args: "Any", **kwargs: "Any") -> "Optional[str]":
 
         # 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()
+            return self.span._to_traceparent()
 
         # else return traceparent from the propagation context
         return self.get_active_propagation_context().to_traceparent()
@@ -598,7 +598,7 @@ def get_baggage(self, *args: "Any", **kwargs: "Any") -> "Optional[Baggage]":
 
         # 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()
+            return self.span._to_baggage()
 
         # else return baggage from the propagation context
         return self.get_active_propagation_context().get_baggage()
@@ -608,7 +608,7 @@ 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()
+            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 +673,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 +720,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 +774,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 +811,39 @@ 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:
+            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 self._span and 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
 
         if source:
             self._transaction_info["source"] = source
@@ -836,12 +866,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 +1180,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 +1208,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

@@ -332,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:
@@ -563,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,8 @@ def _set_initial_sampling_decision(
                 )
             )
 
+    _get_baggage = get_baggage
+
 
 class NoOpSpan(Span):
     def __repr__(self) -> str:
@@ -1324,6 +1332,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:
 

sentry_sdk/tracing_utils.py

@@ -935,6 +935,14 @@ async def async_wrapper(*args: "Any", **kwargs: "Any") -> "Any":
                 )
                 return await f(*args, **kwargs)
 
+            if isinstance(current_span, StreamedSpan):
+                warnings.warn(
+                    "Use the @sentry_sdk.traces.trace decorator in span streaming mode.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+                return await f(*args, **kwargs)
+
             span_op = op or _get_span_op(template)
             function_name = name or qualname_from_function(f) or ""
             span_name = _get_span_name(template, function_name, kwargs)
@@ -972,6 +980,14 @@ def sync_wrapper(*args: "Any", **kwargs: "Any") -> "Any":
                 )
                 return f(*args, **kwargs)
 
+            if isinstance(current_span, StreamedSpan):
+                warnings.warn(
+                    "Use the @sentry_sdk.traces.trace decorator in span streaming mode.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+                return f(*args, **kwargs)
+
             span_op = op or _get_span_op(template)
             function_name = name or qualname_from_function(f) or ""
             span_name = _get_span_name(template, function_name, kwargs)
@@ -1072,7 +1088,9 @@ def sync_wrapper(*args: "Any", **kwargs: "Any") -> "Any":
     return span_decorator
 
 
-def get_current_span(scope: "Optional[sentry_sdk.Scope]" = None) -> "Optional[Span]":
+def get_current_span(
+    scope: "Optional[sentry_sdk.Scope]" = None,
+) -> "Optional[Union[Span, StreamedSpan]]":
     """
     Returns the currently active span if there is one running, otherwise `None`
     """
@@ -1081,16 +1099,24 @@ def get_current_span(scope: "Optional[sentry_sdk.Scope]" = None) -> "Optional[Sp
     return current_span
 
 
-def set_span_errored(span: "Optional[Span]" = None) -> None:
+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
+
     span = span or get_current_span()
+
     if span is not None:
-        span.set_status(SPANSTATUS.INTERNAL_ERROR)
-        if span.containing_transaction is not None:
-            span.containing_transaction.set_status(SPANSTATUS.INTERNAL_ERROR)
+        if isinstance(span, Span):
+            span.set_status(SPANSTATUS.INTERNAL_ERROR)
+            if span.containing_transaction is not None:
+                span.containing_transaction.set_status(SPANSTATUS.INTERNAL_ERROR)
+        elif isinstance(span, StreamedSpan):
+            span.status = SpanStatus.ERROR
+            if span._segment is not None:
+                span._segment.status = SpanStatus.ERROR
 
 
 def _generate_sample_rand(

tests/tracing/test_span_streaming.py

@@ -678,6 +678,31 @@ def test_continue_trace_no_sample_rand(sentry_init, capture_envelopes):
     assert segment["trace_id"] == trace_id
 
 
+def test_outgoing_traceparent_and_baggage(sentry_init, capture_envelopes):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    sentry_sdk.traces.new_trace()
+
+    with sentry_sdk.traces.start_span(name="span") as span:
+        assert span.sampled is True
+
+        trace_id = span.trace_id
+        span_id = span.span_id
+
+        traceparent = sentry_sdk.get_traceparent()
+        assert traceparent == f"{trace_id}-{span_id}-1"
+
+        baggage = sentry_sdk.get_baggage()
+        baggage_items = dict(tuple(item.split("=")) for item in baggage.split(","))
+        assert "sentry-trace_id" in baggage_items
+        assert baggage_items["sentry-trace_id"] == trace_id
+        assert "sentry-sampled" in baggage_items
+        assert baggage_items["sentry-sampled"] == "true"
+
+
 def test_trace_decorator(sentry_init, capture_envelopes):
     sentry_init(
         traces_sample_rate=1.0,