add try except and more tests

Author: alexander-alderman-webb

sentry_sdk/integrations/fastapi.py

@@ -11,7 +11,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Dict
+    from typing import Any, Callable, Dict, Optional
     from sentry_sdk._types import Event
 
 try:
@@ -112,33 +112,38 @@ async def _sentry_app(*args, **kwargs):
             sentry_scope = sentry_sdk.get_isolation_scope()
             sentry_scope._name = FastApiIntegration.identifier
 
+            def _make_request_event_processor(info):
+                # type: (Optional[Dict[str, Any]]) -> Callable[[Event, Dict[str, Any]], Event]
+                def event_processor(event, hint):
+                    # type: (Event, Dict[str, Any]) -> Event
+
+                    # Extract information from request
+                    request_info = event.get("request", {})
+                    if info:
+                        if "cookies" in info and should_send_default_pii():
+                            request_info["cookies"] = info["cookies"]
+                        if "data" in info:
+                            request_info["data"] = info["data"]
+                    event["request"] = deepcopy(request_info)
+
+                    return event
+
+                return event_processor
+
             try:
                 response = await old_app(*args, **kwargs)
-            finally:
+            except Exception as exception:
                 extractor = StarletteRequestExtractor(request)
                 info = await extractor.extract_request_info()
 
-                def _make_request_event_processor(req, integration):
-                    # type: (Any, Any) -> Callable[[Event, Dict[str, Any]], Event]
-                    def event_processor(event, hint):
-                        # type: (Event, Dict[str, Any]) -> Event
-
-                        # Extract information from request
-                        request_info = event.get("request", {})
-                        if info:
-                            if "cookies" in info and should_send_default_pii():
-                                request_info["cookies"] = info["cookies"]
-                            if "data" in info:
-                                request_info["data"] = info["data"]
-                        event["request"] = deepcopy(request_info)
+                sentry_scope.add_event_processor(_make_request_event_processor(info))
 
-                        return event
+                raise exception
 
-                    return event_processor
+            extractor = StarletteRequestExtractor(request)
+            info = await extractor.extract_request_info()
 
-                sentry_scope.add_event_processor(
-                    _make_request_event_processor(request, integration)
-                )
+            sentry_scope.add_event_processor(_make_request_event_processor(info))
 
             return response
 

sentry_sdk/integrations/starlette.py

@@ -501,20 +501,15 @@ async def _sentry_async_func(*args, **kwargs):
                 sentry_scope = sentry_sdk.get_isolation_scope()
                 sentry_scope._name = StarletteIntegration.identifier
 
-                response = await old_func(*args, **kwargs)
-
-                extractor = StarletteRequestExtractor(request)
-                info = await extractor.extract_request_info()
-
-                def _make_request_event_processor(req, integration):
-                    # type: (Any, Any) -> Callable[[Event, dict[str, Any]], Event]
+                def _make_request_event_processor(info):
+                    # type: (Optional[Dict[str, Any]]) -> Callable[[Event, Dict[str, Any]], Event]
                     def event_processor(event, hint):
                         # type: (Event, Dict[str, Any]) -> Event
 
-                        # Add info from request to event
+                        # Extract information from request
                         request_info = event.get("request", {})
                         if info:
-                            if "cookies" in info:
+                            if "cookies" in info and should_send_default_pii():
                                 request_info["cookies"] = info["cookies"]
                             if "data" in info:
                                 request_info["data"] = info["data"]
@@ -524,9 +519,22 @@ def event_processor(event, hint):
 
                     return event_processor
 
-                sentry_scope.add_event_processor(
-                    _make_request_event_processor(request, integration)
-                )
+                try:
+                    response = await old_func(*args, **kwargs)
+                except Exception as exception:
+                    extractor = StarletteRequestExtractor(request)
+                    info = await extractor.extract_request_info()
+
+                    sentry_scope.add_event_processor(
+                        _make_request_event_processor(info)
+                    )
+
+                    raise exception
+
+                extractor = StarletteRequestExtractor(request)
+                info = await extractor.extract_request_info()
+
+                sentry_scope.add_event_processor(_make_request_event_processor(info))
 
                 return response
 

tests/integrations/fastapi/test_fastapi.py

@@ -12,7 +12,7 @@
 from fastapi.middleware.trustedhost import TrustedHostMiddleware
 
 import sentry_sdk
-from sentry_sdk import capture_message
+from sentry_sdk import capture_message, capture_exception
 from sentry_sdk.feature_flags import add_feature_flag
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 from sentry_sdk.integrations.fastapi import FastApiIntegration
@@ -75,29 +75,6 @@ async def _thread_ids_async():
     return app
 
 
-def test_stream_available_in_handler(sentry_init):
-    sentry_init(
-        integrations=[StarletteIntegration(), FastApiIntegration()],
-    )
-
-    app = FastAPI()
-
-    @app.post("/consume")
-    async def _consume_stream_body(request):
-        # Avoid cache by constructing new request
-        wrapped_request = Request(request.scope, request.receive)
-
-        assert await asyncio.wait_for(wrapped_request.json(), timeout=1.0) == BODY_JSON
-
-        return {"status": "ok"}
-
-    client = TestClient(app)
-    client.post(
-        "/consume",
-        json=BODY_JSON,
-    )
-
-
 @pytest.mark.asyncio
 async def test_response(sentry_init, capture_events):
     # FastAPI is heavily based on Starlette so we also need
