fix(starlette): Stop eagerly consuming request bodies for streamed spans

sentry_sdk/consts.py

@@ -819,6 +819,12 @@ class SPANDATA:
     Example: GET
     """
 
+    HTTP_REQUEST_BODY_DATA = "http.request.body.data"
+    """
+    HTTP request body data. Can be given as string or structural data of any format.
+    Example: "[{\"role\": \"user\", \"message\": \"hello\"}]"
+    """
+
     HTTP_REQUEST_METHOD = "http.request.method"
     """
     The HTTP method used.

sentry_sdk/integrations/starlette.py

@@ -8,7 +8,8 @@
 from typing import TYPE_CHECKING
 
 import sentry_sdk
-from sentry_sdk.consts import OP
+from sentry_sdk._types import OVER_SIZE_LIMIT_SUBSTITUTE
+from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import (
     _DEFAULT_FAILED_REQUEST_STATUS_CODES,
     DidNotEnable,
@@ -46,7 +47,10 @@
     import starlette  # type: ignore
     from starlette import __version__ as STARLETTE_VERSION
     from starlette.applications import Starlette  # type: ignore
-    from starlette.datastructures import UploadFile  # type: ignore
+    from starlette.datastructures import (  # type: ignore
+        FormData,
+        UploadFile,
+    )
     from starlette.middleware import Middleware  # type: ignore
     from starlette.middleware.authentication import (  # type: ignore
         AuthenticationMiddleware,
@@ -88,6 +92,9 @@
 
 TRANSACTION_STYLE_VALUES = ("endpoint", "url")
 
+_SCOPE_STATE_JSON_REQUEST_BODY_KEY = "sentry_sdk.json_request_body"
+_SCOPE_STATE_FORMDATA_REQUEST_BODY_KEY = "sentry_sdk.formdata_request_body"
+
 
 class StarletteIntegration(Integration):
     identifier = "starlette"
@@ -146,6 +153,16 @@ def setup_once() -> None:
         if version >= (0, 24):
             patch_templates()
 
+    def setup_once_with_options(
+        self: "StarletteIntegration", options: "Optional[dict[str, Any]]" = None
+    ) -> None:
+        is_span_streaming_enabled = has_span_streaming_enabled(options)
+        if not is_span_streaming_enabled:
+            return
+
+        _patch_json_request_body_accessor()
+        _patch_formdata_request_body_accessor()
+
 
 def _enable_span_for_middleware(middleware_class: "Any") -> type:
     old_call = middleware_class.__call__
@@ -478,61 +495,152 @@ def _is_async_callable(obj: "Any") -> bool:
     )
 
 
-def patch_request_response() -> None:
-    old_request_response = starlette.routing.request_response
+def _patch_json_request_body_accessor() -> None:
+    """
+    Caches request body data on the ASGI scope, so that the body can be attached to telemetry after the request handler runs.
+    Without the cache, consuming the stream can cause the application to hang.
+    """
+    _original_json = Request.json
 
-    def _sentry_request_response(func: "Callable[[Any], Any]") -> "ASGIApp":
-        old_func = func
+    @functools.wraps(_original_json)
+    async def sentry_json(self: "Request", *args: "Any", **kwargs: "Any") -> "Any":
+        request_json = await _original_json(self, *args, **kwargs)
+        self.scope.setdefault("state", {})[_SCOPE_STATE_JSON_REQUEST_BODY_KEY] = (
+            request_json
+        )
+        return request_json
 
-        is_coroutine = _is_async_callable(old_func)
-        if is_coroutine:
+    Request.json = sentry_json
 
-            async def _sentry_async_func(*args: "Any", **kwargs: "Any") -> "Any":
-                client = sentry_sdk.get_client()
-                integration = client.get_integration(StarletteIntegration)
-                if integration is None:
-                    return await old_func(*args, **kwargs)
 
-                request = args[0]
+def _patch_formdata_request_body_accessor() -> None:
+    """
+    Caches request body data on the ASGI scope, so that the body can be attached to telemetry after the request handler runs.
+    Without the cache, consuming the stream can cause the application to hang.
+    """
+    _original_form = Request.form
 
-                _set_transaction_name_and_source(
-                    sentry_sdk.get_current_scope(),
-                    integration.transaction_style,
-                    request,
-                )
+    @functools.wraps(_original_form)
+    async def sentry_form(self: "Request", *args: "Any", **kwargs: "Any") -> "FormData":
+        request_formdata = await _original_form(self, *args, **kwargs)
+        self.scope.setdefault("state", {})[_SCOPE_STATE_FORMDATA_REQUEST_BODY_KEY] = (
+            request_formdata
+        )
+        return request_formdata
 
-                sentry_scope = sentry_sdk.get_isolation_scope()
-                extractor = StarletteRequestExtractor(request)
-                info = await extractor.extract_request_info()
+    Request.form = sentry_form
 
-                def _make_request_event_processor(
-                    req: "Any", integration: "Any"
-                ) -> "Callable[[Event, dict[str, Any]], Event]":
-                    def event_processor(
-                        event: "Event", hint: "Dict[str, Any]"
-                    ) -> "Event":
-                        # Add info from request to event
-                        request_info = event.get("request", {})
-                        if info:
-                            if "cookies" in info:
-                                request_info["cookies"] = info["cookies"]
-                            if "data" in info:
-                                request_info["data"] = info["data"]
-                        event["request"] = deepcopy(request_info)
 
-                        return event
+def _serialize_cached_request_body_attribute(
+    client: "sentry_sdk.client.BaseClient", request: "Request"
+) -> "Optional[str]":
+    """
+    Returns a stringified JSON representation of the request body if the request body is cached on the ASGI scope and within size bounds.
+    """
+    scope_state = request.scope.get("state", {})
+    if (
+        "content-length" not in request.headers
+        or _SCOPE_STATE_JSON_REQUEST_BODY_KEY not in scope_state
+        and _SCOPE_STATE_FORMDATA_REQUEST_BODY_KEY not in scope_state
+    ):
+        return None
 
-                    return event_processor
+    try:
+        content_length = int(request.headers["content-length"])
+    except ValueError:
+        return None
 
-                sentry_scope._name = StarletteIntegration.identifier
-                sentry_scope.add_event_processor(
-                    _make_request_event_processor(request, integration)
+    if content_length and not request_body_within_bounds(client, content_length):
+        return OVER_SIZE_LIMIT_SUBSTITUTE
+
+    if _SCOPE_STATE_JSON_REQUEST_BODY_KEY in request.scope["state"]:
+        return json.dumps(request.scope["state"][_SCOPE_STATE_JSON_REQUEST_BODY_KEY])
+
+    form = request.scope["state"][_SCOPE_STATE_FORMDATA_REQUEST_BODY_KEY]
+
+    form_data = {}
+    for key, val in form.items():
+        is_file = isinstance(val, UploadFile)
+        form_data[key] = val if not is_file else "[Unparsable]"
+
+    return json.dumps(form_data)
+
+
+async def _wrap_async_handler(
+    handler: "Callable[..., Awaitable[Any]]", *args: "Any", **kwargs: "Any"
+) -> "Any":
+    """
+    Wraps an asynchronous handler function to attach request info to the server segment span.
+    The request body cached on the ASGI scope is attached to streamed spans, but consuming the request body in the event
+    processor can still cause application hangs.
+    """
+    client = sentry_sdk.get_client()
+    integration = client.get_integration(StarletteIntegration)
+    if integration is None:
+        return await handler(*args, **kwargs)
+
+    request = args[0]
+
+    _set_transaction_name_and_source(
+        sentry_sdk.get_current_scope(),
+        integration.transaction_style,
+        request,
+    )
+
+    sentry_scope = sentry_sdk.get_isolation_scope()
+    extractor = StarletteRequestExtractor(request)
+
+    info = await extractor.extract_request_info()
+
+    def _make_request_event_processor(
+        req: "Any", integration: "Any"
+    ) -> "Callable[[Event, dict[str, Any]], Event]":
+        def event_processor(event: "Event", hint: "Dict[str, Any]") -> "Event":
+            # Add info from request to event
+            request_info = event.get("request", {})
+            if info:
+                if "cookies" in info:
+                    request_info["cookies"] = info["cookies"]
+                if "data" in info:
+                    request_info["data"] = info["data"]
+            event["request"] = deepcopy(request_info)
+
+            return event
+
+        return event_processor
+
+    sentry_scope._name = StarletteIntegration.identifier
+    sentry_scope.add_event_processor(
+        _make_request_event_processor(request, integration)
+    )
+
+    try:
+        return await handler(*args, **kwargs)
+    finally:
+        current_span = _get_current_streamed_span()
+
+        if type(current_span) is StreamedSpan:
+            serialized_request_body = _serialize_cached_request_body_attribute(
+                client=client, request=request
+            )
+            if serialized_request_body:
+                current_span._segment.set_attribute(
+                    SPANDATA.HTTP_REQUEST_BODY_DATA,
+                    serialized_request_body,
                 )
 
-                if has_span_streaming_enabled(client.options):
-                    _set_request_body_data_on_streaming_segment(info)
 
-                return await old_func(*args, **kwargs)
+def patch_request_response() -> None:
+    old_request_response = starlette.routing.request_response
+
+    def _sentry_request_response(func: "Callable[[Any], Any]") -> "ASGIApp":
+        old_func = func
+
+        is_coroutine = _is_async_callable(old_func)
+        if is_coroutine:
+
+            async def _sentry_async_func(*args: "Any", **kwargs: "Any") -> "Any":
+                return await _wrap_async_handler(old_func, *args, **kwargs)
 
             func = _sentry_async_func