import asyncio import json import logging from time import gmtime, perf_counter, strftime from typing import cast import aiofiles from fastapi import Request from starlette.concurrency import iterate_in_threadpool from starlette.middleware.base import BaseHTTPMiddleware from starlette.responses import Response, StreamingResponse 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 dispatch(self, request: Request, call_next) -> Response: current_time = strftime("%a, %d %b %Y %H:%M:%S +0000", gmtime()) file_time=strftime("%b_%Y", gmtime()) response = await call_next(request) streaming_response = cast(StreamingResponse, response) chunks = [] async for chunk in streaming_response.body_iterator: if isinstance(chunk, str): chunks.append(chunk.encode()) else: chunks.append(bytes(chunk)) body_bytes = b"".join(chunks) streaming_response.body_iterator = iterate_in_threadpool(iter([body_bytes])) try: body = json.loads(body_bytes) body=body.get("detail", None) except (json.JSONDecodeError, TypeError): body = None client_ip = request.headers.get('x-forwarded-for', '').split(',')[0].strip() or (request.client.host if request.client else 'unknown') async with aiofiles.open(f"./logs/endpoints_log_{file_time}.txt", "a") as file: await file.write(f"[{current_time}] [{response.status_code}] [{body}] [{client_ip}]\n") return response class LoggerDB(logging.Handler): def emit(self, record: logging.LogRecord) -> None: msg = self.format(record) asyncio.create_task(self._write(msg)) async def _write(self, msg: str) -> None: file_time = strftime("%b_%Y", gmtime()) current_time = strftime("%a, %d %b %Y %H:%M:%S +0000", gmtime()) async with aiofiles.open(f"./logs/sql_log_{file_time}.txt", "a") as f: await f.write(f"[{current_time}] {msg}\n")