✨ 匿名cookie和用户cookie一起调度

nonebot_bison/config/db_config.py

@@ -272,7 +272,20 @@ class DBConfig:
                 res = [cookie for cookie in res if cookie.id in ids]
             return res
 
-    async def add_cookie(self, platform_name: str, content: str) -> int:
+    async def get_unviersal_cookie(self, platform_name: str = None, target: T_Target = None) -> list[Cookie]:
+        async with create_session() as sess:
+            query = select(Cookie).distinct().where(Cookie.is_universal == True)  # noqa: E712
+            if platform_name:
+                query = query.where(Cookie.platform_name == platform_name)
+            query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets))
+            res = (await sess.scalars(query)).all()
+            if target:
+                query = select(CookieTarget.cookie_id).join(Target).where(Target.target == target)
+                ids = set((await sess.scalars(query)).all())
+                res = [cookie for cookie in res if cookie.id in ids]
+            return res
+
+    async def add_cookie_with_content(self, platform_name: str, content: str) -> int:
         async with create_session() as sess:
             cookie = Cookie(platform_name=platform_name, content=content)
             sess.add(cookie)
@@ -280,6 +293,13 @@ class DBConfig:
             await sess.refresh(cookie)
             return cookie.id
 
+    async def add_cookie(self, cookie: Cookie) -> int:
+        async with create_session() as sess:
+            sess.add(cookie)
+            await sess.commit()
+            await sess.refresh(cookie)
+            return cookie.id
+
     async def update_cookie(self, cookie: Cookie):
         async with create_session() as sess:
             cookie_in_db: Cookie | None = await sess.scalar(select(Cookie).where(Cookie.id == cookie.id))
@@ -306,6 +326,15 @@ class DBConfig:
             )
             return list((await sess.scalars(query)).all())
 
+    async def get_universal_cookie(self, platform_name: str) -> list[Cookie]:
+        async with create_session() as sess:
+            query = (
+                select(Cookie)
+                .where(Cookie.platform_name == platform_name)
+                .where(Cookie.is_universal == True)  # noqa: E712
+            )
+            return list((await sess.scalars(query)).all())
+
     async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int):
         async with create_session() as sess:
             target_obj = await sess.scalar(

nonebot_bison/config/db_model.py

@@ -74,10 +74,15 @@ class Cookie(Model):
     id: Mapped[int] = mapped_column(primary_key=True)
     platform_name: Mapped[str] = mapped_column(String(20))
     content: Mapped[str] = mapped_column(String(1024))
+    # 最后使用的时刻
     last_usage: Mapped[datetime.datetime] = mapped_column(DateTime, default=datetime.datetime(1970, 1, 1))
+    # Cookie 当前的状态
     status: Mapped[str] = mapped_column(String(20), default="")
+    # 使用一次之后，需要的冷却时间
     cd: Mapped[int] = mapped_column(default=0)
+    # 是否是通用 Cookie，默认用于匿名 Cookie
     is_universal: Mapped[bool] = mapped_column(default=False)
+    # 标签，扩展用
     tags: Mapped[dict[str, Any]] = mapped_column(JSON().with_variant(JSONB, "postgresql"), default={})
 
     targets: Mapped[list["CookieTarget"]] = relationship(back_populates="cookie")