diff --git a/src/elexam_core/middleware.py b/src/elexam_core/middleware.py new file mode 100644 index 0000000..b004a18 --- /dev/null +++ b/src/elexam_core/middleware.py @@ -0,0 +1,24 @@ +from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint +from starlette.requests import Request +from starlette.responses import Response +from ulid import ULID + +from elexam_core.context import trace_id_context + + +TRACE_HEADER: str = "X-Trace-Id" + + +class TraceMiddleware(BaseHTTPMiddleware): + async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: + header_trace_id = request.headers.get(TRACE_HEADER, default=None) + trace_id = header_trace_id or str(ULID()) + token = trace_id_context.set(trace_id) + + try: + response = await call_next(request) + finally: + trace_id_context.reset(token) + + response.headers[TRACE_HEADER] = trace_id + return response \ No newline at end of file