feat(quart): Add span streaming support to Quart integration

sentry_sdk/integrations/quart.py

@@ -9,7 +9,10 @@
 from sentry_sdk.integrations._wsgi_common import _filter_headers
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.tracing import SOURCE_FOR_STYLE
+from sentry_sdk.traces import SOURCE_FOR_STYLE as SEGMENT_SOURCE_FOR_STYLE
+from sentry_sdk.traces import StreamedSpan, get_current_span
+from sentry_sdk.tracing import SOURCE_FOR_STYLE as TRANSACTION_SOURCE_FOR_STYLE
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     ensure_integration_enabled,
@@ -144,9 +147,16 @@ def _set_transaction_name_and_source(
             "url": request.url_rule.rule,
             "endpoint": request.url_rule.endpoint,
         }
+
+        source = (
+            SEGMENT_SOURCE_FOR_STYLE[transaction_style]
+            if has_span_streaming_enabled(sentry_sdk.get_client().options)
+            else TRANSACTION_SOURCE_FOR_STYLE[transaction_style]
+        )
+
         scope.set_transaction_name(
-            name_for_style[transaction_style],
-            source=SOURCE_FOR_STYLE[transaction_style],
+            name=name_for_style[transaction_style],
+            source=source,
         )
     except Exception:
         pass
@@ -169,6 +179,39 @@ async def _request_websocket_started(app: "Quart", **kwargs: "Any") -> None:
     )
 
     scope = sentry_sdk.get_isolation_scope()
+
+    if has_span_streaming_enabled(sentry_sdk.get_client().options):
+        current_span = get_current_span()
+        if type(current_span) is StreamedSpan:
+            segment = current_span._segment
+
+            segment.set_attribute("http.request.method", request_websocket.method)
+            header_attributes: "dict[str, Any]" = {}
+
+            for header, header_value in _filter_headers(
+                dict(request_websocket.headers), use_annotated_value=False
+            ).items():
+                header_attributes[f"http.request.header.{header.lower()}"] = (
+                    header_value
+                )
+
+            segment.set_attributes(header_attributes)
+
+            if should_send_default_pii():
+                segment.set_attribute("url.full", request_websocket.url)
+                segment.set_attribute(
+                    "url.query",
+                    request_websocket.query_string.decode("utf-8", errors="replace"),
+                )
+
+                if len(request_websocket.access_route) >= 1:
+                    segment.set_attribute(
+                        "client.address", request_websocket.access_route[0]
+                    )
+                    segment.set_attribute(
+                        "user.ip_address", request_websocket.access_route[0]
+                    )
+
     evt_processor = _make_request_event_processor(app, request_websocket, integration)
     scope.add_event_processor(evt_processor)
 
@@ -194,7 +237,8 @@ def inner(event: "Event", hint: "dict[str, Any]") -> "Event":
             request_info["headers"] = _filter_headers(dict(request.headers))
 
             if should_send_default_pii():
-                request_info["env"] = {"REMOTE_ADDR": request.access_route[0]}
+                if len(request.access_route) >= 1:
+                    request_info["env"] = {"REMOTE_ADDR": request.access_route[0]}
                 _add_user_to_event(event)
 
         return event

tests/integrations/quart/test_quart.py

@@ -14,6 +14,7 @@
     set_tag,
 )
 from sentry_sdk.integrations.logging import LoggingIntegration
+from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE
 
 
 def quart_app_factory():
@@ -647,3 +648,246 @@ async def test_span_origin(sentry_init, capture_events):
     (_, event) = events
 
     assert event["contexts"]["trace"]["origin"] == "auto.http.quart"
