|
@@ -1,9 +1,8 @@
|
|
|
# -*- coding: utf-8 -*-
|
|
# -*- coding: utf-8 -*-
|
|
|
|
|
|
|
|
-import json
|
|
|
|
|
from starlette.responses import HTMLResponse
|
|
from starlette.responses import HTMLResponse
|
|
|
-from typing import Any, AsyncGenerator, Callable
|
|
|
|
|
-from fastapi import FastAPI, Request, Response, WebSocket
|
|
|
|
|
|
|
+from typing import Any, AsyncGenerator
|
|
|
|
|
+from fastapi import Depends, FastAPI, Request, Response
|
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.staticfiles import StaticFiles
|
|
|
from fastapi.concurrency import asynccontextmanager
|
|
from fastapi.concurrency import asynccontextmanager
|
|
|
from fastapi.openapi.docs import (
|
|
from fastapi.openapi.docs import (
|
|
@@ -12,10 +11,8 @@ from fastapi.openapi.docs import (
|
|
|
get_swagger_ui_oauth2_redirect_html
|
|
get_swagger_ui_oauth2_redirect_html
|
|
|
)
|
|
)
|
|
|
from fastapi_limiter import FastAPILimiter
|
|
from fastapi_limiter import FastAPILimiter
|
|
|
|
|
+from fastapi_limiter.depends import RateLimiter
|
|
|
from math import ceil
|
|
from math import ceil
|
|
|
-from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
|
|
|
|
-from starlette.types import ASGIApp
|
|
|
|
|
-import time
|
|
|
|
|
|
|
|
|
|
from app.config.setting import settings
|
|
from app.config.setting import settings
|
|
|
from app.core.logger import log
|
|
from app.core.logger import log
|
|
@@ -29,138 +26,40 @@ from app.api.v1.module_system.params.service import ParamsService
|
|
|
from app.api.v1.module_system.dict.service import DictDataService
|
|
from app.api.v1.module_system.dict.service import DictDataService
|
|
|
|
|
|
|
|
|
|
|
|
|
-# ✅ 核心修复:手动实现Redis限流计数(适配0.1.6版本,确保限流生效)
|
|
|
|
|
-class CustomLimiterMiddleware(BaseHTTPMiddleware):
|
|
|
|
|
- """
|
|
|
|
|
- 手动实现限流逻辑(直接操作Redis,替代RateLimiter的隐式调用)
|
|
|
|
|
- 核心:按IP+路径生成唯一key,Redis计数,超过阈值触发限流
|
|
|
|
|
- """
|
|
|
|
|
-
|
|
|
|
|
- def __init__(
|
|
|
|
|
- self,
|
|
|
|
|
- app: ASGIApp,
|
|
|
|
|
- times: int = 5, # 限流次数
|
|
|
|
|
- seconds: int = 10, # 限流时间窗口
|
|
|
|
|
- prefix: str = settings.REQUEST_LIMITER_REDIS_PREFIX or "fastapi-limiter:" # Redis key前缀
|
|
|
|
|
- ):
|
|
|
|
|
- super().__init__(app)
|
|
|
|
|
- self.times = times
|
|
|
|
|
- self.seconds = seconds
|
|
|
|
|
- self.prefix = prefix
|
|
|
|
|
-
|
|
|
|
|
- # 生成唯一限流key(IP + 请求路径)
|
|
|
|
|
- def _get_limit_key(self, request: Request) -> str:
|
|
|
|
|
- client_ip = request.client.host or "unknown"
|
|
|
|
|
- path = request.url.path
|
|
|
|
|
- return f"{self.prefix}:{client_ip}:{path}"
|
|
|
|
|
-
|
|
|
|
|
- # 限流触发回调
|
|
|
|
|
- async def _limit_callback(self, expire: int):
|
|
|
|
|
- """返回指定格式的429响应"""
|
|
|
|
|
- expires = ceil(expire / 30) # 动态计算Retry-After值(你示例中的222是占位,实际是expires)
|
|
|
|
|
- # 构造严格匹配的响应体
|
|
|
|
|
- response_body = {
|
|
|
|
|
- "code": -1,
|
|
|
|
|
- "msg": "请求过于频繁,请稍后重试",
|
|
|
|
|
- "data": {
|
|
|
|
|
- "Retry-After": str(expires) # 动态值,替换示例中的222
|
|
|
|
|
- },
|
|
|
|
|
- "status_code": 429,
|
|
|
|
|
- "success": False
|
|
|
|
|
- }
|
|
|
|
|
- # 返回Response对象
|
|
|
|
|
- return Response(
|
|
|
|
|
- content=json.dumps(response_body, ensure_ascii=False), # 确保中文正常显示
|
|
|
|
|
- status_code=429, # HTTP状态码
|
|
|
|
|
- headers={"Retry-After": str(expires)}, # 响应头也保留(可选)
|
|
|
|
|
- media_type="application/json" # 声明JSON格式
|
|
|
|
|
- )
|
|
|
|
|
-
|
|
|
|
|
- async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
|
|
|
|
|
- # 1. WebSocket请求 → 跳过限流
|
|
|
|
|
- if request.scope.get("type") == "websocket":
|
|
|
|
|
- return await call_next(request)
|
|
|
|
|
-
|
|
|
|
|
- # 2. HTTP请求 → 执行手动限流逻辑
|
|
|
|
|
- try:
|
|
|
|
|
- # 获取Redis连接(fastapi-limiter已初始化)
|
|
|
|
|
- redis_client = FastAPILimiter.redis
|
|
|
|
|
- if not redis_client:
|
|
|
|
|
- log.warning("Redis未初始化,跳过限流")
|
|
|
|
|
- return await call_next(request)
|
|
|
|
|
-
|
|
|
|
|
- # 生成限流key
|
|
|
|
|
- limit_key = self._get_limit_key(request)
|
|
|
|
|
- # 当前时间戳
|
|
|
|
|
- now = int(time.time())
|
|
|
|
|
- # 时间窗口起始(now - seconds)
|
|
|
|
|
- window_start = now - self.seconds
|
|
|
|
|
-
|
|
|
|
|
- # ✅ 核心Redis操作(原子计数,避免并发问题)
|
|
|
|
|
- async with redis_client.pipeline(transaction=True) as pipe:
|
|
|
|
|
- # 1. 删除时间窗口外的旧计数
|
|
|
|
|
- await pipe.zremrangebyscore(limit_key, 0, window_start)
|
|
|
|
|
- # 2. 添加当前请求时间戳到有序集合
|
|
|
|
|
- await pipe.zadd(limit_key, {now: now})
|
|
|
|
|
- # 3. 设置key过期时间(避免内存泄漏)
|
|
|
|
|
- await pipe.expire(limit_key, self.seconds * 2)
|
|
|
|
|
- # 4. 获取当前窗口内的请求数
|
|
|
|
|
- await pipe.zcard(limit_key)
|
|
|
|
|
- # 执行管道
|
|
|
|
|
- results = await pipe.execute()
|
|
|
|
|
- # 提取请求数(第四个操作的结果)
|
|
|
|
|
- request_count = results[3]
|
|
|
|
|
-
|
|
|
|
|
- # ✅ 判断是否超过限流阈值
|
|
|
|
|
- if request_count > self.times:
|
|
|
|
|
- # 获取过期时间,调用回调返回响应
|
|
|
|
|
- ttl = await redis_client.ttl(limit_key)
|
|
|
|
|
- # 关键:返回响应,不是 raise
|
|
|
|
|
- return await self._limit_callback(ttl or self.seconds)
|
|
|
|
|
- except CustomException:
|
|
|
|
|
- # 限流触发 → 抛异常(全局处理器捕获)
|
|
|
|
|
- raise
|
|
|
|
|
- except Exception as e:
|
|
|
|
|
- # 非限流异常 → 日志记录,放行请求(避免阻断业务)
|
|
|
|
|
- log.error(f"限流中间件执行异常: {str(e)}")
|
|
|
|
|
-
|
|
|
|
|
- # 3. 执行后续处理
|
|
|
|
|
- response = await call_next(request)
|
|
|
|
|
- return response
|
|
|
|
|
-
|
|
|
|
|
-
|
|
|
|
|
-# ✅ 生命周期函数(仅初始化,无中间件操作)
|
|
|
|
|
@asynccontextmanager
|
|
@asynccontextmanager
|
|
|
async def lifespan(app: FastAPI) -> AsyncGenerator[Any, Any]:
|
|
async def lifespan(app: FastAPI) -> AsyncGenerator[Any, Any]:
|
|
|
|
|
+ """
|
|
|
|
|
+ 自定义 FastAPI 应用生命周期。
|
|
|
|
|
+
|
|
|
|
|
+ 参数:
|
|
|
|
|
+ - app (FastAPI): FastAPI 应用实例。
|
|
|
|
|
+
|
|
|
|
|
+ 返回:
|
|
|
|
|
+ - AsyncGenerator[Any, Any]: 生命周期上下文生成器。
|
|
|
|
|
+ """
|
|
|
try:
|
|
try:
|
|
|
- # 数据库初始化
|
|
|
|
|
await InitializeData().init_db()
|
|
await InitializeData().init_db()
|
|
|
log.info(f"✅ {settings.DATABASE_TYPE}数据库初始化完成")
|
|
log.info(f"✅ {settings.DATABASE_TYPE}数据库初始化完成")
|
|
|
-
|
|
|
|
|
- # 全局事件加载
|
|
|
|
|
await import_modules_async(modules=settings.EVENT_LIST, desc="全局事件", app=app, status=True)
|
|
await import_modules_async(modules=settings.EVENT_LIST, desc="全局事件", app=app, status=True)
|
|
|
log.info("✅ 全局事件模块加载完成")
|
|
log.info("✅ 全局事件模块加载完成")
|
|
|
-
|
|
|
|
|
- # Redis配置/字典初始化
|
|
|
|
|
await ParamsService().init_config_service(redis=app.state.redis)
|
|
await ParamsService().init_config_service(redis=app.state.redis)
|
|
|
log.info("✅ Redis系统配置初始化完成")
|
|
log.info("✅ Redis系统配置初始化完成")
|
|
|
await DictDataService().init_dict_service(redis=app.state.redis)
|
|
await DictDataService().init_dict_service(redis=app.state.redis)
|
|
|
log.info("✅ Redis数据字典初始化完成")
|
|
log.info("✅ Redis数据字典初始化完成")
|
|
|
-
|
|
|
|
|
- # 定时任务初始化
|
|
|
|
|
await SchedulerUtil.init_system_scheduler()
|
|
await SchedulerUtil.init_system_scheduler()
|
|
|
scheduler_jobs_count = len(SchedulerUtil.get_all_jobs())
|
|
scheduler_jobs_count = len(SchedulerUtil.get_all_jobs())
|
|
|
scheduler_status = SchedulerUtil.get_job_status()
|
|
scheduler_status = SchedulerUtil.get_job_status()
|
|
|
log.info(f"✅ 定时任务调度器初始化完成 ({scheduler_jobs_count} 个任务)")
|
|
log.info(f"✅ 定时任务调度器初始化完成 ({scheduler_jobs_count} 个任务)")
|
|
|
|
|
|
|
|
- # ✅ 初始化fastapi-limiter(仅获取Redis连接)
|
|
|
|
|
|
|
+ # 6. 初始化请求限制器
|
|
|
await FastAPILimiter.init(
|
|
await FastAPILimiter.init(
|
|
|
redis=app.state.redis,
|
|
redis=app.state.redis,
|
|
|
prefix=settings.REQUEST_LIMITER_REDIS_PREFIX,
|
|
prefix=settings.REQUEST_LIMITER_REDIS_PREFIX,
|
|
|
|
|
+ http_callback=http_limit_callback,
|
|
|
)
|
|
)
|
|
|
log.info("✅ 请求限制器初始化完成")
|
|
log.info("✅ 请求限制器初始化完成")
|
|
|
-
|
|
|
|
|
- # 启动信息面板
|
|
|
|
|
|
|
+
|
|
|
|
|
+ # 导入并显示最终的启动信息面板
|
|
|
from app.utils.console import run as console_run
|
|
from app.utils.console import run as console_run
|
|
|
from app.common.enums import EnvironmentEnum
|
|
from app.common.enums import EnvironmentEnum
|
|
|
console_run(
|
|
console_run(
|
|
@@ -171,14 +70,13 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[Any, Any]:
|
|
|
scheduler_jobs=scheduler_jobs_count,
|
|
scheduler_jobs=scheduler_jobs_count,
|
|
|
scheduler_status=scheduler_status,
|
|
scheduler_status=scheduler_status,
|
|
|
)
|
|
)
|
|
|
-
|
|
|
|
|
|
|
+
|
|
|
except Exception as e:
|
|
except Exception as e:
|
|
|
log.error(f"❌ 应用初始化失败: {str(e)}")
|
|
log.error(f"❌ 应用初始化失败: {str(e)}")
|
|
|
raise
|
|
raise
|
|
|
|
|
|
|
|
yield
|
|
yield
|
|
|
-
|
|
|
|
|
- # 关闭逻辑
|
|
|
|
|
|
|
+
|
|
|
try:
|
|
try:
|
|
|
await import_modules_async(modules=settings.EVENT_LIST, desc="全局事件", app=app, status=False)
|
|
await import_modules_async(modules=settings.EVENT_LIST, desc="全局事件", app=app, status=False)
|
|
|
log.info("✅ 全局事件模块卸载完成")
|
|
log.info("✅ 全局事件模块卸载完成")
|
|
@@ -186,45 +84,78 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[Any, Any]:
|
|
|
log.info("✅ 定时任务调度器已关闭")
|
|
log.info("✅ 定时任务调度器已关闭")
|
|
|
await FastAPILimiter.close()
|
|
await FastAPILimiter.close()
|
|
|
log.info("✅ 请求限制器已关闭")
|
|
log.info("✅ 请求限制器已关闭")
|
|
|
|
|
+
|
|
|
except Exception as e:
|
|
except Exception as e:
|
|
|
log.error(f"❌ 应用关闭过程中发生错误: {str(e)}")
|
|
log.error(f"❌ 应用关闭过程中发生错误: {str(e)}")
|
|
|
|
|
+
|
|
|
|
|
|
|
|
-
|
|
|
|
|
-# ✅ 中间件注册(添加自定义限流中间件)
|
|
|
|
|
def register_middlewares(app: FastAPI) -> None:
|
|
def register_middlewares(app: FastAPI) -> None:
|
|
|
- # 1. 原有中间件
|
|
|
|
|
|
|
+ """
|
|
|
|
|
+ 注册全局中间件。
|
|
|
|
|
+
|
|
|
|
|
+ 参数:
|
|
|
|
|
+ - app (FastAPI): FastAPI 应用实例。
|
|
|
|
|
+
|
|
|
|
|
+ 返回:
|
|
|
|
|
+ - None
|
|
|
|
|
+ """
|
|
|
for middleware in settings.MIDDLEWARE_LIST[::-1]:
|
|
for middleware in settings.MIDDLEWARE_LIST[::-1]:
|
|
|
if not middleware:
|
|
if not middleware:
|
|
|
continue
|
|
continue
|
|
|
middleware = import_module(middleware, desc="中间件")
|
|
middleware = import_module(middleware, desc="中间件")
|
|
|
app.add_middleware(middleware)
|
|
app.add_middleware(middleware)
|
|
|
|
|
|
|
|
- # 2. 限流中间件(核心:5次/10秒)
|
|
|
|
|
- app.add_middleware(
|
|
|
|
|
- CustomLimiterMiddleware,
|
|
|
|
|
- times=5,
|
|
|
|
|
- seconds=10,
|
|
|
|
|
- prefix=settings.REQUEST_LIMITER_REDIS_PREFIX or "fastapi-limiter:"
|
|
|
|
|
- )
|
|
|
|
|
- log.info("✅ 限流中间件注册完成")
|
|
|
|
|
|
|
+def register_exceptions(app: FastAPI) -> None:
|
|
|
|
|
+ """
|
|
|
|
|
+ 统一注册异常处理器。
|
|
|
|
|
|
|
|
|
|
+ 参数:
|
|
|
|
|
+ - app (FastAPI): FastAPI 应用实例。
|
|
|
|
|
|
|
|
-# ✅ 其他函数(不变)
|
|
|
|
|
-def register_exceptions(app: FastAPI) -> None:
|
|
|
|
|
|
|
+ 返回:
|
|
|
|
|
+ - None
|
|
|
|
|
+ """
|
|
|
handle_exception(app)
|
|
handle_exception(app)
|
|
|
|
|
|
|
|
-
|
|
|
|
|
def register_routers(app: FastAPI) -> None:
|
|
def register_routers(app: FastAPI) -> None:
|
|
|
- app.include_router(router=router)
|
|
|
|
|
|
|
+ """
|
|
|
|
|
+ 注册根路由。
|
|
|
|
|
+
|
|
|
|
|
+ 参数:
|
|
|
|
|
+ - app (FastAPI): FastAPI 应用实例。
|
|
|
|
|
|
|
|
|
|
+ 返回:
|
|
|
|
|
+ - None
|
|
|
|
|
+ """
|
|
|
|
|
+ app.include_router(router=router, dependencies=[Depends(RateLimiter(times=5, seconds=10))])
|
|
|
|
|
|
|
|
def register_files(app: FastAPI) -> None:
|
|
def register_files(app: FastAPI) -> None:
|
|
|
|
|
+ """
|
|
|
|
|
+ 注册静态资源挂载和文件相关配置。
|
|
|
|
|
+
|
|
|
|
|
+ 参数:
|
|
|
|
|
+ - app (FastAPI): FastAPI 应用实例。
|
|
|
|
|
+
|
|
|
|
|
+ 返回:
|
|
|
|
|
+ - None
|
|
|
|
|
+ """
|
|
|
|
|
+ # 挂载静态文件目录
|
|
|
if settings.STATIC_ENABLE:
|
|
if settings.STATIC_ENABLE:
|
|
|
|
|
+ # 确保日志目录存在
|
|
|
settings.STATIC_ROOT.mkdir(parents=True, exist_ok=True)
|
|
settings.STATIC_ROOT.mkdir(parents=True, exist_ok=True)
|
|
|
app.mount(path=settings.STATIC_URL, app=StaticFiles(directory=settings.STATIC_ROOT), name=settings.STATIC_DIR)
|
|
app.mount(path=settings.STATIC_URL, app=StaticFiles(directory=settings.STATIC_ROOT), name=settings.STATIC_DIR)
|
|
|
|
|
|
|
|
-
|
|
|
|
|
def reset_api_docs(app: FastAPI) -> None:
|
|
def reset_api_docs(app: FastAPI) -> None:
|
|
|
|
|
+ """
|
|
|
|
|
+ 使用本地静态资源自定义 API 文档页面(Swagger UI 与 ReDoc)。
|
|
|
|
|
+
|
|
|
|
|
+ 参数:
|
|
|
|
|
+ - app (FastAPI): FastAPI 应用实例。
|
|
|
|
|
+
|
|
|
|
|
+ 返回:
|
|
|
|
|
+ - None
|
|
|
|
|
+ """
|
|
|
|
|
+
|
|
|
@app.get(settings.DOCS_URL, include_in_schema=False)
|
|
@app.get(settings.DOCS_URL, include_in_schema=False)
|
|
|
async def custom_swagger_ui_html() -> HTMLResponse:
|
|
async def custom_swagger_ui_html() -> HTMLResponse:
|
|
|
return get_swagger_ui_html(
|
|
return get_swagger_ui_html(
|
|
@@ -247,4 +178,20 @@ def reset_api_docs(app: FastAPI) -> None:
|
|
|
title=app.title + " - ReDoc",
|
|
title=app.title + " - ReDoc",
|
|
|
redoc_js_url=settings.REDOC_JS_URL,
|
|
redoc_js_url=settings.REDOC_JS_URL,
|
|
|
redoc_favicon_url=settings.FAVICON_URL,
|
|
redoc_favicon_url=settings.FAVICON_URL,
|
|
|
- )
|
|
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+async def http_limit_callback(request: Request, response: Response, expire: int):
|
|
|
|
|
+ """
|
|
|
|
|
+ 请求限制时的默认回调函数
|
|
|
|
|
+
|
|
|
|
|
+ :param request: FastAPI 请求对象
|
|
|
|
|
+ :param response: FastAPI 响应对象
|
|
|
|
|
+ :param expire: 剩余毫秒数
|
|
|
|
|
+ :return:
|
|
|
|
|
+ """
|
|
|
|
|
+ expires = ceil(expire / 30)
|
|
|
|
|
+ raise CustomException(
|
|
|
|
|
+ status_code=429,
|
|
|
|
|
+ msg='请求过于频繁,请稍后重试',
|
|
|
|
|
+ data={'Retry-After': str(expires)},
|
|
|
|
|
+ )
|