添加接口调用限制

Context/common.py

@@ -10,6 +10,7 @@ from Utils.SqlAlchemyUtils import SqlalchemyConnect
 from configparser import ConfigParser
 from Utils.VerifyCodeUtils import EmailCodeVerify, ImageCaptchaVerify, PhoneVerifyCode
 from Utils.PhoneMsgUtils import PhoneMsgSend
+from Utils.ApiLimitUtils import ApiLimit
 
 # 日志
 uvicorn_log = getLogger('uvicorn')
@@ -58,3 +59,7 @@ phone_verify_code = PhoneVerifyCode(redis_pool, phone_msg_send)
 tianyancha_api = TianyanchaApi(conf['tianyancha']['token'], mongo_connect=mg_db)
 # 预警通接口
 yujingtong_api = YujingtongApi(conf['yujingtong']['token'], mongo_connect=mg_db)
+
+# 接口使用统计与限制
+api_limit = ApiLimit(redis_pool=redis_pool)
+

Mods/BasicInformation/Router.py

@@ -8,7 +8,7 @@ from .Mods.CompanyIndexMain import Router as CompanyIndexMainRouter
 from .Mods.CompanyShareholder import Router as CompanyShareholderRouter
 from .Mods.CompanyEquityPenetration import Router as CompanyEquityPenetrationRouter
 from .Mods.CompanySearchMain import Router as CompanySearchMainRouter
-from Context.common import auth_util, common_db, tianyancha_api
+from Context.common import auth_util, common_db,  api_limit
 from . import Crud, Schemas
 
 router = APIRouter(
@@ -34,6 +34,7 @@ def company_all_info_query(req: Schemas.CompanyAllInfoQueryReq, db: Session = De
 @router.post("/auto_create_company", summary="自动新建城投信息", tags=["城投详细信息"],
              response_model=Schemas.AutoCreateCompanyRes)
 def auto_create_company(req: Schemas.AutoCreateCompanyReq, db: Session = Depends(common_db.get_db)):
+    api_limit.limit_by_time_count_check("auto_create_company", 10, 30)
     Crud.auto_create_company(db, req)
     return "创建成功"
 

Mods/OtherInformation/Router.py

@@ -14,7 +14,6 @@ from .Mods.CompanyOtherLease import Router as CompanyOtherLeaseRouter
 from .Mods.CompanyRentalLease import Router as CompanyRentalLeaseRouter
 from .Mods.CompanyRestrictedAssets import Router as CompanyRestrictedAssetsRouter
 from .Mods.CompanyTrustLease import Router as CompanyTrustLeaseRouter
-
 from Context.common import auth_util
 
 router = APIRouter(prefix="/other_Information",

Utils/ApiLimitUtils.py

@@ -0,0 +1,58 @@
+from fastapi import HTTPException
+
+from Utils.RedisUtils import RedisPool
+
+
+class ApiLimit:
+    """
+    重要接口使用统计与限制
+    """
+    def __init__(self, redis_pool: RedisPool):
+        self.redis_pool = redis_pool
+
+    def limit_by_time_count_check(self, api_key, count, time_s):
+        """
+        限制接口在一段时间内的调用次数
+        @return:
+        """
+        api_id = f"limit_by_time_count_check_{api_key}"
+        last_count = self.redis_pool.conn.get(api_id)
+        if not last_count:
+            self.redis_pool.conn.set(api_id, 1)
+            self.redis_pool.conn.expire(api_id, time_s)
+        else:
+            last_count = int(last_count)
+            if last_count > count:
+                raise HTTPException(status_code=303, detail=f"{api_key}接口调用次数已达该时段上限")
+            else:
+                self.redis_pool.conn.set(last_count)
+
+    def limit_by_time_ip_count_check(self, api_key, count, time_s, ip):
+        """
+        限制某一ip在一段时间内接口的调用次数
+        @return:
+        """
+        api_id = f"limit_by_time_ip_count_check_{api_key}_{ip}"
+        last_count = self.redis_pool.conn.get(api_id)
+        if not last_count:
+            self.redis_pool.conn.set(api_id, 1)
+            self.redis_pool.conn.expire(api_id, time_s)
+        else:
+            last_count = int(last_count)
+            if last_count > count:
+                raise HTTPException(status_code=303, detail=f"{api_key}接口{ip}调用次数已达该时段上限")
+            else:
+                self.redis_pool.conn.set(last_count)
+
+    def limit_by_time_check(self, api_key, time_s: int):
+        """
+        限制某一接口在一段时间内接口的调用次数
+        @return:
+        """
+        api_id = f"limit_by_time_check_{api_key}"
+        last_count = self.redis_pool.conn.get(api_id)
+        if not last_count:
+            self.redis_pool.conn.set(api_id, 1)
+            self.redis_pool.conn.expire(api_id, time_s)
+        else:
+            raise HTTPException(status_code=303, detail=f"{api_key}接口调用过快")