From b9d91c5cd7291bb9c5fd017df195f2914ae2f710 Mon Sep 17 00:00:00 2001
From: InfinityPacer <160988576+InfinityPacer@users.noreply.github.com>
Date: Sun, 30 Jun 2024 11:48:29 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20DoubanModule=E8=A7=A6=E5=8F=91=E9=99=90?=
 =?UTF-8?q?=E6=B5=81=E6=97=B6=E6=94=AF=E6=8C=81=E7=AB=8B=E5=8D=B3=E6=8A=9B?=
 =?UTF-8?q?=E5=87=BA=E9=99=90=E6=B5=81=E5=BC=82=E5=B8=B8?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/modules/douban/__init__.py | 23 ++++++++++++++++++-----
 app/schemas/exception.py       |  8 ++++++++
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/app/modules/douban/__init__.py b/app/modules/douban/__init__.py
index f63ed112..337ec39f 100644
--- a/app/modules/douban/__init__.py
+++ b/app/modules/douban/__init__.py
@@ -15,6 +15,7 @@ from app.modules.douban.apiv2 import DoubanApi
 from app.modules.douban.douban_cache import DoubanCache
 from app.modules.douban.scraper import DoubanScraper
 from app.schemas import MediaPerson
+from app.schemas.exception import APIRateLimitException
 from app.schemas.types import MediaType
 from app.utils.common import retry
 from app.utils.http import RequestUtils
@@ -147,11 +148,12 @@ class DoubanModule(_ModuleBase):
 
         return None
 
-    def douban_info(self, doubanid: str, mtype: MediaType = None) -> Optional[dict]:
+    def douban_info(self, doubanid: str, mtype: MediaType = None, raise_exception: bool = True) -> Optional[dict]:
         """
         获取豆瓣信息
         :param doubanid: 豆瓣ID
         :param mtype:    媒体类型
+        :param raise_exception: 触发速率限制时是否抛出异常
         :return: 豆瓣信息
         """
         """
@@ -427,7 +429,10 @@ class DoubanModule(_ModuleBase):
             info = self.doubanapi.tv_detail(doubanid)
             if info:
                 if "subject_ip_rate_limit" in info.get("msg", ""):
-                    logger.warn(f"触发豆瓣IP速率限制，错误信息：{info} ...")
+                    msg = f"触发豆瓣IP速率限制，错误信息：{info} ..."
+                    logger.warn(msg)
+                    if raise_exception:
+                        raise APIRateLimitException(msg)
                     return None
                 celebrities = self.doubanapi.tv_celebrities(doubanid)
                 if celebrities:
@@ -442,7 +447,10 @@ class DoubanModule(_ModuleBase):
             info = self.doubanapi.movie_detail(doubanid)
             if info:
                 if "subject_ip_rate_limit" in info.get("msg", ""):
-                    logger.warn(f"触发豆瓣IP速率限制，错误信息：{info} ...")
+                    msg = f"触发豆瓣IP速率限制，错误信息：{info} ..."
+                    logger.warn(msg)
+                    if raise_exception:
+                        raise APIRateLimitException(msg)
                     return None
                 celebrities = self.doubanapi.movie_celebrities(doubanid)
                 if celebrities:
@@ -601,7 +609,8 @@ class DoubanModule(_ModuleBase):
 
     @retry(Exception, 5, 3, 3, logger=logger)
     def match_doubaninfo(self, name: str, imdbid: str = None,
-                         mtype: MediaType = None, year: str = None, season: int = None) -> dict:
+                         mtype: MediaType = None, year: str = None, season: int = None,
+                         raise_exception: bool = False) -> dict:
         """
         搜索和匹配豆瓣信息
         :param name:  名称
@@ -609,6 +618,7 @@ class DoubanModule(_ModuleBase):
         :param mtype:  类型
         :param year:  年份
         :param season:  季号
+        :param raise_exception: 触发速率限制时是否抛出异常
         """
         if imdbid:
             # 优先使用IMDBID查询
@@ -629,7 +639,10 @@ class DoubanModule(_ModuleBase):
             return {}
         # 触发rate limit
         if "search_access_rate_limit" in result.values():
-            logger.warn(f"触发豆瓣API速率限制，错误信息：{result} ...")
+            msg = f"触发豆瓣API速率限制，错误信息：{result} ..."
+            logger.warn(msg)
+            if raise_exception:
+                raise APIRateLimitException(msg)
             return {}
         if not result.get("items"):
             logger.warn(f"未找到 {name} 的豆瓣信息")
diff --git a/app/schemas/exception.py b/app/schemas/exception.py
index 8a536d71..63ef784e 100644
--- a/app/schemas/exception.py
+++ b/app/schemas/exception.py
@@ -4,3 +4,11 @@ class ImmediateException(Exception):
     当不希望使用重试机制时，可以抛出此异常。
     """
     pass
+
+
+class APIRateLimitException(ImmediateException):
+    """
+    用于表示API速率限制的异常类。
+    当API调用触发速率限制时，可以抛出此异常以立即终止操作并报告错误。
+    """
+    pass