63 lines
2.6 KiB
Python
63 lines
2.6 KiB
Python
import json
|
|
from time import perf_counter
|
|
from typing import cast
|
|
from uuid import uuid4
|
|
|
|
from fastapi import Request
|
|
from fastapi.responses import JSONResponse
|
|
from starlette.concurrency import iterate_in_threadpool
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.responses import Response, StreamingResponse
|
|
|
|
from src.logging.logger import log_queue, request_id_ctx
|
|
|
|
|
|
class ProcessingTimeMiddleware(BaseHTTPMiddleware, ):
|
|
async def dispatch(self, request: Request, call_next)->Response:
|
|
start_time = perf_counter()
|
|
response = await call_next(request)
|
|
process_time = perf_counter() - start_time
|
|
response.headers["X-Process-Time"] = str(process_time)
|
|
return response
|
|
|
|
|
|
class LoggingMiddleware(BaseHTTPMiddleware):
|
|
|
|
async def build_log_line(self, request_id, method, path, status_code, detail, client_ip) -> str:
|
|
return f"[{request_id}] [{method}] [{path}] [{status_code}] [{detail}] [{client_ip}]"
|
|
|
|
async def dispatch(self, request: Request, call_next) -> Response:
|
|
request_id=str(uuid4())
|
|
request_id_ctx.set(request_id)
|
|
client_ip = request.headers.get('x-forwarded-for', '').split(',')[0].strip() or (request.client.host if request.client else 'unknown')
|
|
method = request.method
|
|
path=request.url.path
|
|
try:
|
|
response = await call_next(request)
|
|
except Exception as exc: # noqa: BLE001
|
|
line=await self.build_log_line(request_id=request_id, method=method, path=path, status_code=500, detail=repr(exc), client_ip=client_ip)
|
|
log_queue.put_nowait(("endpoints", line))
|
|
return JSONResponse(
|
|
status_code=500,
|
|
content={"detail": "Internal Server Error", "request_id": request_id}
|
|
)
|
|
|
|
streaming_response = cast(StreamingResponse, response)
|
|
|
|
chunks = []
|
|
async for chunk in streaming_response.body_iterator:
|
|
chunks.append(chunk.encode() if isinstance(chunk, str) else bytes(chunk))
|
|
|
|
body_bytes = b"".join(chunks)
|
|
streaming_response.body_iterator = iterate_in_threadpool(iter([body_bytes]))
|
|
|
|
try:
|
|
parsed = json.loads(body_bytes)
|
|
body = parsed.get("detail", None) if not isinstance(parsed, bool) else None
|
|
except (json.JSONDecodeError, TypeError):
|
|
body = None
|
|
|
|
line=await self.build_log_line(request_id=request_id, method=method, path=path, status_code=response.status_code, detail=body, client_ip=client_ip)
|
|
log_queue.put_nowait(("endpoints", line))
|
|
|
|
return response |