35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
|
"""Middleware to propagate or generate request_id for every request."""
|
||
|
|
|
||
|
|
import uuid
|
||
|
|
|
||
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
||
|
|
from starlette.requests import Request
|
||
|
|
from starlette.responses import Response
|
||
|
|
|
||
|
|
from app.core.logging_config import request_id_ctx
|
||
|
|
|
||
|
|
REQUEST_ID_HEADER = "X-Request-ID"
|
||
|
|
|
||
|
|
|
||
|
|
class RequestIDMiddleware(BaseHTTPMiddleware):
|
||
|
|
"""Extract X-Request-ID from incoming request headers or generate one.
|
||
|
|
|
||
|
|
The request_id is stored in a ContextVar so that all log records emitted
|
||
|
|
during the request are automatically annotated with it, without needing to
|
||
|
|
pass it explicitly through every call.
|
||
|
|
|
||
|
|
The same request_id is also echoed back in the response header so that
|
||
|
|
callers can correlate logs.
|
||
|
|
"""
|
||
|
|
|
||
|
|
async def dispatch(self, request: Request, call_next) -> Response:
|
||
|
|
request_id = request.headers.get(REQUEST_ID_HEADER) or str(uuid.uuid4())
|
||
|
|
token = request_id_ctx.set(request_id)
|
||
|
|
try:
|
||
|
|
response = await call_next(request)
|
||
|
|
finally:
|
||
|
|
request_id_ctx.reset(token)
|
||
|
|
|
||
|
|
response.headers[REQUEST_ID_HEADER] = request_id
|
||
|
|
return response
|