From 42e8606a6a2efe9ab3e4715511754d41fa60fbe6 Mon Sep 17 00:00:00 2001 From: Andrew Weiland Date: Thu, 23 Jul 2026 08:00:13 -0400 Subject: [PATCH 1/2] Adding error to datadog span (#2080) --- src/infrastructure/dependency/structured_logs.py | 6 ++++-- src/infrastructure/http/exceptions.py | 14 ++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/infrastructure/dependency/structured_logs.py b/src/infrastructure/dependency/structured_logs.py index d3ecbf77b6b..bb2c39b81ba 100644 --- a/src/infrastructure/dependency/structured_logs.py +++ b/src/infrastructure/dependency/structured_logs.py @@ -187,9 +187,11 @@ async def dispatch(self, request: Request, call_next) -> Response: try: response = await call_next(request) - except Exception: + except Exception as exc: structlog.stdlib.get_logger("api.error").exception("Unhandled exception") - + span = tracer.current_span() + if span: + span.set_exc_info(type(exc), exc, exc.__traceback__) finally: access_logger = structlog.stdlib.get_logger("api.access") process_time = time.perf_counter_ns() - start_time diff --git a/src/infrastructure/http/exceptions.py b/src/infrastructure/http/exceptions.py index 91bfa23283e..f33071fccda 100644 --- a/src/infrastructure/http/exceptions.py +++ b/src/infrastructure/http/exceptions.py @@ -1,4 +1,5 @@ from asyncpg import InvalidPasswordError +from ddtrace import tracer from fastapi.encoders import jsonable_encoder from fastapi.exception_handlers import http_exception_handler from fastapi.exceptions import RequestValidationError @@ -14,10 +15,17 @@ from infrastructure.logger import logger +def _set_trace_exception(exc: Exception) -> None: + span = tracer.current_span() + if span: + span.set_exc_info(type(exc), exc, exc.__traceback__) + + def custom_base_errors_handler(_: Request, error: BaseError) -> JSONResponse: """This function is called if the BaseError was raised.""" logger.error(error.error, exc_info=error) + _set_trace_exception(error) response = ErrorResponseMulti( result=[ @@ -77,6 +85,7 @@ def python_base_error_handler(_: Request, error: Exception) -> JSONResponse: response = ErrorResponseMulti(result=[ErrorResponse(message=f"Unhandled error: {error_message}")]) logger.error(error_message, exc_info=error) + _set_trace_exception(error) return JSONResponse( content=jsonable_encoder(response.model_dump(by_alias=True)), @@ -90,6 +99,9 @@ def pydantic_validation_errors_handler(request: Request, error: RequestValidatio this_logger = logger.bind( error_location={"file": error.endpoint_file, "line": error.endpoint_line, "function": error.endpoint_function} ) + + _set_trace_exception(error) + for err in error.errors(): if isinstance(err, dict): message = err["msg"] @@ -118,6 +130,8 @@ def sqlalchemy_database_error_handler( ) -> JSONResponse: """This function is called if the SQLAlchemy database error was raised.""" logger.error(str(error), exc_info=error) + _set_trace_exception(error) + response = ErrorResponseMulti(result=[ErrorResponse(message="Internal server error")]) return JSONResponse( From 2a415f3d6e43ed728b26404abf77aeb63772ffa8 Mon Sep 17 00:00:00 2001 From: sricharan varanasi <59170910+sricharan-varanasi@users.noreply.github.com> Date: Thu, 23 Jul 2026 18:07:11 -0400 Subject: [PATCH 2/2] fix: handle WebSocket connections in audit log exception handlers (#2106) --- src/apps/audit/fields.py | 19 ++++----- src/apps/authentication/tests/test_auth.py | 47 ++++++++++++++++++++++ src/infrastructure/http/exceptions.py | 12 ++++-- 3 files changed, 65 insertions(+), 13 deletions(-) diff --git a/src/apps/audit/fields.py b/src/apps/audit/fields.py index f6a02258f04..526c8c03a36 100644 --- a/src/apps/audit/fields.py +++ b/src/apps/audit/fields.py @@ -1,26 +1,27 @@ from asgi_correlation_id.context import correlation_id from ddtrace.trace import tracer -from fastapi import Request from fastapi.routing import APIRoute from starlette.exceptions import HTTPException as StarletteHTTPException +from starlette.requests import HTTPConnection, Request from apps.shared.exception import BaseError from .enums import EventOutcome -def http_audit_fields(request: Request, error: BaseError | StarletteHTTPException | None = None) -> dict: - """Audit fields derived from HTTP request/error.""" - route = request.scope.get("route") +def http_audit_fields(conn: HTTPConnection, error: BaseError | StarletteHTTPException | None = None) -> dict: + """Audit fields derived from an HTTP or WebSocket connection and error.""" + route = conn.scope.get("route") span = tracer.current_span() fields = { - "client_ip": request.client and request.client.host, + "client_ip": conn.client and conn.client.host, "http_request_id": correlation_id.get(), - "http_request_method": request.method, + # WebSocket connections have no HTTP method. + "http_request_method": conn.method if isinstance(conn, Request) else None, "http_response_status_code": isinstance(route, APIRoute) and route.status_code or 200, - "url_path": request.url.path, - "url_query": request.url.query or None, - "user_agent": request.headers.get("user-agent"), + "url_path": conn.url.path, + "url_query": conn.url.query or None, + "user_agent": conn.headers.get("user-agent"), "trace_id": span and str(span.trace_id), } if error is not None: diff --git a/src/apps/authentication/tests/test_auth.py b/src/apps/authentication/tests/test_auth.py index d99395e685c..857fd2b3514 100644 --- a/src/apps/authentication/tests/test_auth.py +++ b/src/apps/authentication/tests/test_auth.py @@ -7,6 +7,8 @@ import pytest from pytest_mock import MockerFixture from sqlalchemy.ext.asyncio import AsyncSession +from starlette.types import Message +from starlette.websockets import WebSocket from apps.audit import EventAction, EventOutcome from apps.authentication.domain.login import UserLoginRequest @@ -25,6 +27,7 @@ from apps.users.cruds.user import UsersCRUD from apps.users.domain import User, UserCreate, UserCreateRequest from config import settings +from infrastructure.http.exceptions import session_token_invalid_error_handler TEST_PASSWORD = "Test12345!" @@ -113,6 +116,50 @@ async def test_user_not_found(self, client: TestClient, mocker: MockerFixture): assert resp.status_code == http.HTTPStatus.UNAUTHORIZED assert resp.json()["result"][0]["message"] == SessionTokenInvalidError.message + async def test_ws_session_invalid_audit_event(self, mocker: MockerFixture): + """A `SessionTokenInvalidError` raised from a WebSocket connection (e.g. `/ws/alerts`) must + + log the `user:session:invalid` audit event without crashing. `http_audit_fields` used to + read the HTTP-only `request.method`, which a `WebSocket` lacks. + """ + audit_log = mocker.patch("infrastructure.http.exceptions.log") + user_id = uuid.uuid4() + + async def receive() -> Message: + return {"type": "websocket.connect"} + + async def send(message: Message) -> None: + return None + + websocket = WebSocket( + { + "type": "websocket", + "scheme": "ws", + "server": ("test.com", 80), + "path": "/ws/alerts", + "query_string": b"", + "root_path": "", + "headers": [(b"host", b"test.com"), (b"user-agent", b"pytest-ws-client")], + "client": ("10.1.2.3", 54321), + }, + receive=receive, + send=send, + ) + error = SessionTokenInvalidError(user_id=user_id) + + resp = await session_token_invalid_error_handler(websocket, error) + + audit_log.assert_awaited_once() + event = audit_log.call_args[0][0] + assert event.event_action == EventAction.USER_SESSION_INVALID + assert event.event_outcome == EventOutcome.FAILURE + assert event.user_id == user_id + assert event.http_request_method is None # WebSocket has no HTTP method + assert event.url_path == "/ws/alerts" + assert event.client_ip == "10.1.2.3" + assert event.user_agent == "pytest-ws-client" + assert resp.status_code == http.HTTPStatus.UNAUTHORIZED + async def test_delete_access_token(self, client: TestClient, user: User, mocker: MockerFixture): audit_log = mocker.patch("apps.authentication.api.auth.log") client.login(user) diff --git a/src/infrastructure/http/exceptions.py b/src/infrastructure/http/exceptions.py index f33071fccda..44034cdf519 100644 --- a/src/infrastructure/http/exceptions.py +++ b/src/infrastructure/http/exceptions.py @@ -5,7 +5,7 @@ from fastapi.exceptions import RequestValidationError from starlette import status from starlette.exceptions import HTTPException as StarletteHTTPException -from starlette.requests import Request +from starlette.requests import HTTPConnection, Request from starlette.responses import JSONResponse, Response from apps.audit import AuditEvent, EventAction, http_audit_fields, log @@ -21,7 +21,7 @@ def _set_trace_exception(exc: Exception) -> None: span.set_exc_info(type(exc), exc, exc.__traceback__) -def custom_base_errors_handler(_: Request, error: BaseError) -> JSONResponse: +def custom_base_errors_handler(_: HTTPConnection, error: BaseError) -> JSONResponse: """This function is called if the BaseError was raised.""" logger.error(error.error, exc_info=error) @@ -53,8 +53,12 @@ def custom_base_errors_handler(_: Request, error: BaseError) -> JSONResponse: ) -async def session_token_invalid_error_handler(request: Request, error: SessionTokenInvalidError) -> JSONResponse: - """user:session:invalid audit event on 401 from invalid session token in `Authorization` header.""" +async def session_token_invalid_error_handler(request: HTTPConnection, error: SessionTokenInvalidError) -> JSONResponse: + """user:session:invalid audit event on 401 from invalid session token. + + ``request`` may be a ``WebSocket`` (e.g. the `/ws/alerts` handshake) as well as an HTTP + ``Request``, so only fields common to both connection types may be used here. + """ await log( AuditEvent( event_action=EventAction.USER_SESSION_INVALID,