style: normalize python formatting gates
This commit is contained in:
@@ -1,8 +1,10 @@
|
||||
"""
|
||||
请求日志中间件
|
||||
"""
|
||||
import time
|
||||
|
||||
import logging
|
||||
import time
|
||||
|
||||
from fastapi import Request
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
|
||||
@@ -11,58 +13,57 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
class RequestLoggingMiddleware(BaseHTTPMiddleware):
|
||||
"""请求日志中间件"""
|
||||
|
||||
|
||||
async def dispatch(self, request: Request, call_next):
|
||||
# 记录请求开始时间
|
||||
start_time = time.time()
|
||||
|
||||
|
||||
# 记录请求信息
|
||||
logger.info(f"Request: {request.method} {request.url.path}")
|
||||
|
||||
|
||||
# 处理请求
|
||||
response = await call_next(request)
|
||||
|
||||
|
||||
# 计算处理时间
|
||||
process_time = time.time() - start_time
|
||||
|
||||
|
||||
# 记录响应信息
|
||||
logger.info(
|
||||
f"Response: {request.method} {request.url.path} "
|
||||
f"status={response.status_code} time={process_time:.3f}s"
|
||||
f"Response: {request.method} {request.url.path} " f"status={response.status_code} time={process_time:.3f}s"
|
||||
)
|
||||
|
||||
|
||||
# 添加响应头
|
||||
response.headers["X-Process-Time"] = str(process_time)
|
||||
|
||||
|
||||
return response
|
||||
|
||||
|
||||
class RateLimitMiddleware(BaseHTTPMiddleware):
|
||||
"""简单的速率限制中间件(基于内存)"""
|
||||
|
||||
|
||||
def __init__(self, app, max_requests: int = 100, window_seconds: int = 60):
|
||||
super().__init__(app)
|
||||
self.max_requests = max_requests
|
||||
self.window_seconds = window_seconds
|
||||
self.requests = {} # {ip: [(timestamp, ...)]}
|
||||
|
||||
|
||||
async def dispatch(self, request: Request, call_next):
|
||||
# 获取客户端 IP
|
||||
client_ip = request.client.host
|
||||
current_time = time.time()
|
||||
|
||||
|
||||
# 清理过期记录
|
||||
if client_ip in self.requests:
|
||||
self.requests[client_ip] = [
|
||||
ts for ts in self.requests[client_ip]
|
||||
if current_time - ts < self.window_seconds
|
||||
ts for ts in self.requests[client_ip] if current_time - ts < self.window_seconds
|
||||
]
|
||||
|
||||
|
||||
# 检查速率限制
|
||||
request_count = len(self.requests.get(client_ip, []))
|
||||
|
||||
|
||||
if request_count >= self.max_requests:
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
return JSONResponse(
|
||||
status_code=429,
|
||||
content={
|
||||
@@ -72,19 +73,17 @@ class RateLimitMiddleware(BaseHTTPMiddleware):
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# 记录请求
|
||||
if client_ip not in self.requests:
|
||||
self.requests[client_ip] = []
|
||||
self.requests[client_ip].append(current_time)
|
||||
|
||||
|
||||
# 处理请求
|
||||
response = await call_next(request)
|
||||
|
||||
|
||||
# 添加速率限制信息到响应头
|
||||
response.headers["X-RateLimit-Limit"] = str(self.max_requests)
|
||||
response.headers["X-RateLimit-Remaining"] = str(
|
||||
self.max_requests - len(self.requests[client_ip])
|
||||
)
|
||||
|
||||
response.headers["X-RateLimit-Remaining"] = str(self.max_requests - len(self.requests[client_ip]))
|
||||
|
||||
return response
|
||||
|
||||
Reference in New Issue
Block a user