use cached attributes directly

Author: alexander-alderman-webb

sentry_sdk/integrations/starlette.py

@@ -48,7 +48,6 @@
     from starlette import __version__ as STARLETTE_VERSION
     from starlette.applications import Starlette  # type: ignore
     from starlette.datastructures import (  # type: ignore
-        FormData,
         UploadFile,
     )
     from starlette.middleware import Middleware  # type: ignore
@@ -92,9 +91,6 @@
 
 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"
@@ -153,16 +149,6 @@ 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__
@@ -495,61 +481,13 @@ def _is_async_callable(obj: "Any") -> bool:
     )
 
 
-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 causes applications to hang when middleware or handlers consume the raw
-    `receive()` callable exposed by Starlette.
-    """
-    _original_json = Request.json
-
-    @functools.wraps(_original_json)
-    async def wrapped_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
-
-    Request.json = wrapped_json
-
-
-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 causes applications to hang when middleware or handlers consume the raw
-    `receive()` callable exposed by Starlette.
-    """
-    if not hasattr(Request, "_get_form"):
-        return
-
-    _original_form = Request._get_form
-
-    @functools.wraps(_original_form)
-    async def wrapped_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
-
-    Request._get_form = wrapped_form
-
-
-def _serialize_cached_request_body_attribute(
+def _get_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.
+    Returns a stringified JSON representation of the request body if the request body is cached 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
-    ):
+    if "content-length" not in request.headers:
         return None
 
     try:
@@ -560,13 +498,16 @@ def _serialize_cached_request_body_attribute(
     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])
+    json_body = getattr(request, "_json", None)
+    if json_body is not None:
+        return json.dumps(json_body)
 
-    form = request.scope["state"][_SCOPE_STATE_FORMDATA_REQUEST_BODY_KEY]
+    formdata_body = getattr(request, "_form", None)
+    if formdata_body is None:
+        return
 
     form_data = {}
-    for key, val in form.items():
+    for key, val in formdata_body.items():
         is_file = isinstance(val, UploadFile)
         form_data[key] = val if not is_file else "[Unparsable]"
 
@@ -627,13 +568,13 @@ def event_processor(event: "Event", hint: "Dict[str, Any]") -> "Event":
         current_span = _get_current_streamed_span()
 
         if type(current_span) is StreamedSpan:
-            serialized_request_body = _serialize_cached_request_body_attribute(
+            request_body = _get_cached_request_body_attribute(
                 client=client, request=request
             )
-            if serialized_request_body:
+            if request_body:
                 current_span._segment.set_attribute(
                     SPANDATA.HTTP_REQUEST_BODY_DATA,
-                    serialized_request_body,
+                    request_body,
                 )
 
 

tests/integrations/starlette/test_starlette.py

@@ -21,14 +21,8 @@
 from starlette.middleware import Middleware
 from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.middleware.trustedhost import TrustedHostMiddleware
-from starlette.requests import Request
 from starlette.testclient import TestClient
 
-try:
-    from starlette.middleware.exceptions import ExceptionMiddleware
-except ImportError:
-    from starlette.exceptions import ExceptionMiddleware
-
 import sentry_sdk
 from sentry_sdk import capture_message, get_baggage, get_traceparent
 from sentry_sdk.consts import SPANDATA
@@ -279,31 +273,6 @@ async def my_send(*args, **kwargs):
         await self.app(scope, partial_receive, partial_send)
 
 
-@pytest.fixture(autouse=True)
-def reset_starlette_integration(uninstall_integration):
-    original_request_json = Request.json
-    original_request_form = Request.form
-    original_starlette_call = starlette.applications.Starlette.__call__
-    original_request_response = starlette.routing.request_response
-    original_middleware_init = Middleware.__init__
-    original_authentication_call = AuthenticationMiddleware.__call__
-    original_exception_middleware_init = ExceptionMiddleware.__init__
-    original_exception_middleware_call = ExceptionMiddleware.__call__
-
-    yield
-
-    Request.json = original_request_json
-    Request.form = original_request_form
-    starlette.applications.Starlette.__call__ = original_starlette_call
-    starlette.routing.request_response = original_request_response
-    Middleware.__init__ = original_middleware_init
-    AuthenticationMiddleware.__call__ = original_authentication_call
-    ExceptionMiddleware.__init__ = original_exception_middleware_init
-    ExceptionMiddleware.__call__ = original_exception_middleware_call
-
-    uninstall_integration("starlette")
-
-
 @pytest.mark.asyncio
 @pytest.mark.parametrize("span_streaming", [True, False])
 async def test_request_info_json_body(
@@ -375,10 +344,6 @@ async def test_request_info_json_body(
 
 
 @pytest.mark.asyncio
-@pytest.mark.skipif(
-    STARLETTE_VERSION < (0, 24),
-    reason="Patched `_get_form()` was added in version 0.24 with https://github.com/Kludex/starlette/commit/c568b55dff8be94b9c917e186e512ab53d7310e1",
-)
 @pytest.mark.parametrize("span_streaming", [True, False])
 async def test_formdata_request_body(
     sentry_init, capture_events, capture_items, span_streaming
@@ -460,10 +425,6 @@ async def test_formdata_request_body(
 
 
 @pytest.mark.asyncio
-@pytest.mark.skipif(
-    STARLETTE_VERSION < (0, 24),
-    reason="Patched `_get_form()` was added in version 0.24 with https://github.com/Kludex/starlette/commit/c568b55dff8be94b9c917e186e512ab53d7310e1",
-)
 @pytest.mark.parametrize("span_streaming", [True, False])
 async def test_request_body_too_big(
     sentry_init, capture_events, capture_items, span_streaming