+
+
+@pytest.mark.asyncio
+async def test_span_streaming_basic(sentry_init, capture_items):
+    sentry_init(
+        integrations=[quart_sentry.QuartIntegration()],
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+    items = capture_items("span")
+
+    app = quart_app_factory()
+    client = app.test_client()
+    response = await client.get("/message")
+    assert response.status_code == 200
+
+    sentry_sdk.flush()
+
+    spans = [item.payload for item in items]
+    assert len(spans) == 1
+
+    segment = spans[0]
+    assert segment["is_segment"] is True
+    assert "parent_span_id" not in segment
+    assert segment["status"] == "ok"
+    assert segment["attributes"]["sentry.op"] == "http.server"
+    assert segment["attributes"]["sentry.origin"] == "auto.http.quart"
+    assert segment["attributes"]["http.request.method"] == "GET"
+    assert segment["name"] == "hi"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "url,transaction_style,expected_name,expected_source",
+    [
+        ("/message", "endpoint", "hi", "component"),
+        ("/message", "url", "/message", "route"),
+        ("/message/123456", "endpoint", "hi_with_id", "component"),
+        ("/message/123456", "url", "/message/<message_id>", "route"),
+    ],
+)
+async def test_span_streaming_transaction_style(
+    sentry_init,
+    capture_items,
+    url,
+    transaction_style,
+    expected_name,
+    expected_source,
+):
+    sentry_init(
+        integrations=[
+            quart_sentry.QuartIntegration(transaction_style=transaction_style)
+        ],
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+    items = capture_items("span")
+
+    app = quart_app_factory()
+    client = app.test_client()
+    response = await client.get(url)
+    assert response.status_code == 200
+
+    sentry_sdk.flush()
+
+    spans = [item.payload for item in items]
+    assert len(spans) == 1
+
+    segment = spans[0]
+    assert segment["is_segment"] is True
+    assert segment["name"] == expected_name
+    assert segment["attributes"]["sentry.span.source"] == expected_source
+
+
+@pytest.mark.asyncio
+async def test_span_streaming_with_error(sentry_init, capture_items):
+    sentry_init(
+        integrations=[quart_sentry.QuartIntegration()],
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+    items = capture_items("event", "span")
+
+    app = quart_app_factory()
+
+    @app.route("/error")
+    async def error():
+        1 / 0
+
+    client = app.test_client()
+    try:
+        await client.get("/error")
+    except ZeroDivisionError:
+        pass
+
+    sentry_sdk.flush()
+
+    events = [item.payload for item in items if item.type == "event"]
+    spans = [item.payload for item in items if item.type == "span"]
+
+    assert len(events) == 1
+    assert len(spans) == 1
+
+    error_event = events[0]
+    segment = spans[0]
+
+    assert segment["trace_id"] == error_event["contexts"]["trace"]["trace_id"]
+    assert segment["is_segment"] is True
+    assert segment["status"] == "error"
+
+    assert "parent_span_id" not in segment
+
+    assert error_event["contexts"]["trace"]["span_id"] == segment["span_id"]
+    assert error_event["exception"]["values"][0]["mechanism"]["type"] == "quart"
+    assert error_event["exception"]["values"][0]["mechanism"]["handled"] is False
+
+
+@pytest.mark.asyncio
+async def test_span_streaming_request_attributes_no_pii(sentry_init, capture_items):
+    sentry_init(
+        integrations=[quart_sentry.QuartIntegration()],
+        traces_sample_rate=1.0,
+        send_default_pii=False,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+    items = capture_items("span")
+
+    app = quart_app_factory()
+    client = app.test_client()
+    response = await client.get("/message?foo=bar")
+    assert response.status_code == 200
+
+    sentry_sdk.flush()
+
+    spans = [item.payload for item in items]
+    assert len(spans) == 1
+
+    segment = spans[0]
+    assert segment["attributes"]["http.request.method"] == "GET"
+    assert "http.request.header.host" in segment["attributes"]
+
+    assert "url.full" not in segment["attributes"]
+    assert "url.query" not in segment["attributes"]
+    assert "client.address" not in segment["attributes"]
+    assert "user.ip_address" not in segment["attributes"]
+
+
+@pytest.mark.asyncio
+async def test_span_streaming_request_attributes_with_pii(sentry_init, capture_items):
+    sentry_init(
+        integrations=[quart_sentry.QuartIntegration()],
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+    items = capture_items("span")
+
+    app = quart_app_factory()
+    client = app.test_client()
+    response = await client.get("/message?foo=bar&baz=qux")
+    assert response.status_code == 200
+
+    sentry_sdk.flush()
+
+    spans = [item.payload for item in items]
+    assert len(spans) == 1
+
+    segment = spans[0]
+    assert segment["attributes"]["http.request.method"] == "GET"
+    assert "http.request.header.host" in segment["attributes"]
+
+    assert (
+        segment["attributes"]["url.full"] == "http://localhost/message?foo=bar&baz=qux"
+    )
+    assert segment["attributes"]["url.query"] == "foo=bar&baz=qux"
+    assert "client.address" in segment["attributes"]
+    assert "user.ip_address" in segment["attributes"]
+
+
+@pytest.mark.asyncio
+async def test_span_streaming_sensitive_header_scrubbing(sentry_init, capture_items):
+    sentry_init(
+        integrations=[quart_sentry.QuartIntegration()],
+        traces_sample_rate=1.0,
+        send_default_pii=False,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+    items = capture_items("span")
+
+    app = quart_app_factory()
+    client = app.test_client()
+    response = await client.get(
+        "/message",
+        headers={
+            "Authorization": "Bearer secret-token",
+            "X-Custom-Header": "passthrough",
+        },
+    )
+    assert response.status_code == 200
+
+    sentry_sdk.flush()
+
+    spans = [item.payload for item in items]
+    assert len(spans) == 1
+
+    segment = spans[0]
+    assert (
+        segment["attributes"]["http.request.header.authorization"]
+        == SENSITIVE_DATA_SUBSTITUTE
+    )
+    assert segment["attributes"]["http.request.header.x-custom-header"] == "passthrough"
+
+
+@pytest.mark.asyncio
+async def test_span_streaming_sensitive_header_passthrough_with_pii(
+    sentry_init, capture_items
+):
+    sentry_init(
+        integrations=[quart_sentry.QuartIntegration()],
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+    items = capture_items("span")
+
+    app = quart_app_factory()
+    client = app.test_client()
+    response = await client.get(
+        "/message",
+        headers={"Authorization": "Bearer secret-token"},
+    )
+    assert response.status_code == 200
+
+    sentry_sdk.flush()
+
+    spans = [item.payload for item in items]
+    assert len(spans) == 1
+
+    segment = spans[0]
+    assert (
+        segment["attributes"]["http.request.header.authorization"]
+        == "Bearer secret-token"
+    )