feat(starlette): Support span streaming

sentry_sdk/integrations/starlette.py

@@ -1,5 +1,6 @@
 import asyncio
 import functools
+import json
 import warnings
 from collections.abc import Set
 from copy import deepcopy
@@ -20,10 +21,12 @@
 )
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan
 from sentry_sdk.tracing import (
     SOURCE_FOR_STYLE,
     TransactionSource,
 )
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
     AnnotatedValue,
     capture_internal_exceptions,
@@ -147,7 +150,8 @@
         send: "Callable[[Dict[str, Any]], Awaitable[None]]",
         **kwargs: "Any",
     ) -> None:
-        integration = sentry_sdk.get_client().get_integration(StarletteIntegration)
+        client = sentry_sdk.get_client()
+        integration = client.get_integration(StarletteIntegration)
         if integration is None:
             return await old_call(app, scope, receive, send, **kwargs)
 
@@ -164,22 +168,38 @@
             return await old_call(app, scope, receive, send, **kwargs)
 
         middleware_name = app.__class__.__name__
+        is_span_streaming_enabled = has_span_streaming_enabled(client.options)
+
+        def _start_middleware_span(op: str, name: str) -> "Any":
+            if is_span_streaming_enabled:
+                return sentry_sdk.traces.start_span(
+                    name=name,
+                    attributes={
+                        "sentry.op": op,
+                        "sentry.origin": StarletteIntegration.origin,
+                        "middleware.name": middleware_name,
+                    },
+                )
+            return sentry_sdk.start_span(
+                op=op,
+                name=name,
+                origin=StarletteIntegration.origin,
+            )
 
-        with sentry_sdk.start_span(
-            op=OP.MIDDLEWARE_STARLETTE,
-            name=middleware_name,
-            origin=StarletteIntegration.origin,
+        with _start_middleware_span(
+            op=OP.MIDDLEWARE_STARLETTE, name=middleware_name
         ) as middleware_span:
-            middleware_span.set_tag("starlette.middleware_name", middleware_name)
+            if not is_span_streaming_enabled:
+                middleware_span.set_tag("starlette.middleware_name", middleware_name)
 
             # Creating spans for the "receive" callback
             async def _sentry_receive(*args: "Any", **kwargs: "Any") -> "Any":
-                with sentry_sdk.start_span(
+                with _start_middleware_span(
                     op=OP.MIDDLEWARE_STARLETTE_RECEIVE,
                     name=getattr(receive, "__qualname__", str(receive)),
-                    origin=StarletteIntegration.origin,
                 ) as span:
-                    span.set_tag("starlette.middleware_name", middleware_name)
+                    if not is_span_streaming_enabled:
+                        span.set_tag("starlette.middleware_name", middleware_name)
                     return await receive(*args, **kwargs)
 
             receive_name = getattr(receive, "__name__", str(receive))
@@ -188,12 +208,12 @@
 
             # Creating spans for the "send" callback
             async def _sentry_send(*args: "Any", **kwargs: "Any") -> "Any":
-                with sentry_sdk.start_span(
+                with _start_middleware_span(
                     op=OP.MIDDLEWARE_STARLETTE_SEND,
                     name=getattr(send, "__qualname__", str(send)),
-                    origin=StarletteIntegration.origin,
                 ) as span:
-                    span.set_tag("starlette.middleware_name", middleware_name)
+                    if not is_span_streaming_enabled:
+                        span.set_tag("starlette.middleware_name", middleware_name)
                     return await send(*args, **kwargs)
 
             send_name = getattr(send, "__name__", str(send))
@@ -214,6 +234,16 @@
     return middleware_class
 
 
+def _serialize_body_data(data: "Any") -> str:
+    # data may be a JSON-serializable value, an AnnotatedValue, or a dict with AnnotatedValue values
+    def _default(value: "Any") -> "Any":
+        if isinstance(value, AnnotatedValue):
+            return {"value": value.value, "metadata": value.metadata}
+        return str(value)
+
+    return json.dumps(data, default=_default)
+
+
 @ensure_integration_enabled(StarletteIntegration)
 def _capture_exception(exception: BaseException, handled: "Any" = False) -> None:
     event, hint = event_from_exception(
@@ -439,9 +469,8 @@
         if is_coroutine:
 
             async def _sentry_async_func(*args: "Any", **kwargs: "Any") -> "Any":
-                integration = sentry_sdk.get_client().get_integration(
-                    StarletteIntegration
-                )
+                client = sentry_sdk.get_client()
+                integration = client.get_integration(StarletteIntegration)
                 if integration is None:
                     return await old_func(*args, **kwargs)
 
@@ -481,6 +510,22 @@
                     _make_request_event_processor(request, integration)
                 )
 
+                is_span_streaming_enabled = has_span_streaming_enabled(client.options)
+                if is_span_streaming_enabled:
+                    current_span = sentry_sdk.get_current_span()
+
+                    if (
+                        info
+                        and "data" in info
+                        and isinstance(current_span, StreamedSpan)
+                        and not isinstance(current_span, NoOpStreamedSpan)
+                    ):
+                        data = info["data"]
+                        current_span._segment.set_attribute(
+                            "http.request.body.data",
+                            _serialize_body_data(data),
+                        )
+
                 return await old_func(*args, **kwargs)
 
             func = _sentry_async_func
@@ -496,7 +541,13 @@
                     return old_func(*args, **kwargs)
 
                 current_scope = sentry_sdk.get_current_scope()
-                if current_scope.transaction is not None:
+                current_span = current_scope.span
+
+                if isinstance(current_span, StreamedSpan) and not isinstance(
+                    current_span, NoOpStreamedSpan
+                ):
+                    current_span._segment._update_active_thread()
+                elif current_scope.transaction is not None:
                     current_scope.transaction.update_active_thread()
 
                 sentry_scope = sentry_sdk.get_isolation_scope()