fix(tracing): Restore scope transaction after streamed segment ends

Author: kcons

sentry_sdk/scope.py

@@ -933,6 +933,8 @@ def streamed_span(self, span: "Optional[StreamedSpan]") -> None:
                 self._transaction_info["source"] = str(
                     span._attributes["sentry.segment.name.source"]
                 )
+            else:
+                self._transaction_info.pop("source", None)
             return
 
         if type(span) is NoOpStreamedSpan:
@@ -942,6 +944,8 @@ def streamed_span(self, span: "Optional[StreamedSpan]") -> None:
                 self._transaction_info["source"] = str(
                     span._attributes["sentry.segment.name.source"]
                 )
+            else:
+                self._transaction_info.pop("source", None)
 
     @property
     def profile(self) -> "Optional[Profile]":

sentry_sdk/traces.py

@@ -254,6 +254,8 @@ class StreamedSpan:
         "_status",
         "_scope",
         "_previous_span_on_scope",
+        "_previous_transaction_on_scope",
+        "_previous_transaction_info_on_scope",
         "_baggage",
         "_sample_rand",
         "_sample_rate",
@@ -364,6 +366,11 @@ def finish(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> No
     def _start(self) -> None:
         if self._active:
             old_span = self._scope.streamed_span
+            if self._is_segment():
+                self._previous_transaction_on_scope = self._scope._transaction
+                self._previous_transaction_info_on_scope = (
+                    self._scope._transaction_info.copy()
+                )
             self._scope.streamed_span = self
             self._previous_span_on_scope = old_span
 
@@ -383,6 +390,13 @@ def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None
                 old_span = self._previous_span_on_scope
                 del self._previous_span_on_scope
                 self._scope.streamed_span = old_span
+                if self._is_segment():
+                    self._scope._transaction = self._previous_transaction_on_scope
+                    del self._previous_transaction_on_scope
+                    self._scope._transaction_info = (
+                        self._previous_transaction_info_on_scope
+                    )
+                    del self._previous_transaction_info_on_scope
 
         # 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
@@ -691,6 +705,11 @@ def _start(self) -> None:
             return
 
         old_span = self._scope.streamed_span
+        if self._is_segment():
+            self._previous_transaction_on_scope = self._scope._transaction
+            self._previous_transaction_info_on_scope = (
+                self._scope._transaction_info.copy()
+            )
         self._scope.streamed_span = self
         self._previous_span_on_scope = old_span
 
@@ -715,6 +734,13 @@ def _end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None
                 old_span = self._previous_span_on_scope
                 del self._previous_span_on_scope
                 self._scope.streamed_span = old_span
+                if self._is_segment():
+                    self._scope._transaction = self._previous_transaction_on_scope
+                    del self._previous_transaction_on_scope
+                    self._scope._transaction_info = (
+                        self._previous_transaction_info_on_scope
+                    )
+                    del self._previous_transaction_info_on_scope
 
         self._finished = True
 

tests/tracing/test_integration_tests.py

@@ -110,6 +110,69 @@ def test_error_event_linked_without_performance_span_streaming(
     assert event["transaction_info"] == {"source": "custom segment source"}
 
 
+@pytest.mark.parametrize("sample_rate", [0.0, 1.0])
+def test_span_streaming_clears_transaction_name(
+    sentry_init, capture_events, sample_rate
+):
+    sentry_init(traces_sample_rate=sample_rate, trace_lifecycle="stream")
+    events = capture_events()
+    scope = sentry_sdk.get_current_scope()
+
+    with sentry_sdk.traces.start_span(
+        name="streamed segment",
+        attributes={"sentry.segment.name.source": "component"},
+    ):
+        assert scope._transaction == "streamed segment"
+        assert scope._transaction_info == {"source": "component"}
+
+    assert scope._transaction is None
+    assert scope._transaction_info == {}
+
+    sentry_sdk.capture_message("after streamed segment")
+
+    (event,) = events
+    assert "transaction" not in event
+
+
+@pytest.mark.parametrize("sample_rate", [0.0, 1.0])
+def test_span_streaming_restores_transaction_name(
+    sentry_init, capture_events, sample_rate
+):
+    sentry_init(traces_sample_rate=sample_rate, trace_lifecycle="stream")
+    events = capture_events()
+    scope = sentry_sdk.get_current_scope()
+    scope.set_transaction_name("previous transaction", source="route")
+
+    with sentry_sdk.traces.start_span(name="streamed segment"):
+        assert scope._transaction == "streamed segment"
+        assert scope._transaction_info == ({"source": "custom"} if sample_rate else {})
+
+    assert scope._transaction == "previous transaction"
+    assert scope._transaction_info == {"source": "route"}
+
+    sentry_sdk.capture_message("after streamed segment")
+
+    (event,) = events
+    assert event["transaction"] == "previous transaction"
+    assert event["transaction_info"] == {"source": "route"}
+
+
+@pytest.mark.parametrize("sample_rate", [0.0, 1.0])
+def test_span_streaming_keeps_name_changed_during_child_span(sentry_init, sample_rate):
+    sentry_init(traces_sample_rate=sample_rate, trace_lifecycle="stream")
+    scope = sentry_sdk.get_current_scope()
+
+    with sentry_sdk.traces.start_span(name="streamed segment"):
+        with sentry_sdk.traces.start_span(name="child"):
+            scope.set_transaction_name("renamed segment", source="component")
+
+        assert scope._transaction == "renamed segment"
+        assert scope._transaction_info == {"source": "component"}
+
+    assert scope._transaction is None
+    assert scope._transaction_info == {}
+
+
 @pytest.mark.parametrize("parent_sampled", [True, False, None])
 @pytest.mark.parametrize("sample_rate", [0.0, 1.0])
 def test_continue_trace(sentry_init, capture_envelopes, parent_sampled, sample_rate):