feat(litestar): Set name and source on request span when streaming (#6553)

Call `Scope.set_transaction_name()` on isolation and current scopes instead of manually overwriting event fields in an event processor.

This ensures that the span name and the `sentry.span.source` attribute are set on the segment span when the streaming trace lifecycle is enabled.

Author: alexander-alderman-webb

sentry_sdk/integrations/litestar.py

@@ -287,40 +287,37 @@ async def handle_wrapper(
 
         request_data = await body
 
-        def event_processor(event: "Event", _: "Hint") -> "Event":
-            route_handler = scope.get("route_handler")
+        route_handler = scope.get("route_handler")
+
+        func = None
+        if route_handler.name is not None:
+            name = route_handler.name
+        # Accounts for use of type `Ref` in earlier versions of litestar without the need to reference it as a type
+        elif hasattr(route_handler.fn, "value"):
+            func = route_handler.fn.value
+        else:
+            func = route_handler.fn
+        if func is not None:
+            name = transaction_from_function(func)
+
+        source = SOURCE_FOR_STYLE["endpoint"]
 
+        if not name:
+            name = _DEFAULT_TRANSACTION_NAME
+            source = TransactionSource.ROUTE
+
+        sentry_sdk.set_transaction_name(name, source)
+        sentry_scope.set_transaction_name(name, source)
+
+        def event_processor(event: "Event", _: "Hint") -> "Event":
             request_info = event.get("request", {})
             request_info["content_length"] = len(scope.get("_body", b""))
             if should_send_default_pii():
                 request_info["cookies"] = extracted_request_data["cookies"]
             if request_data is not None:
                 request_info["data"] = request_data
 
-            func = None
-            if route_handler.name is not None:
-                tx_name = route_handler.name
-            # Accounts for use of type `Ref` in earlier versions of litestar without the need to reference it as a type
-            elif hasattr(route_handler.fn, "value"):
-                func = route_handler.fn.value
-            else:
-                func = route_handler.fn
-            if func is not None:
-                tx_name = transaction_from_function(func)
-
-            tx_info = {"source": SOURCE_FOR_STYLE["endpoint"]}
-
-            if not tx_name:
-                tx_name = _DEFAULT_TRANSACTION_NAME
-                tx_info = {"source": TransactionSource.ROUTE}
-
-            event.update(
-                {
-                    "request": deepcopy(request_info),
-                    "transaction": tx_name,
-                    "transaction_info": tx_info,
-                }
-            )
+            event["request"] = deepcopy(request_info)
             return event
 
         sentry_scope._name = LitestarIntegration.identifier

tests/integrations/litestar/test_litestar.py

@@ -138,6 +138,69 @@ def test_catch_exceptions(
     assert event["exception"]["values"][0]["mechanism"]["type"] == "litestar"
 
 
+@pytest.mark.parametrize(
+    "test_url,expected_tx_name",
+    [
+        (
+            "/some_url",
+            "tests.integrations.litestar.test_litestar.litestar_app_factory.<locals>.homepage_handler",
+        ),
+        (
+            "/custom_error",
+            "custom_name",
+        ),
+        (
+            "/controller/error",
+            "tests.integrations.litestar.test_litestar.litestar_app_factory.<locals>.MyController.controller_error",
+        ),
+    ],
+)
+@pytest.mark.parametrize("span_streaming", [True, False])
+def test_transaction_name_and_source(
+    sentry_init,
+    capture_events,
+    test_url,
+    expected_tx_name,
+    capture_items,
+    span_streaming,
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[LitestarIntegration()],
+        _experiments={
+            "trace_lifecycle": "stream" if span_streaming else "static",
+        },
+    )
+    litestar_app = litestar_app_factory()
+    client = TestClient(litestar_app)
+
+    if span_streaming:
+        items = capture_items("span")
+
+        try:
+            client.get(test_url)
+        except Exception:
+            pass
+
+        sentry_sdk.flush()
+        spans = [item.payload for item in items]
+
+        spans = [span for span in spans if expected_tx_name in span["name"]]
+        assert len(spans) == 1
+        assert spans[0]["attributes"]["sentry.span.source"] == "component"
+    else:
+        events = capture_events()
+
+        try:
+            client.get(test_url)
+        except Exception:
+            pass
+
+        (_, transaction) = events
+        assert expected_tx_name in transaction["transaction"]
+        assert transaction["transaction_info"] == {"source": "component"}
+
+
 @pytest.mark.parametrize("span_streaming", [True, False])
 def test_middleware_spans(
     sentry_init,