Set response status code in transaction "response" context. (#2312)

Make sure that the HTTP response status code be set in the transactions "response" context.

This works in WSGI (was already calling set_http_status.) Also added this to ASGI projects.

Fixes #2289

Author: antonpirker

sentry_sdk/integrations/asgi.py

@@ -132,20 +132,24 @@ def _run_asgi2(self, scope):
         # type: (Any) -> Any
         async def inner(receive, send):
             # type: (Any, Any) -> Any
-            return await self._run_app(scope, lambda: self.app(scope)(receive, send))
+            return await self._run_app(scope, receive, send, asgi_version=2)
 
         return inner
 
     async def _run_asgi3(self, scope, receive, send):
         # type: (Any, Any, Any) -> Any
-        return await self._run_app(scope, lambda: self.app(scope, receive, send))
+        return await self._run_app(scope, receive, send, asgi_version=3)
 
-    async def _run_app(self, scope, callback):
-        # type: (Any, Any) -> Any
+    async def _run_app(self, scope, receive, send, asgi_version):
+        # type: (Any, Any, Any, Any, int) -> Any
         is_recursive_asgi_middleware = _asgi_middleware_applied.get(False)
         if is_recursive_asgi_middleware:
             try:
-                return await callback()
+                if asgi_version == 2:
+                    return await self.app(scope)(receive, send)
+                else:
+                    return await self.app(scope, receive, send)
+
             except Exception as exc:
                 _capture_exception(Hub.current, exc, mechanism_type=self.mechanism_type)
                 raise exc from None
@@ -178,11 +182,28 @@ async def _run_app(self, scope, callback):
                     with hub.start_transaction(
                         transaction, custom_sampling_context={"asgi_scope": scope}
                     ):
-                        # XXX: Would be cool to have correct span status, but we
-                        # would have to wrap send(). That is a bit hard to do with
-                        # the current abstraction over ASGI 2/3.
                         try:
-                            return await callback()
+
+                            async def _sentry_wrapped_send(event):
+                                # type: (Dict[str, Any]) -> Any
+                                is_http_response = (
+                                    event.get("type") == "http.response.start"
+                                    and transaction is not None
+                                    and "status" in event
+                                )
+                                if is_http_response:
+                                    transaction.set_http_status(event["status"])
+
+                                return await send(event)
+
+                            if asgi_version == 2:
+                                return await self.app(scope)(
+                                    receive, _sentry_wrapped_send
+                                )
+                            else:
+                                return await self.app(
+                                    scope, receive, _sentry_wrapped_send
+                                )
                         except Exception as exc:
                             _capture_exception(
                                 hub, exc, mechanism_type=self.mechanism_type

sentry_sdk/tracing.py

@@ -663,6 +663,11 @@ def set_context(self, key, value):
         # type: (str, Any) -> None
         self._contexts[key] = value
 
+    def set_http_status(self, http_status):
+        # type: (int) -> None
+        super(Transaction, self).set_http_status(http_status)
+        self.set_context("response", {"status_code": http_status})
+
     def to_json(self):
         # type: () -> Dict[str, Any]
         rv = super(Transaction, self).to_json()

tests/integrations/asgi/test_asgi.py

@@ -48,8 +48,11 @@ async def app(scope, receive, send):
 
 @pytest.fixture
 def asgi3_app_with_error():
+    async def send_with_error(event):
+        1 / 0
+
     async def app(scope, receive, send):
-        await send(
+        await send_with_error(
             {
                 "type": "http.response.start",
                 "status": 200,
@@ -58,10 +61,7 @@ async def app(scope, receive, send):
                 ],
             }
         )
-
-        1 / 0
-
-        await send(
+        await send_with_error(
             {
                 "type": "http.response.body",
                 "body": b"Hello, world!",
@@ -167,9 +167,9 @@ async def test_capture_transaction_with_error(
     sentry_init(send_default_pii=True, traces_sample_rate=1.0)
     app = SentryAsgiMiddleware(asgi3_app_with_error)
 
+    events = capture_events()
     with pytest.raises(ZeroDivisionError):
         async with TestClient(app) as client:
-            events = capture_events()
             await client.get("/")
 
     (error_event, transaction_event) = events
@@ -395,38 +395,35 @@ async def test_auto_session_tracking_with_aggregates(
         (
             "/message",
             "endpoint",
-            "tests.integrations.asgi.test_asgi.asgi3_app_with_error.<locals>.app",
+            "tests.integrations.asgi.test_asgi.asgi3_app.<locals>.app",
             "component",
         ),
     ],
 )
 @pytest.mark.asyncio
 async def test_transaction_style(
     sentry_init,
-    asgi3_app_with_error,
+    asgi3_app,
     capture_events,
     url,
     transaction_style,
     expected_transaction,
     expected_source,
 ):
     sentry_init(send_default_pii=True, traces_sample_rate=1.0)
-    app = SentryAsgiMiddleware(
-        asgi3_app_with_error, transaction_style=transaction_style
-    )
+    app = SentryAsgiMiddleware(asgi3_app, transaction_style=transaction_style)
 
     scope = {
-        "endpoint": asgi3_app_with_error,
+        "endpoint": asgi3_app,
         "route": url,
         "client": ("127.0.0.1", 60457),
     }
 
-    with pytest.raises(ZeroDivisionError):
-        async with TestClient(app, scope=scope) as client:
-            events = capture_events()
-            await client.get(url)
+    async with TestClient(app, scope=scope) as client:
+        events = capture_events()
+        await client.get(url)
 
-    (_, transaction_event) = events
+    (transaction_event,) = events
 
     assert transaction_event["transaction"] == expected_transaction
     assert transaction_event["transaction_info"] == {"source": expected_source}

tests/integrations/fastapi/test_fastapi.py

@@ -22,6 +22,12 @@
 def fastapi_app_factory():
     app = FastAPI()
 
+    @app.get("/error")
+    async def _error():
+        capture_message("Hi")
+        1 / 0
+        return {"message": "Hi"}
+
     @app.get("/message")
     async def _message():
         capture_message("Hi")
@@ -218,3 +224,101 @@ async def _error(request: Request):
     event = events[0]
     assert event["request"]["data"] == {"password": "[Filtered]"}
     assert event["request"]["headers"]["authorization"] == "[Filtered]"
+
+
+@pytest.mark.asyncio
+def test_response_status_code_ok_in_transaction_context(sentry_init, capture_envelopes):
+    """
+    Tests that the response status code is added to the transaction "response" context.
+    """
+    sentry_init(
+        integrations=[StarletteIntegration(), FastApiIntegration()],
+        traces_sample_rate=1.0,
+        release="demo-release",
+    )
+
+    envelopes = capture_envelopes()
+
+    app = fastapi_app_factory()
+
+    client = TestClient(app)
+    client.get("/message")
+
+    (_, transaction_envelope) = envelopes
+    transaction = transaction_envelope.get_transaction_event()
+
+    assert transaction["type"] == "transaction"
+    assert len(transaction["contexts"]) > 0
+    assert (
+        "response" in transaction["contexts"].keys()
+    ), "Response context not found in transaction"
+    assert transaction["contexts"]["response"]["status_code"] == 200
+
+
+@pytest.mark.asyncio
+def test_response_status_code_error_in_transaction_context(
+    sentry_init,
+    capture_envelopes,
+):
+    """
+    Tests that the response status code is added to the transaction "response" context.
+    """
+    sentry_init(
+        integrations=[StarletteIntegration(), FastApiIntegration()],
+        traces_sample_rate=1.0,
+        release="demo-release",
+    )
+
+    envelopes = capture_envelopes()
+
+    app = fastapi_app_factory()
+
+    client = TestClient(app)
+    with pytest.raises(ZeroDivisionError):
+        client.get("/error")
+
+    (
+        _,
+        _,
+        transaction_envelope,
+    ) = envelopes
+    transaction = transaction_envelope.get_transaction_event()
+
+    assert transaction["type"] == "transaction"
+    assert len(transaction["contexts"]) > 0
+    assert (
+        "response" in transaction["contexts"].keys()
+    ), "Response context not found in transaction"
+    assert transaction["contexts"]["response"]["status_code"] == 500
+
+
+@pytest.mark.asyncio
+def test_response_status_code_not_found_in_transaction_context(
+    sentry_init,
+    capture_envelopes,
+):
+    """
+    Tests that the response status code is added to the transaction "response" context.
+    """
+    sentry_init(
+        integrations=[StarletteIntegration(), FastApiIntegration()],
+        traces_sample_rate=1.0,
+        release="demo-release",
+    )
+
+    envelopes = capture_envelopes()
+
+    app = fastapi_app_factory()
+
+    client = TestClient(app)
+    client.get("/non-existing-route-123")
+
+    (transaction_envelope,) = envelopes
+    transaction = transaction_envelope.get_transaction_event()
+
+    assert transaction["type"] == "transaction"
+    assert len(transaction["contexts"]) > 0
+    assert (
+        "response" in transaction["contexts"].keys()
+    ), "Response context not found in transaction"
+    assert transaction["contexts"]["response"]["status_code"] == 404

tests/integrations/flask/test_flask.py

@@ -912,3 +912,61 @@ def error():
     assert (
         event["contexts"]["replay"]["replay_id"] == "12312012123120121231201212312012"
     )
+
+
+def test_response_status_code_ok_in_transaction_context(
+    sentry_init, capture_envelopes, app
+):
+    """
+    Tests that the response status code is added to the transaction context.
+    This also works for when there is an Exception during the request, but somehow the test flask app doesn't seem to trigger that.
+    """
+    sentry_init(
+        integrations=[flask_sentry.FlaskIntegration()],
+        traces_sample_rate=1.0,
+        release="demo-release",
+    )
+
+    envelopes = capture_envelopes()
+
+    client = app.test_client()
+    client.get("/message")
+
+    Hub.current.client.flush()
+
+    (_, transaction_envelope, _) = envelopes
+    transaction = transaction_envelope.get_transaction_event()
+
+    assert transaction["type"] == "transaction"
+    assert len(transaction["contexts"]) > 0
+    assert (
+        "response" in transaction["contexts"].keys()
+    ), "Response context not found in transaction"
+    assert transaction["contexts"]["response"]["status_code"] == 200
+
+
+def test_response_status_code_not_found_in_transaction_context(
+    sentry_init, capture_envelopes, app
+):
+    sentry_init(
+        integrations=[flask_sentry.FlaskIntegration()],
+        traces_sample_rate=1.0,
+        release="demo-release",
+    )
+
+    envelopes = capture_envelopes()
+
+    client = app.test_client()
+    client.get("/not-existing-route")
+
+    Hub.current.client.flush()
+
+    (transaction_envelope, _) = envelopes
+    transaction = transaction_envelope.get_transaction_event()
+
+    assert transaction["type"] == "transaction"
+    assert len(transaction["contexts"]) > 0
+    assert (
+        "response" in transaction["contexts"].keys()
+    ), "Response context not found in transaction"
+    assert transaction["contexts"]["response"]["status_code"] == 404

tests/integrations/starlette/test_starlette.py

@@ -700,9 +700,7 @@ def test_middleware_callback_spans(sentry_init, capture_events):
         },
         {
             "op": "middleware.starlette.send",
-            "description": "_ASGIAdapter.send.<locals>.send"
-            if STARLETTE_VERSION < (0, 21)
-            else "_TestClientTransport.handle_request.<locals>.send",
+            "description": "SentryAsgiMiddleware._run_app.<locals>._sentry_wrapped_send",
             "tags": {"starlette.middleware_name": "ServerErrorMiddleware"},
         },
         {
@@ -717,9 +715,7 @@ def test_middleware_callback_spans(sentry_init, capture_events):
         },
         {
             "op": "middleware.starlette.send",
-            "description": "_ASGIAdapter.send.<locals>.send"
-            if STARLETTE_VERSION < (0, 21)
-            else "_TestClientTransport.handle_request.<locals>.send",
+            "description": "SentryAsgiMiddleware._run_app.<locals>._sentry_wrapped_send",
             "tags": {"starlette.middleware_name": "ServerErrorMiddleware"},
         },
     ]
@@ -793,9 +789,7 @@ def test_middleware_partial_receive_send(sentry_init, capture_events):
         },
         {
             "op": "middleware.starlette.send",
-            "description": "_ASGIAdapter.send.<locals>.send"
-            if STARLETTE_VERSION < (0, 21)
-            else "_TestClientTransport.handle_request.<locals>.send",
+            "description": "SentryAsgiMiddleware._run_app.<locals>._sentry_wrapped_send",
             "tags": {"starlette.middleware_name": "ServerErrorMiddleware"},
         },
         {

tests/integrations/starlite/test_starlite.py

@@ -221,12 +221,12 @@ def test_middleware_callback_spans(sentry_init, capture_events):
         },
         {
             "op": "middleware.starlite.send",
-            "description": "TestClientTransport.create_send.<locals>.send",
+            "description": "SentryAsgiMiddleware._run_app.<locals>._sentry_wrapped_send",
             "tags": {"starlite.middleware_name": "SampleMiddleware"},
         },
         {
             "op": "middleware.starlite.send",
-            "description": "TestClientTransport.create_send.<locals>.send",
+            "description": "SentryAsgiMiddleware._run_app.<locals>._sentry_wrapped_send",
             "tags": {"starlite.middleware_name": "SampleMiddleware"},
         },
     ]
@@ -286,12 +286,11 @@ def test_middleware_partial_receive_send(sentry_init, capture_events):
         },
         {
             "op": "middleware.starlite.send",
-            "description": "TestClientTransport.create_send.<locals>.send",
+            "description": "SentryAsgiMiddleware._run_app.<locals>._sentry_wrapped_send",
             "tags": {"starlite.middleware_name": "SamplePartialReceiveSendMiddleware"},
         },
     ]
 
-    print(transaction_event["spans"])
     idx = 0
     for span in transaction_event["spans"]:
         assert span["op"] == expected[idx]["op"]