ref: Replace _get_current_streamed_span with get_current_span

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: sentrivana

sentry_sdk/api.py

@@ -8,7 +8,8 @@
 from sentry_sdk.consts import INSTRUMENTER
 from sentry_sdk.crons import monitor
 from sentry_sdk.scope import Scope, _ScopeManager, isolation_scope, new_scope
-from sentry_sdk.traces import StreamedSpan, _get_current_streamed_span
+from sentry_sdk.traces import StreamedSpan
+from sentry_sdk.traces import get_current_span as _get_current_streamed_span
 from sentry_sdk.tracing import NoOpSpan, Transaction, trace
 
 if TYPE_CHECKING:

sentry_sdk/feature_flags.py

@@ -65,7 +65,7 @@ def add_feature_flag(flag: str, result: bool) -> None:
     flags.set(flag, result)
 
     if has_span_streaming_enabled(client.options):
-        span = sentry_sdk.traces._get_current_streamed_span()
+        span = sentry_sdk.traces.get_current_span()
         if span and isinstance(span, sentry_sdk.traces.StreamedSpan):
             span.set_attribute(f"flag.evaluation.{flag}", result)
 

sentry_sdk/integrations/celery/__init__.py

@@ -16,7 +16,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 Scope, should_send_default_pii
-from sentry_sdk.traces import StreamedSpan, _get_current_streamed_span
+from sentry_sdk.traces import StreamedSpan, get_current_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 (
@@ -286,7 +286,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 _get_current_streamed_span() is not None:
+            if not task_started_from_beat and get_current_span() is not None:
                 span_mgr = sentry_sdk.traces.start_span(
                     name=task_name,
                     attributes={
@@ -567,7 +567,7 @@ def sentry_publish(self: "Producer", *args: "Any", **kwargs: "Any") -> "Any":
 
         span: "Union[StreamedSpan, Span, None]" = None
         if span_streaming:
-            if _get_current_streamed_span() is not None:
+            if get_current_span() is not None:
                 span = sentry_sdk.traces.start_span(
                     name=task_name,
                     attributes={

sentry_sdk/integrations/starlette.py

@@ -22,7 +22,7 @@
 )
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.traces import StreamedSpan, _get_current_streamed_span
+from sentry_sdk.traces import StreamedSpan, get_current_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 = _get_current_streamed_span()
+    current_span = get_current_span()
     if info and "data" in info and type(current_span) is StreamedSpan:
         with capture_internal_exceptions():
             current_span._segment.set_attribute(

tests/integrations/celery/test_celery.py

@@ -8,13 +8,13 @@
 from celery.bin import worker
 
 import sentry_sdk
+import sentry_sdk.traces
 from sentry_sdk import get_current_span, start_transaction
 from sentry_sdk.integrations.celery import (
     CeleryIntegration,
     _wrap_task_run,
 )
 from sentry_sdk.integrations.celery.beat import _get_headers
-from sentry_sdk.traces import _get_current_streamed_span
 from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE
 from tests.conftest import ApproxDict
 
@@ -667,7 +667,7 @@ 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_streamed_span().trace_id
+            sentry_sdk.traces.get_current_span().trace_id
             if span_streaming
             else get_current_span().trace_id
         )

tests/integrations/rust_tracing/test_rust_tracing.py

@@ -87,15 +87,13 @@ def test_on_new_span_on_close(
         with sentry_sdk.traces.start_span(name="custom parent"):
             rust_tracing.new_span(RustTracingLevel.Info, 3)
 
-            sentry_first_rust_span = sentry_sdk.traces._get_current_streamed_span()
+            sentry_first_rust_span = sentry_sdk.traces.get_current_span()
             rust_first_rust_span = rust_tracing.spans[3]
 
             assert sentry_first_rust_span == rust_first_rust_span
 
             rust_tracing.close_span(3)
-            assert (
-                sentry_sdk.traces._get_current_streamed_span() != sentry_first_rust_span
-            )
+            assert sentry_sdk.traces.get_current_span() != sentry_first_rust_span
 
         sentry_sdk.flush()
         spans = [item.payload for item in items if item.type == "span"]
@@ -173,32 +171,30 @@ def test_nested_on_new_span_on_close(
     if span_streaming:
         items = capture_items("event", "transaction", "span")
         with sentry_sdk.traces.start_span(name="custom parent"):
-            original_sentry_span = sentry_sdk.traces._get_current_streamed_span()
+            original_sentry_span = sentry_sdk.traces.get_current_span()
 
             rust_tracing.new_span(RustTracingLevel.Info, 3, index_arg=10)
-            sentry_first_rust_span = sentry_sdk.traces._get_current_streamed_span()
+            sentry_first_rust_span = sentry_sdk.traces.get_current_span()
             rust_first_rust_span = rust_tracing.spans[3]
 
             # Use a different `index_arg` value for the inner span to help
             # distinguish the two at the end of the test
             rust_tracing.new_span(RustTracingLevel.Info, 5, index_arg=9)
-            sentry_second_rust_span = sentry_sdk.traces._get_current_streamed_span()
+            sentry_second_rust_span = sentry_sdk.traces.get_current_span()
             rust_second_rust_span = rust_tracing.spans[5]
 
             assert rust_second_rust_span == sentry_second_rust_span
 
             rust_tracing.close_span(5)
 
             # Ensure the current sentry span was moved back to the parent
-            sentry_span_after_close = sentry_sdk.traces._get_current_streamed_span()
+            sentry_span_after_close = sentry_sdk.traces.get_current_span()
             assert sentry_span_after_close == sentry_first_rust_span
             assert sentry_span_after_close == rust_first_rust_span
 
             rust_tracing.close_span(3)
 
-            assert (
-                sentry_sdk.traces._get_current_streamed_span() == original_sentry_span
-            )
+            assert sentry_sdk.traces.get_current_span() == original_sentry_span
 
         sentry_sdk.flush()
         spans = [item.payload for item in items if item.type == "span"]
@@ -312,10 +308,10 @@ def test_on_new_span_without_transaction(sentry_init, span_streaming):
     )
 
     if span_streaming:
-        assert sentry_sdk.traces._get_current_streamed_span() is None
+        assert sentry_sdk.traces.get_current_span() is None
 
         rust_tracing.new_span(RustTracingLevel.Info, 3)
-        current_span = sentry_sdk.traces._get_current_streamed_span()
+        current_span = sentry_sdk.traces.get_current_span()
         assert current_span is not None
         assert current_span._segment is current_span
     else:
@@ -584,24 +580,22 @@ def span_filter(metadata: Dict[str, object]) -> bool:
     if span_streaming:
         items = capture_items("event", "transaction", "span")
         with sentry_sdk.traces.start_span(name="custom parent"):
-            original_sentry_span = sentry_sdk.traces._get_current_streamed_span()
+            original_sentry_span = sentry_sdk.traces.get_current_span()
 
             # Span is not ignored
             rust_tracing.new_span(RustTracingLevel.Info, 3, index_arg=10)
-            info_span = sentry_sdk.traces._get_current_streamed_span()
+            info_span = sentry_sdk.traces.get_current_span()
 
             # Span is ignored, current span should remain the same
             rust_tracing.new_span(RustTracingLevel.Trace, 5, index_arg=9)
-            assert sentry_sdk.traces._get_current_streamed_span() == info_span
+            assert sentry_sdk.traces.get_current_span() == info_span
 
             # Closing the filtered span should leave the current span alone
             rust_tracing.close_span(5)
-            assert sentry_sdk.traces._get_current_streamed_span() == info_span
+            assert sentry_sdk.traces.get_current_span() == info_span
 
             rust_tracing.close_span(3)
-            assert (
-                sentry_sdk.traces._get_current_streamed_span() == original_sentry_span
-            )
+            assert sentry_sdk.traces.get_current_span() == original_sentry_span
 
         sentry_sdk.flush()
         spans = [item.payload for item in items if item.type == "span"]
@@ -652,16 +646,12 @@ def test_record(sentry_init, span_streaming):
         with sentry_sdk.traces.start_span(name="custom parent"):
             rust_tracing.new_span(RustTracingLevel.Info, 3)
 
-            span_before_record = (
-                sentry_sdk.traces._get_current_streamed_span()._to_json()
-            )
+            span_before_record = sentry_sdk.traces.get_current_span()._to_json()
             assert span_before_record["attributes"]["version"] == "None"
 
             rust_tracing.record(3)
 
-            span_after_record = (
-                sentry_sdk.traces._get_current_streamed_span()._to_json()
-            )
+            span_after_record = sentry_sdk.traces.get_current_span()._to_json()
             assert span_after_record["attributes"]["version"] == "memoized"
     else:
         with start_transaction():
@@ -699,18 +689,14 @@ def span_filter(metadata: Dict[str, object]) -> bool:
         with sentry_sdk.traces.start_span(name="custom parent"):
             rust_tracing.new_span(RustTracingLevel.Info, 3)
 
-            span_before_record = (
-                sentry_sdk.traces._get_current_streamed_span()._to_json()
-            )
+            span_before_record = sentry_sdk.traces.get_current_span()._to_json()
             assert span_before_record["attributes"]["version"] == "None"
 
             rust_tracing.new_span(RustTracingLevel.Trace, 5)
             rust_tracing.record(5)
 
             # `on_record()` should not do anything to the current Sentry span if the associated Rust span was ignored
-            span_after_record = (
-                sentry_sdk.traces._get_current_streamed_span()._to_json()
-            )
+            span_after_record = sentry_sdk.traces.get_current_span()._to_json()
             assert span_after_record["attributes"]["version"] == "None"
     else:
         with start_transaction():
@@ -764,19 +750,15 @@ def test_include_tracing_fields(
         with sentry_sdk.traces.start_span(name="custom parent"):
             rust_tracing.new_span(RustTracingLevel.Info, 3)
 
-            span_before_record = (
-                sentry_sdk.traces._get_current_streamed_span()._to_json()
-            )
+            span_before_record = sentry_sdk.traces.get_current_span()._to_json()
             if tracing_fields_expected:
                 assert span_before_record["attributes"]["version"] == "None"
             else:
                 assert span_before_record["attributes"]["version"] == "[Filtered]"
 
             rust_tracing.record(3)
 
-            span_after_record = (
-                sentry_sdk.traces._get_current_streamed_span()._to_json()
-            )
+            span_after_record = sentry_sdk.traces.get_current_span()._to_json()
 
             if tracing_fields_expected:
                 assert span_after_record["attributes"] == {