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

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"]