88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
"""Request correlation and access logging middleware."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import time
|
|
import uuid
|
|
from collections.abc import Awaitable, Callable
|
|
|
|
from fastapi import Request, Response
|
|
from starlette.datastructures import MutableHeaders
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.types import Message
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
RequestHandler = Callable[[Request], Awaitable[Response]]
|
|
|
|
|
|
class RequestContextMiddleware(BaseHTTPMiddleware):
|
|
"""Attach a bounded request ID and log one structured completion event."""
|
|
|
|
def __init__(self, app: object, header_name: str) -> None:
|
|
super().__init__(app) # type: ignore[arg-type]
|
|
self._header_name = header_name
|
|
|
|
async def dispatch(self, request: Request, call_next: RequestHandler) -> Response:
|
|
supplied = request.headers.get(self._header_name, "")
|
|
request_id = supplied if 0 < len(supplied) <= 128 else str(uuid.uuid4())
|
|
request.state.request_id = request_id
|
|
started = time.monotonic()
|
|
response = await call_next(request)
|
|
response.headers[self._header_name] = request_id
|
|
logger.info(
|
|
"request completed",
|
|
extra={
|
|
"request_id": request_id,
|
|
"method": request.method,
|
|
"path": request.url.path,
|
|
"status": response.status_code,
|
|
"duration_ms": round((time.monotonic() - started) * 1000, 3),
|
|
},
|
|
)
|
|
return response
|
|
|
|
|
|
class HeadAsGetMiddleware(BaseHTTPMiddleware):
|
|
"""Serve HEAD for every GET route (deep + stub) with an empty body.
|
|
|
|
FastAPI ``add_api_route(methods=['GET'])`` and some router setups omit HEAD;
|
|
contract matrix probes expect synthetic HEAD on each GET path.
|
|
"""
|
|
|
|
async def dispatch(self, request: Request, call_next: RequestHandler) -> Response:
|
|
if request.method != "HEAD":
|
|
return await call_next(request)
|
|
|
|
# Replay as GET, then strip the body while preserving status/headers.
|
|
request.scope["method"] = "GET"
|
|
response = await call_next(request)
|
|
|
|
body = bytearray()
|
|
async for chunk in response.body_iterator:
|
|
if isinstance(chunk, str):
|
|
body.extend(chunk.encode(response.charset or "utf-8"))
|
|
else:
|
|
body.extend(chunk)
|
|
|
|
headers = MutableHeaders(scope={"type": "http", "headers": []})
|
|
for key, value in response.headers.items():
|
|
if key.lower() in {"content-length", "content-type", "transfer-encoding"}:
|
|
continue
|
|
headers.append(key, value)
|
|
headers["content-length"] = str(len(body))
|
|
if response.media_type:
|
|
headers["content-type"] = response.media_type
|
|
|
|
async def _empty_receive() -> Message:
|
|
return {"type": "http.request", "body": b"", "more_body": False}
|
|
|
|
del _empty_receive
|
|
return Response(
|
|
content=b"",
|
|
status_code=response.status_code,
|
|
headers=headers,
|
|
media_type=response.media_type,
|
|
)
|