✨ 数据库Cookie表添加is_universal属性

nonebot_bison/config/db_model.py

@@ -77,6 +77,7 @@ class Cookie(Model):
     last_usage: Mapped[datetime.datetime] = mapped_column(DateTime, default=datetime.datetime(1970, 1, 1))
     status: Mapped[str] = mapped_column(String(20), default="")
     cd: Mapped[int] = mapped_column(default=0)
+    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")

nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py

@@ -0,0 +1,61 @@
+"""empty message
+
+Revision ID: fc2b8481bdde
+Revises: f9baef347cc8
+Create Date: 2024-09-05 19:31:59.366754
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import Text
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "fc2b8481bdde"
+down_revision = "f9baef347cc8"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "nonebot_bison_cookie",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("platform_name", sa.String(length=20), nullable=False),
+        sa.Column("content", sa.String(length=1024), nullable=False),
+        sa.Column("last_usage", sa.DateTime(), nullable=False),
+        sa.Column("status", sa.String(length=20), nullable=False),
+        sa.Column("cd", sa.Integer(), nullable=False),
+        sa.Column("is_universal", sa.Boolean(), nullable=False),
+        sa.Column("tags", sa.JSON().with_variant(postgresql.JSONB(astext_type=Text()), "postgresql"), nullable=False),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookie")),
+    )
+    op.create_table(
+        "nonebot_bison_cookietarget",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("target_id", sa.Integer(), nullable=False),
+        sa.Column("cookie_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["cookie_id"],
+            ["nonebot_bison_cookie.id"],
+            name=op.f("fk_nonebot_bison_cookietarget_cookie_id_nonebot_bison_cookie"),
+            ondelete="CASCADE",
+        ),
+        sa.ForeignKeyConstraint(
+            ["target_id"],
+            ["nonebot_bison_target.id"],
+            name=op.f("fk_nonebot_bison_cookietarget_target_id_nonebot_bison_target"),
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookietarget")),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("nonebot_bison_cookietarget")
+    op.drop_table("nonebot_bison_cookie")
+    # ### end Alembic commands ###

nonebot_bison/utils/site.py

@@ -71,11 +71,16 @@ class CookieClientManager(ClientManager):
         cookies = (
             cookie for cookie in cookies if cookie.last_usage + timedelta(seconds=self._cookie_cd) < datetime.now()
         )
+        cookies = list(cookies)
         if not cookies:
             return Cookie(content="{}")
         cookie = max(cookies, key=lambda x: x.last_usage)
         return cookie
 
+    async def _check_cookie(self, cookie: Cookie) -> Cookie:
+        """检查Cookie，可以做一些自定义的逻辑，比如说Site的统一风控"""
+        return cookie
+
     async def get_client(self, target: Target | None) -> AsyncClient:
         """获取 client，根据 target 选择 cookie"""
         client = http_client()
@@ -85,9 +90,9 @@ class CookieClientManager(ClientManager):
         else:
             logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用空cookie")
 
-        return await self.assemble_client(client, cookie)
+        return await self._assemble_client(client, cookie)
 
-    async def assemble_client(self, client, cookie):
+    async def _assemble_client(self, client, cookie):
         cookies = httpx.Cookies()
         if cookie:
             cookies.update(json.loads(cookie.content))