diff --git a/sentry_sdk/integrations/rust_tracing.py b/sentry_sdk/integrations/rust_tracing.py index e2b203a286..697cf7db60 100644 --- a/sentry_sdk/integrations/rust_tracing.py +++ b/sentry_sdk/integrations/rust_tracing.py @@ -32,7 +32,7 @@ import json from enum import Enum, auto -from typing import Any, Callable, Dict, Tuple, Optional +from typing import Any, Callable, Dict, Optional import sentry_sdk from sentry_sdk.integrations import Integration @@ -40,8 +40,6 @@ from sentry_sdk.tracing import Span as SentrySpan from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE -TraceState = Optional[Tuple[Optional[SentrySpan], SentrySpan]] - class RustTracingLevel(Enum): Trace = "TRACE" @@ -170,7 +168,7 @@ def _include_tracing_fields(self) -> bool: else self.include_tracing_fields ) - def on_event(self, event: str, _span_state: "TraceState") -> None: + def on_event(self, event: str, sentry_span: "SentrySpan") -> None: deserialized_event = json.loads(event) metadata = deserialized_event.get("metadata", {}) @@ -184,7 +182,7 @@ def on_event(self, event: str, _span_state: "TraceState") -> None: elif event_type == EventTypeMapping.Event: process_event(deserialized_event) - def on_new_span(self, attrs: str, span_id: str) -> "TraceState": + def on_new_span(self, attrs: str, span_id: str) -> "Optional[SentrySpan]": attrs = json.loads(attrs) metadata = attrs.get("metadata", {}) @@ -210,13 +208,7 @@ def on_new_span(self, attrs: str, span_id: str) -> "TraceState": "origin": self.origin, } - scope = sentry_sdk.get_current_scope() - parent_sentry_span = scope.span - if parent_sentry_span: - sentry_span = parent_sentry_span.start_child(**kwargs) - else: - sentry_span = scope.start_span(**kwargs) - + sentry_span = sentry_sdk.start_span(**kwargs) fields = metadata.get("fields", []) for field in fields: if self._include_tracing_fields(): @@ -224,21 +216,18 @@ def on_new_span(self, attrs: str, span_id: str) -> "TraceState": else: sentry_span.set_data(field, SENSITIVE_DATA_SUBSTITUTE) - scope.span = sentry_span - return (parent_sentry_span, sentry_span) + sentry_span.__enter__() + return sentry_span - def on_close(self, span_id: str, span_state: "TraceState") -> None: - if span_state is None: + def on_close(self, span_id: str, sentry_span: "SentrySpan") -> None: + if sentry_span is None: return - parent_sentry_span, sentry_span = span_state - sentry_span.finish() - sentry_sdk.get_current_scope().span = parent_sentry_span + sentry_span.__exit__(None, None, None) - def on_record(self, span_id: str, values: str, span_state: "TraceState") -> None: - if span_state is None: + def on_record(self, span_id: str, values: str, sentry_span: "SentrySpan") -> None: + if sentry_span is None: return - _parent_sentry_span, sentry_span = span_state deserialized_values = json.loads(values) for key, value in deserialized_values.items(): diff --git a/tests/integrations/rust_tracing/test_rust_tracing.py b/tests/integrations/rust_tracing/test_rust_tracing.py index 893fc86966..a3e367f7e9 100644 --- a/tests/integrations/rust_tracing/test_rust_tracing.py +++ b/tests/integrations/rust_tracing/test_rust_tracing.py @@ -78,7 +78,7 @@ def test_on_new_span_on_close(sentry_init, capture_events): rust_tracing.new_span(RustTracingLevel.Info, 3) sentry_first_rust_span = sentry_sdk.get_current_span() - _, rust_first_rust_span = rust_tracing.spans[3] + rust_first_rust_span = rust_tracing.spans[3] assert sentry_first_rust_span == rust_first_rust_span @@ -120,24 +120,22 @@ def test_nested_on_new_span_on_close(sentry_init, capture_events): rust_tracing.new_span(RustTracingLevel.Info, 3, index_arg=10) sentry_first_rust_span = sentry_sdk.get_current_span() - _, rust_first_rust_span = rust_tracing.spans[3] + 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.get_current_span() - rust_parent_span, rust_second_rust_span = rust_tracing.spans[5] + rust_second_rust_span = rust_tracing.spans[5] assert rust_second_rust_span == sentry_second_rust_span - assert rust_parent_span == sentry_first_rust_span - assert rust_parent_span == rust_first_rust_span - assert rust_parent_span != rust_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.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)