Add correlation ID to responses and logs (#673)

* Integrate asgi-correlation-id. * Register CorrelationIdMiddleware last. * Custom header name and shorter token * Make headers match * Revert interactive testing code * Improve comment * Improve comment * Improve clarity of code and comment Co-authored-by: Padraic Shafer <[email protected]> --------- Co-authored-by: Padraic Shafer <[email protected]>
bluesky · Feb 29, 2024 · 95d973b · 95d973b
1 parent 2e160b5
commit 95d973b
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 5 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -47,6 +47,7 @@ all = [
     "aiosqlite",
     "alembic",
     "anyio",
+    "asgi-correlation-id",
     "asyncpg",
     "appdirs",
     "awkward >=2.4.3",
@@ -189,6 +190,7 @@ minimal-server = [
     "alembic",
     "anyio",
     "appdirs",
+    "asgi-correlation-id",
     "cachey",
     "cachetools",
     "click !=8.1.0",
@@ -225,6 +227,7 @@ server = [
     "alembic",
     "anyio",
     "appdirs",
+    "asgi-correlation-id",
     "asyncpg",
     "awkward >=2.4.3",
     "blosc",

diff --git a/tiled/commandline/_serve.py b/tiled/commandline/_serve.py
@@ -98,6 +98,9 @@ def serve_directory(
         ),
     ),
     port: int = typer.Option(8000, help="Bind to a socket with this port."),
+    log_config: Optional[str] = typer.Option(
+        None, help="Custom uvicorn logging configuration file"
+    ),
     object_cache_available_bytes: Optional[float] = typer.Option(
         None,
         "--data-cache",
@@ -243,7 +246,7 @@ async def walk_and_serve():
 
         import uvicorn
 
-        uvicorn.run(web_app, host=host, port=port)
+        uvicorn.run(web_app, host=host, port=port, log_config=log_config)
 
 
 def serve_catalog(

diff --git a/tiled/server/app.py b/tiled/server/app.py
@@ -12,7 +12,9 @@
 import anyio
 import packaging.version
 import yaml
+from asgi_correlation_id import CorrelationIdMiddleware, correlation_id
 from fastapi import APIRouter, FastAPI, HTTPException, Request, Response, Security
+from fastapi.exception_handlers import http_exception_handler
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.openapi.utils import get_openapi
 from fastapi.responses import HTMLResponse, JSONResponse
@@ -298,8 +300,25 @@ async def unsupported_query_type_exception_handler(
         allow_credentials=True,
         allow_methods=["*"],
         allow_headers=["*"],
+        # If we restrict the allowed_headers in future, remember to include
+        # exemptions for these, related to asgi_correlation_id.
+        # allow_headers=["X-Requested-With", "X-Tiled-Request-ID"],
+        expose_headers=["X-Tiled-Request-ID"],
     )
 
+    @app.exception_handler(Exception)
+    async def unhandled_exception_handler(
+        request: Request, exc: Exception
+    ) -> JSONResponse:
+        return await http_exception_handler(
+            request,
+            HTTPException(
+                500,
+                "Internal server error",
+                headers={"X-Tiled-Request-ID": correlation_id.get() or ""},
+            ),
+        )
+
     app.include_router(router, prefix="/api/v1")
 
     # The Tree and Authenticator have the opportunity to add custom routes to
@@ -848,16 +867,27 @@ async def capture_metrics_prometheus(request: Request, call_next):
             try:
                 response = await call_next(request)
             except Exception:
-                # Make a placeholder to feed to capture_request_metrics.
-                # The actual error response will be created by FastAPI/Starlette
-                # further up the call stack, where this exception will propagate.
-                response = Response(status_code=500)
+                # Make an ephemeral response solely for 'capture_request_metrics'.
+                # It will only be used in the 'finally' clean-up block.
+                only_for_metrics = Response(status_code=500)
+                response = only_for_metrics
+                # Now re-raise the exception so that the server can generate and
+                # send an appropriate response to the client.
                 raise
             finally:
                 response.__class__ = PatchedStreamingResponse  # tolerate memoryview
                 metrics.capture_request_metrics(request, response)
+
+            # This is a *real* response (i.e., not the 'only_for_metrics' response).
+            # An exception above would have triggered an early exit.
             return response
 
+    app.add_middleware(
+        CorrelationIdMiddleware,
+        header_name="X-Tiled-Request-ID",
+        generator=lambda: secrets.token_hex(8),
+    )
+
     return app