67 lines
2.3 KiB
Python
67 lines
2.3 KiB
Python
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")
|