30 lines
981 B
Python
30 lines
981 B
Python
import contextvars
|
|
from fastapi import Request
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
import ulid
|
|
|
|
# ContextVar to hold request_id globally in the thread/coroutine context
|
|
request_id_var = contextvars.ContextVar("request_id", default="")
|
|
|
|
def get_request_id() -> str:
|
|
return request_id_var.get()
|
|
|
|
class RequestTraceMiddleware(BaseHTTPMiddleware):
|
|
async def dispatch(self, request: Request, call_next):
|
|
# Generate new trace ID
|
|
req_id = str(ulid.ULID())
|
|
|
|
# Set context variable
|
|
token = request_id_var.set(req_id)
|
|
|
|
# Store in state for easy route access
|
|
request.state.request_id = req_id
|
|
|
|
try:
|
|
response = await call_next(request)
|
|
# Append to response headers
|
|
response.headers["X-Request-ID"] = req_id
|
|
return response
|
|
finally:
|
|
# Reset context variable to prevent leakage
|
|
request_id_var.reset(token)
|