diff --git a/sentry_sdk/integrations/fastapi.py b/sentry_sdk/integrations/fastapi.py index 10d4fd81ac..713c6d9efa 100644 --- a/sentry_sdk/integrations/fastapi.py +++ b/sentry_sdk/integrations/fastapi.py @@ -5,7 +5,7 @@ import sentry_sdk from sentry_sdk.integrations import DidNotEnable from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan +from sentry_sdk.traces import _is_sampled_streamed_span from sentry_sdk.tracing import SOURCE_FOR_STYLE, TransactionSource from sentry_sdk.tracing_utils import has_span_streaming_enabled from sentry_sdk.utils import transaction_from_function @@ -95,9 +95,7 @@ def _sentry_call(*args: "Any", **kwargs: "Any") -> "Any": if has_span_streaming_enabled(client.options): current_span = current_scope.streamed_span - if isinstance(current_span, StreamedSpan) and not isinstance( - current_span, NoOpStreamedSpan - ): + if _is_sampled_streamed_span(current_span): segment = current_span._segment segment._update_active_thread() diff --git a/sentry_sdk/integrations/starlette.py b/sentry_sdk/integrations/starlette.py index 68ff98cc02..b1134a436a 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, _get_current_streamed_span +from sentry_sdk.traces import _get_current_streamed_span, _is_sampled_streamed_span from sentry_sdk.tracing import ( SOURCE_FOR_STYLE, TransactionSource, @@ -255,12 +255,7 @@ def _set_request_body_data_on_streaming_segment( info: "Optional[Dict[str, Any]]", ) -> None: current_span = _get_current_streamed_span() - if ( - info - and "data" in info - and isinstance(current_span, StreamedSpan) - and not isinstance(current_span, NoOpStreamedSpan) - ): + if info and "data" in info and _is_sampled_streamed_span(current_span): with capture_internal_exceptions(): current_span._segment.set_attribute( "http.request.body.data", @@ -557,9 +552,7 @@ def _sentry_sync_func(*args: "Any", **kwargs: "Any") -> "Any": if span_streaming: current_span = current_scope.streamed_span - if isinstance(current_span, StreamedSpan) and not isinstance( - current_span, NoOpStreamedSpan - ): + if _is_sampled_streamed_span(current_span): 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 656fabc564..2da13158e3 100644 --- a/sentry_sdk/scope.py +++ b/sentry_sdk/scope.py @@ -34,7 +34,12 @@ _make_sampling_decision, PropagationContext, ) -from sentry_sdk.traces import _DEFAULT_PARENT_SPAN, StreamedSpan, NoOpStreamedSpan +from sentry_sdk.traces import ( + _DEFAULT_PARENT_SPAN, + StreamedSpan, + NoOpStreamedSpan, + _is_sampled_streamed_span, +) from sentry_sdk.tracing import ( BAGGAGE_HEADER_NAME, SENTRY_TRACE_HEADER_NAME, @@ -585,11 +590,7 @@ def get_traceparent(self, *args: "Any", **kwargs: "Any") -> "Optional[str]": span_streaming = has_span_streaming_enabled(client.options) # If we have an active span, return traceparent from there - if ( - span_streaming - and self.streamed_span is not None - and not isinstance(self.streamed_span, NoOpStreamedSpan) - ): + if span_streaming and _is_sampled_streamed_span(self.streamed_span): return self.streamed_span._to_traceparent() elif not span_streaming and self.span is not None: return self.span._to_traceparent() @@ -609,11 +610,7 @@ def get_baggage(self, *args: "Any", **kwargs: "Any") -> "Optional[Baggage]": span_streaming = has_span_streaming_enabled(client.options) # If we have an active span, return baggage from there - if ( - span_streaming - and self.streamed_span is not None - and not isinstance(self.streamed_span, NoOpStreamedSpan) - ): + if span_streaming and _is_sampled_streamed_span(self.streamed_span): return self.streamed_span._to_baggage() elif not span_streaming and self.span is not None: return self.span._to_baggage() @@ -918,11 +915,7 @@ def streamed_span(self, span: "Optional[StreamedSpan]") -> None: # Also set _transaction and _transaction_info in streaming mode as this # is used for populating events and linking them to segments - if ( - isinstance(span, StreamedSpan) - and not isinstance(span, NoOpStreamedSpan) - and span._is_segment() - ): + if _is_sampled_streamed_span(span) and span._is_segment(): self._transaction = span.name if span._attributes.get("sentry.span.source"): self._transaction_info["source"] = str( diff --git a/sentry_sdk/traces.py b/sentry_sdk/traces.py index 47c532c523..f31b8960af 100644 --- a/sentry_sdk/traces.py +++ b/sentry_sdk/traces.py @@ -31,7 +31,16 @@ ) if TYPE_CHECKING: - from typing import Any, Callable, Iterator, Optional, ParamSpec, TypeVar, Union + from typing import ( + Any, + Callable, + Iterator, + Optional, + ParamSpec, + TypeGuard, + TypeVar, + Union, + ) from sentry_sdk._types import Attributes, AttributeValue from sentry_sdk.profiler.continuous_profiler import ContinuousProfile @@ -757,6 +766,11 @@ def make_db_query(sql): return decorator +def _is_sampled_streamed_span(span: "Any") -> "TypeGuard[StreamedSpan]": + """Returns True if span is a StreamedSpan that was sampled (not a NoOpStreamedSpan).""" + return isinstance(span, StreamedSpan) and not isinstance(span, NoOpStreamedSpan) + + def _get_current_streamed_span( scope: "Optional[sentry_sdk.Scope]" = None, ) -> "Optional[StreamedSpan]":