limit api to traces.py and mark private

Author: alexander-alderman-webb

sentry_sdk/__init__.py

@@ -32,7 +32,6 @@
     "get_isolation_scope",
     "get_current_scope",
     "get_current_span",
-    "get_current_streamed_span",
     "get_traceparent",
     "is_initialized",
     "isolation_scope",

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
 
@@ -67,7 +67,6 @@ def overload(x: "T") -> "T":
     "get_isolation_scope",
     "get_current_scope",
     "get_current_span",
-    "get_current_streamed_span",
     "get_traceparent",
     "is_initialized",
     "isolation_scope",
@@ -430,15 +429,6 @@ def get_current_span(
     return tracing_utils.get_current_span(scope)
 
 
-def get_current_streamed_span(
-    scope: "Optional[Scope]" = None,
-) -> "Optional[StreamedSpan]":
-    """
-    Returns the currently active streamed span if there is one running, otherwise `None`
-    """
-    return tracing_utils.get_current_streamed_span(scope)
-
-
 def get_traceparent() -> "Optional[str]":
     """
     Returns the traceparent either from the active span or from the scope.
@@ -543,7 +533,7 @@ def update_current_span(
             attributes={"user_id": 123, "batch_size": 50}
         )
     """
-    if isinstance(get_current_streamed_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 "

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 (
@@ -292,10 +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_streamed_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={
@@ -576,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_streamed_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/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_streamed_span()
+    current_span = _get_current_streamed_span()
     if (
         info
         and "data" in info

sentry_sdk/traces.py

@@ -755,3 +755,14 @@ 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 if there is one running, otherwise `None`
+    """
+    scope = scope or sentry_sdk.get_current_scope()
+    current_span = scope.streamed_span
+    return current_span

sentry_sdk/tracing_utils.py

@@ -35,6 +35,7 @@
     _module_in_list,
 )
 from sentry_sdk.tracing import Span as LegacySpan
+from sentry_sdk.traces import _get_current_streamed_span
 
 from typing import TYPE_CHECKING
 
@@ -1203,25 +1204,21 @@ def get_current_span(
     return current_span
 
 
-def get_current_streamed_span(
-    scope: "Optional[sentry_sdk.Scope]" = None,
-) -> "Optional[StreamedSpan]":
-    """
-    Returns the currently active span if there is one running, otherwise `None`
-    """
-    scope = scope or sentry_sdk.get_current_scope()
-    current_span = scope.streamed_span
-    return current_span
-
-
 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()
+    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

@@ -7,7 +7,8 @@
 from celery.bin import worker
 
 import sentry_sdk
-from sentry_sdk import start_transaction, get_current_span, get_current_streamed_span
+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,
@@ -664,7 +665,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
+            _get_current_streamed_span().trace_id
             if span_streaming
             else get_current_span().trace_id
         )