From 8459e0c4470225707fa8b1f7492935ec2f8735b5 Mon Sep 17 00:00:00 2001
From: chenweijia <chenweijia@leelen.cn>
Date: Wed, 14 Aug 2024 14:20:19 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E5=A2=9E=E5=8A=A0=E4=B8=80=E4=B8=AA?=
 =?UTF-8?q?=E8=AF=B7=E6=B1=82=E9=99=90=E5=88=B6=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.py            | 30 ++++++++++++++++++------------
 poetry.lock        | 22 +++++++++++++++++++++-
 pyproject.toml     |  1 +
 src/api/example.py |  6 ++++--
 4 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/main.py b/main.py
index ee35a2b..d7b353e 100644
--- a/main.py
+++ b/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))
 
diff --git a/poetry.lock b/poetry.lock
index a6d4768..6288676 100644
--- a/poetry.lock
+++ b/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"
diff --git a/pyproject.toml b/pyproject.toml
index b8b6b47..d4b5a9f 100644
--- a/pyproject.toml
+++ b/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]]
diff --git a/src/api/example.py b/src/api/example.py
index bf6115b..69787c2 100644
--- a/src/api/example.py
+++ b/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"}