fix: 增加一个请求限制功能

main.py

@@ -1,6 +1,6 @@
 import logging.config as logging_config
 import os
-
+from contextlib import asynccontextmanager
 import fastapi_plugins
 from fastapi import FastAPI, Request
 from fastapi.exceptions import HTTPException, RequestValidationError
@@ -11,17 +11,17 @@ from config import init_config
 # from src.middleware.flask import flask_app
 from src.utils.exception import (http_exception_handler,
                                  request_validation_error_handler)
+# 请求限制
+import redis.asyncio as aio_redis
+from fastapi_limiter import FastAPILimiter
 
+
+    
 def create_app():
-    app = FastAPI()
-
     mysql_config, redis_config = init_config()
-    # 添加sqlalchemy数据库中间件
-    # once the middleware is applied, any route can then access the database session from the global ``db``
-    app.add_middleware(DBSessionMiddleware, db_url=mysql_config.sqlalchemy_db_uri)
 
-    @app.on_event("startup")
-    async def startup_event():
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
         # 创建日志文件夹和临时文件上传文件夹
         if not os.path.exists("files"):
             os.mkdir("files")
@@ -35,15 +35,21 @@ def create_app():
         print(redis_config)
         await fastapi_plugins.redis_plugin.init_app(app, redis_config)
         await fastapi_plugins.redis_plugin.init()
-
-    @app.on_event("shutdown")
-    async def shutdown_event():
+        # 请求限制
+        await FastAPILimiter.init(redis=aio_redis.from_url("redis://localhost:6379", encoding="utf8"))
+        yield
+        # 应用关闭时，关闭redis连接
         await fastapi_plugins.redis_plugin.terminate()
-
+        await FastAPILimiter.close()
+    app = FastAPI(lifespan=lifespan)
+    # 添加sqlalchemy数据库中间件
+    # once the middleware is applied, any route can then access the database session from the global ``db``
+    app.add_middleware(DBSessionMiddleware, db_url=mysql_config.sqlalchemy_db_uri)
     # 添加异常处理
     app.add_exception_handler(HTTPException, http_exception_handler)
     app.add_exception_handler(RequestValidationError, request_validation_error_handler)
 
+
     # 可以在这里挂载Flask的应用，复用之前项目的相关代码 
     # app.mount("/v1", WSGIMiddleware(flask_app))
 

poetry.lock

@@ -419,6 +419,26 @@ type = "legacy"
 url = "https://pypi.tuna.tsinghua.edu.cn/simple"
 reference = "mirrors"
 
+[[package]]
+name = "fastapi-limiter"
+version = "0.1.6"
+description = "A request rate limiter for fastapi"
+optional = false
+python-versions = ">=3.9,<4.0"
+files = [
+    {file = "fastapi_limiter-0.1.6-py3-none-any.whl", hash = "sha256:2e53179a4208b8f2c8795e38bb001324d3dc37d2800ff49fd28ec5caabf7a240"},
+    {file = "fastapi_limiter-0.1.6.tar.gz", hash = "sha256:6f5fde8efebe12eb33861bdffb91009f699369a3c2862cdc7c1d9acf912ff443"},
+]
+
+[package.dependencies]
+fastapi = "*"
+redis = ">=4.2.0rc1"
+
+[package.source]
+type = "legacy"
+url = "https://pypi.tuna.tsinghua.edu.cn/simple"
+reference = "mirrors"
+
 [[package]]
 name = "fastapi-plugins"
 version = "0.13.0"
@@ -2116,4 +2136,4 @@ reference = "mirrors"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "39a4bc6981b275faba68cb572e406a60300d2d868e646fbbf62ad251e2288e37"
+content-hash = "8da45bacc7264c87e92c82edfe2add96cc11746277726691bf55280061402b3f"

pyproject.toml

@@ -19,6 +19,7 @@ pytz = "^2024.1"
 qiniu = "^7.13.2"
 pillow = "^10.4.0"
 captcha = "^0.6.0"
+fastapi-limiter = "^0.1.6"
 
 
 [[tool.poetry.source]]

src/api/example.py

@@ -7,6 +7,8 @@ from fastapi.security import OAuth2PasswordBearer
 from fastapi_plugins import depends_redis
 from fastapi_sqlalchemy import db
 from sqlalchemy.sql import text
+from fastapi_limiter.depends import RateLimiter
+
 from pydantic import BaseModel
 
 from src.service.example import get_user_by_id
@@ -15,8 +17,8 @@ from src.dtos.example import UserExampleListPagesResult
 
 router = APIRouter()
 
-
-@router.get("/")
+# 并发数限制 每5秒最多100次请求
+@router.get("/", dependencies=[Depends(RateLimiter(times=100, seconds=5))])
 def index():
     return {"msg": "This is Index Page"}