32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
"""每请求 trace_id：contextvar 存链路号，中间件生成 / 透传并回写响应头。"""
|
||
import contextvars
|
||
import uuid
|
||
|
||
from starlette.middleware.base import BaseHTTPMiddleware
|
||
from starlette.requests import Request
|
||
from starlette.responses import Response
|
||
|
||
_request_id: contextvars.ContextVar[str | None] = contextvars.ContextVar("request_id", default=None)
|
||
|
||
|
||
def get_request_id() -> str | None:
|
||
"""获取当前请求的 trace_id（请求上下文外返回 None）。"""
|
||
return _request_id.get()
|
||
|
||
|
||
def new_request_id() -> str:
|
||
return uuid.uuid4().hex[:16]
|
||
|
||
|
||
class RequestIdMiddleware(BaseHTTPMiddleware):
|
||
"""透传客户端 X-Trace-Id（无则生成），供日志/响应头/审计链路关联。"""
|
||
|
||
async def dispatch(self, request: Request, call_next):
|
||
rid = request.headers.get("X-Trace-Id") or new_request_id()
|
||
token = _request_id.set(rid)
|
||
try:
|
||
response: Response = await call_next(request)
|
||
response.headers["X-Trace-Id"] = rid
|
||
return response
|
||
finally:
|
||
_request_id.reset(token) |