Skip to content

Commit bb56d30

Browse files
authored
Add request-ID middleware that propagates X-Request-ID via ContextVar (#45)
Closes #44. Adds a `RequestIDMiddleware` (in `src/climatevision/api/middleware.py`) that reads the inbound `X-Request-ID` header, falls back to a fresh UUID4 when absent, stores the value on a `contextvars.ContextVar` (`request_id_var`), and echoes it back on the response. Pairs the middleware with a `RequestIDLogFilter` and wires it into `setup_logging` so every log record emitted during the request lifecycle -- including from `inference/pipeline.py` and helper modules that don't hold the FastAPI `Request` object -- carries `%(request_id)s` in the JSON log format. The middleware is registered last in `create_app` so it sits outermost in the stack, ensuring the ContextVar is set before any other middleware or route handler runs. Tests in `tests/test_request_id_middleware.py` cover the three cases the issue called out: - request without `X-Request-ID` -> response carries a UUID-shaped value - request with explicit `X-Request-ID` -> response echoes it back - log record emitted inside the handler exposes `record.request_id` Plus a leak test asserting the ContextVar resets after the response.
1 parent e3e93a1 commit bb56d30

3 files changed

Lines changed: 186 additions & 3 deletions

File tree

src/climatevision/api/main.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -410,6 +410,14 @@ def create_app() -> FastAPI:
410410
allow_methods=["*"],
411411
allow_headers=["*"],
412412
)
413+
# Register the RequestIDMiddleware LAST so it sits OUTERMOST in the
414+
# middleware stack: it must run before every other middleware so the
415+
# request_id ContextVar is set in time for AuditLogMiddleware's logger
416+
# call (and for any inference-pipeline code further down). Starlette
417+
# wraps middleware in reverse add_middleware order, so the last call
418+
# is the outermost wrapper.
419+
from climatevision.api.middleware import RequestIDMiddleware
420+
app.add_middleware(RequestIDMiddleware)
413421

414422
# Wire OWASP-aligned security controls (rate limiting, payload limits, etc.)
415423
app.add_middleware(SecurityMiddleware)

src/climatevision/api/middleware.py

Lines changed: 69 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import logging
1111
import time
1212
import uuid
13+
from contextvars import ContextVar
1314
from typing import Callable
1415

1516
from fastapi import Request, Response
@@ -18,6 +19,55 @@
1819
logger = logging.getLogger(__name__)
1920

2021

