| """ |
| Correlation ID middleware. |
| Reads X-Request-ID from incoming headers (or generates one) and makes it |
| available via a ContextVar so it can be injected into logs and error responses |
| without threading issues. |
| """ |
| import uuid |
| from contextvars import ContextVar |
|
|
| from starlette.middleware.base import BaseHTTPMiddleware |
| from starlette.requests import Request |
|
|
| _request_id_ctx: ContextVar[str] = ContextVar("request_id", default="") |
|
|
|
|
| def get_request_id() -> str: |
| return _request_id_ctx.get() |
|
|
|
|
| class CorrelationIdMiddleware(BaseHTTPMiddleware): |
| async def dispatch(self, request: Request, call_next): |
| request_id = request.headers.get("X-Request-ID") or str(uuid.uuid4()) |
| token = _request_id_ctx.set(request_id) |
| try: |
| response = await call_next(request) |
| response.headers["X-Request-ID"] = request_id |
| return response |
| finally: |
| _request_id_ctx.reset(token) |
|
|
