diff --git a/sentry_sdk/api.py b/sentry_sdk/api.py index 105f531e5f..23f003557c 100644 --- a/sentry_sdk/api.py +++ b/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 diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py index 6c033bf3d1..6eb295d81c 100644 --- a/sentry_sdk/integrations/celery/__init__.py +++ b/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={ diff --git a/sentry_sdk/integrations/fastapi.py b/sentry_sdk/integrations/fastapi.py index 12874a10b3..10d4fd81ac 100644 --- a/sentry_sdk/integrations/fastapi.py +++ b/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() diff --git a/sentry_sdk/integrations/starlette.py b/sentry_sdk/integrations/starlette.py index 3312604b69..68ff98cc02 100644 --- a/sentry_sdk/integrations/starlette.py +++ b/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() diff --git a/sentry_sdk/scope.py b/sentry_sdk/scope.py index f24edcf137..b880d4c91f 100644 --- a/sentry_sdk/scope.py +++ b/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 diff --git a/sentry_sdk/traces.py b/sentry_sdk/traces.py index 485be2a3e2..47c532c523 100644 --- a/sentry_sdk/traces.py +++ b/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 diff --git a/sentry_sdk/tracing_utils.py b/sentry_sdk/tracing_utils.py index 034af715d8..33b1a33814 100644 --- a/sentry_sdk/tracing_utils.py +++ b/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): diff --git a/tests/integrations/celery/test_celery.py b/tests/integrations/celery/test_celery.py index 8e1689296f..4255dac211 100644 --- a/tests/integrations/celery/test_celery.py +++ b/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: diff --git a/tests/test_api.py b/tests/test_api.py index bd6b245841..a066ce1172 100644 --- a/tests/test_api.py +++ b/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