@@ -246,6 +223,67 @@ def test_active_thread_id(sentry_init, capture_envelopes, teardown_profiling, en
         assert str(data["active"]) == trace_context["data"]["thread.id"]
 
 
+@pytest.mark.asyncio
+def test_request_body_not_cached_with_exception(sentry_init, capture_events):
+    sentry_init(
+        integrations=[StarletteIntegration(), FastApiIntegration()],
+    )
+
+    app = FastAPI()
+
+    @app.post("/exception")
+    async def _exception(request: Request):
+        1 / 0
+        return {"error": "Oh no!"}
+
+    events = capture_events()
+
+    client = TestClient(app)
+
+    try:
+        client.post(
+            "/exception",
+            json=BODY_JSON,
+        )
+    except Exception:
+        capture_exception()
+
+    event = events[0]
+    assert event["request"]["data"] == ""
+
+
+def test_request_body_cached_with_exception(sentry_init, capture_events):
+    sentry_init(
+        integrations=[StarletteIntegration(), FastApiIntegration()],
+    )
+
+    app = FastAPI(debug=True)
+    app.user_middleware = []
+    app.middleware_stack = app.build_middleware_stack()
+
+    @app.post("/exception")
+    async def _exception(request: Request):
+        await request.json()
+        1 / 0
+        return {"error": "Oh no!"}
+
+    events = capture_events()
+
+    client = TestClient(app)
+
+    try:
+        client.post(
+            "/exception",
+            json=BODY_JSON,
+        )
+    except Exception:
+        print("capturing")
+        capture_exception()
+
+    event = events[0]
+    assert event["request"]["data"] == BODY_JSON
+
+
 @pytest.mark.asyncio
 async def test_original_request_not_scrubbed(sentry_init, capture_events):
     sentry_init(

tests/integrations/starlette/test_starlette.py

@@ -11,7 +11,7 @@
 
 import pytest
 
-from sentry_sdk import capture_message, get_baggage, get_traceparent
+from sentry_sdk import capture_message, get_baggage, get_traceparent, capture_exception
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 from sentry_sdk.integrations.starlette import (
     StarletteIntegration,
@@ -509,32 +509,6 @@ async def test_starletterequestextractor_extract_request_info_no_pii(sentry_init
     assert request_info["data"] == BODY_JSON
 
 
-def test_stream_available_in_handler(sentry_init):
-    sentry_init(
-        integrations=[StarletteIntegration()],
-    )
-
-    async def _consume_stream_body(request):
-        # Avoid cache by constructing new request
-        wrapped_request = Request(request.scope, request.receive)
-
-        assert await asyncio.wait_for(wrapped_request.json(), timeout=1.0) == BODY_JSON
-
-        return starlette.responses.JSONResponse({"status": "ok"})
-
-    app = starlette.applications.Starlette(
-        routes=[
-            starlette.routing.Route("/consume", _consume_stream_body, methods=["POST"]),
-        ],
-    )
-
-    client = TestClient(app)
-    client.post(
-        "/consume",
-        json=BODY_JSON,
-    )
-
-
 @pytest.mark.parametrize(
     "url,transaction_style,expected_transaction,expected_source",
     [
@@ -1001,6 +975,69 @@ def test_active_thread_id(sentry_init, capture_envelopes, teardown_profiling, en
         assert str(data["active"]) == trace_context["data"]["thread.id"]
 
 
+def test_request_body_not_cached_with_exception(sentry_init, capture_events):
+    sentry_init(
+        integrations=[StarletteIntegration()],
+    )
+
+    events = capture_events()
+
+    async def _exception(request):
+        1 / 0
+        return {"error": "Oh no!"}
+
+    app = starlette.applications.Starlette(
+        routes=[
+            starlette.routing.Route("/exception", _exception, methods=["POST"]),
+        ],
+    )
+
+    client = TestClient(app)
+
+    try:
+        client.post(
+            "/exception",
+            json=BODY_JSON,
+        )
+    except ZeroDivisionError:
+        capture_exception()
+
+    event = events[0]
+    assert event["request"]["data"] == ""
+
+
+def test_request_body_cached_with_exception(sentry_init, capture_events):
+    sentry_init(
+        integrations=[StarletteIntegration()],
+    )
+
+    events = capture_events()
+
+    async def _exception(request):
+        request.json()
+        1 / 0
+        return {"error": "Oh no!"}
+
+    app = starlette.applications.Starlette(
+        routes=[
+            starlette.routing.Route("/exception", _exception, methods=["POST"]),
+        ],
+    )
+
+    client = TestClient(app)
+
+    try:
+        client.post(
+            "/exception",
+            json=BODY_JSON,
+        )
+    except ZeroDivisionError:
+        capture_exception()
+
+    event = events[0]
+    assert event["request"]["data"] == BODY_JSON
+
+
 def test_original_request_not_scrubbed(sentry_init, capture_events):
     sentry_init(
         default_integrations=False,