53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
"""HTTP middleware: bind a correlation id and log each request."""
|
|
|
|
import time
|
|
import uuid
|
|
|
|
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
|
|
|
from app.core.logging import correlation_id, get_logger
|
|
|
|
log = get_logger("http")
|
|
|
|
_HEADER = "x-correlation-id"
|
|
|
|
|
|
class CorrelationIdMiddleware:
|
|
"""Pure-ASGI middleware: reuse inbound ``X-Correlation-Id`` or mint one,
|
|
bind it for downstream logs, echo it back, and log request completion."""
|
|
|
|
def __init__(self, app: ASGIApp) -> None:
|
|
self.app = app
|
|
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
headers = dict(scope["headers"])
|
|
inbound = headers.get(_HEADER.encode())
|
|
cid = inbound.decode() if inbound else uuid.uuid4().hex
|
|
token = correlation_id.set(cid)
|
|
started = time.perf_counter()
|
|
status_code = 0
|
|
|
|
async def send_wrapper(message: Message) -> None:
|
|
nonlocal status_code
|
|
if message["type"] == "http.response.start":
|
|
status_code = message["status"]
|
|
message.setdefault("headers", [])
|
|
message["headers"].append((_HEADER.encode(), cid.encode()))
|
|
await send(message)
|
|
|
|
try:
|
|
await self.app(scope, receive, send_wrapper)
|
|
finally:
|
|
log.info(
|
|
"request",
|
|
method=scope["method"],
|
|
path=scope["path"],
|
|
status=status_code,
|
|
duration_ms=round((time.perf_counter() - started) * 1000, 1),
|
|
)
|
|
correlation_id.reset(token)
|