22+
# Context variable that carries the X-Request-ID through the entire request
23+
# lifecycle, including non-FastAPI code (inference pipeline, helper modules)
24+
# that does not have access to the FastAPI `Request` object.
25+
request_id_var: ContextVar[str | None] = ContextVar("request_id", default=None)
26+
27+
28+
class RequestIDMiddleware(BaseHTTPMiddleware):
29+
"""
30+
Propagate ``X-Request-ID`` through the request lifecycle.
31+
32+
Reads the inbound ``X-Request-ID`` header (or generates a fresh UUID4 if
33+
absent), stores it on a :class:`~contextvars.ContextVar` so any code that
34+
runs during the request -- inference pipeline, helper modules, background
35+
tasks scheduled with ``asyncio.create_task`` -- can read the same value,
36+
and echoes it back on the response.
37+
38+
Pair this with :class:`RequestIDLogFilter` so log records emitted during
39+
the request automatically carry ``%(request_id)s``.
40+
"""
41+
42+
async def dispatch(self, request: Request, call_next: Callable) -> Response:
43+
request_id = request.headers.get("X-Request-ID") or str(uuid.uuid4())
44+
token = request_id_var.set(request_id)
45+
# Mirror to ``request.state`` so handlers that already read from
46+
# there (e.g. the existing ``RequestLoggingMiddleware`` /
47+
# ``AuditLogMiddleware``) keep working.
48+
request.state.request_id = request_id
49+
try:
50+
response = await call_next(request)
51+
finally:
52+
request_id_var.reset(token)
53+
response.headers["X-Request-ID"] = request_id
54+
return response
55+
56+
57+
class RequestIDLogFilter(logging.Filter):
58+
"""Inject the current request ID into every log record.
59+
60+
Reads :data:`request_id_var` and exposes it as ``record.request_id`` so
61+
logging formatters can reference ``%(request_id)s``. Records emitted
62+
outside of any request (startup, background workers without context)
63+
receive ``"-"`` so the format string never KeyErrors.
64+
"""
65+
66+
def filter(self, record: logging.LogRecord) -> bool: # noqa: A003
67+
record.request_id = request_id_var.get() or "-"
68+
return True
69+
70+
2171
class RequestLoggingMiddleware(BaseHTTPMiddleware):
2272
"""
2373
Middleware for structured request logging and audit trails.
@@ -135,9 +185,25 @@ async def dispatch(self, request: Request, call_next: Callable) -> Response:
135185

136186

137187
def setup_logging(log_level: str = "INFO") -> None:
138-
"""Configure structured JSON logging for the API."""
188+
"""Configure structured JSON logging for the API.
189+
190+
Installs :class:`RequestIDLogFilter` on the root logger so every log
191+
record emitted during a request carries the current ``request_id``.
192+
"""
139193
logging.basicConfig(
140194
level=getattr(logging, log_level.upper()),
141-
format='{"timestamp":"%(asctime)s","level":"%(levelname)s","message":"%(message)s"}',
142-
datefmt="%Y-%m-%dT%H:%M:%S"
195+
format=(
196+
'{"timestamp":"%(asctime)s","level":"%(levelname)s",'
197+
'"request_id":"%(request_id)s","message":"%(message)s"}'
198+
),
199+
datefmt="%Y-%m-%dT%H:%M:%S",
143200
)
201+
request_id_filter = RequestIDLogFilter()
202+
root = logging.getLogger()
203+
# Attach to the root logger and to any handlers already installed by
204+
# ``logging.basicConfig`` so existing handlers also see ``request_id``.
205+
if not any(isinstance(f, RequestIDLogFilter) for f in root.filters):
206+
root.addFilter(request_id_filter)
207+
for handler in root.handlers:
208+
if not any(isinstance(f, RequestIDLogFilter) for f in handler.filters):
209+
handler.addFilter(request_id_filter)
Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
"""Tests for the X-Request-ID middleware.
2+
3+
Covers issue #44: request log lines from the inference pipeline and helper
4+
modules need to carry the same identifier as the inbound HTTP request.
5+
"""
6+
7+
from __future__ import annotations
8+
9+
import logging
10+
import re
11+
import uuid
12+
13+
from fastapi import FastAPI
14+
from fastapi.testclient import TestClient
15+
16+
from climatevision.api.middleware import (
17+
RequestIDLogFilter,
18+
RequestIDMiddleware,
19+
request_id_var,
20+
)
21+
22+
23+
UUID_RE = re.compile(
24+
r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
25+
)
26+
27+
28+
def _build_app() -> FastAPI:
29+
app = FastAPI()
30+
app.add_middleware(RequestIDMiddleware)
31+
32+
logger = logging.getLogger("climatevision.tests.request_id")
33+
34+
@app.get("/echo")
35+
def echo() -> dict[str, str]:
36+
# Read the contextvar the way inference/pipeline.py would
37+
return {"request_id": request_id_var.get() or ""}
38+
39+
@app.get("/log")
40+
def log_route() -> dict[str, str]:
41+
logger.info("inside-handler", extra={"sample": True})
42+
return {"ok": "1"}
43+
44+
return app
45+
46+
47+
def test_response_has_uuid_request_id_when_header_absent() -> None:
48+
"""Without an X-Request-ID header the middleware mints a fresh UUID4."""
49+
client = TestClient(_build_app())
50+
response = client.get("/echo")
51+
assert response.status_code == 200
52+
53+
request_id = response.headers["X-Request-ID"]
54+
assert UUID_RE.match(request_id), f"not a UUID4 shape: {request_id!r}"
55+
# And the same value flowed into the contextvar that the route saw.
56+
assert response.json()["request_id"] == request_id
57+
58+
59+
def test_response_echoes_explicit_request_id() -> None:
60+
"""When the client sends X-Request-ID the middleware must echo it back."""
61+
client = TestClient(_build_app())
62+
sent = "test-id-deadbeef"
63+
response = client.get("/echo", headers={"X-Request-ID": sent})
64+
assert response.status_code == 200
65+
assert response.headers["X-Request-ID"] == sent
66+
assert response.json()["request_id"] == sent
67+
68+
69+
def test_log_records_include_request_id_via_filter(
70+
caplog: object,
71+
) -> None:
72+
"""Logs emitted inside the request show request_id once the filter is on."""
73+
handler = logging.StreamHandler()
74+
handler.addFilter(RequestIDLogFilter())
75+
handler.setFormatter(
76+
logging.Formatter("%(request_id)s | %(message)s")
77+
)
78+
79+
target_logger = logging.getLogger("climatevision.tests.request_id")
80+
target_logger.setLevel(logging.INFO)
81+
target_logger.addHandler(handler)
82+
try:
83+
# Use pytest's caplog with the same filter so we can introspect records.
84+
with caplog.at_level(logging.INFO, logger="climatevision.tests.request_id"): # type: ignore[attr-defined]
85+
for record_filter in [RequestIDLogFilter()]:
86+
caplog.handler.addFilter(record_filter) # type: ignore[attr-defined]
87+
88+
sent = str(uuid.uuid4())
89+
client = TestClient(_build_app())
90+
response = client.get("/log", headers={"X-Request-ID": sent})
91+
assert response.status_code == 200
92+
93+
records = [
94+
r for r in caplog.records # type: ignore[attr-defined]
95+
if r.name == "climatevision.tests.request_id"
96+
]
97+
assert records, "expected the route handler to emit a log record"
98+
record = records[0]
99+
assert getattr(record, "request_id", None) == sent
100+
finally:
101+
target_logger.removeHandler(handler)
102+
103+
104+
def test_request_id_var_resets_after_response() -> None:
105+
"""The ContextVar must not leak across requests."""
106+
client = TestClient(_build_app())
107+
client.get("/echo", headers={"X-Request-ID": "first"})
108+
# Outside any request the var returns its default ("" via ``or ''``).
109+
assert request_id_var.get() is None

0 commit comments

Comments
 (0)