✨ 添加cookie相关的数据库表

nonebot_bison/config/db_config.py

@@ -13,7 +13,7 @@ from nonebot_plugin_datastore import create_session
 from ..types import Tag
 from ..types import Target as T_Target
 from .utils import NoSuchTargetException
-from .db_model import User, Target, Subscribe, ScheduleTimeWeight
+from .db_model import User, Cookie, Target, Subscribe, CookieTarget, ScheduleTimeWeight
 from ..types import Category, UserSubInfo, WeightConfig, TimeWeightConfig, PlatformWeightConfigResp
 
 
@@ -259,5 +259,74 @@ class DBConfig:
             )
         return res
 
+    async def get_cookie_by_user(self, user: PlatformTarget) -> list[Cookie]:
+        async with create_session() as sess:
+            res = await sess.scalar(
+                select(User)
+                .where(User.user_target == model_dump(user))
+                .join(Cookie)
+                .outerjoin(CookieTarget)
+                .options(selectinload(User.cookies))
+            )
+            if not res:
+                return []
+            return res.cookies
+
+    async def add_cookie(self, user: PlatformTarget, platform_name: str, content: str) -> int:
+        async with create_session() as sess:
+            user_obj = await sess.scalar(select(User).where(User.user_target == model_dump(user)))
+            cookie = Cookie(user=user_obj, platform_name=platform_name, content=content)
+            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))
+            if not cookie_in_db:
+                return
+            cookie_in_db.content = cookie.content
+            cookie_in_db.last_usage = cookie.last_usage
+            cookie_in_db.status = cookie.status
+            cookie_in_db.tags = cookie.tags
+            await sess.commit()
+
+    async def delete_cookie(self, cookie_id: int):
+        async with create_session() as sess:
+            await sess.execute(delete(Cookie).where(Cookie.id == cookie_id))
+            await sess.commit()
+
+    async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> list[Cookie]:
+        async with create_session() as sess:
+            query = (
+                select(Cookie)
+                .join(CookieTarget)
+                .join(Target)
+                .where(Target.platform_name == platform_name, Target.target == target)
+            )
+            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(
+                select(Target).where(Target.platform_name == platform_name, Target.target == target)
+            )
+            cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id))
+            cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj)
+            sess.add(cookie_target)
+            await sess.commit()
+
+    async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int):
+        async with create_session() as sess:
+            target_obj = await sess.scalar(
+                select(Target).where(Target.platform_name == platform_name, Target.target == target)
+            )
+            cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id))
+            await sess.execute(
+                delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj)
+            )
+            await sess.commit()
+
 
 config = DBConfig()

nonebot_bison/config/db_model.py

@@ -1,4 +1,5 @@
 import datetime
+from typing import Any
 from pathlib import Path
 
 from nonebot_plugin_saa import PlatformTarget
@@ -6,7 +7,7 @@ from sqlalchemy.dialects.postgresql import JSONB
 from nonebot.compat import PYDANTIC_V2, ConfigDict
 from nonebot_plugin_datastore import get_plugin_data
 from sqlalchemy.orm import Mapped, relationship, mapped_column
-from sqlalchemy import JSON, String, ForeignKey, UniqueConstraint
+from sqlalchemy import JSON, String, DateTime, ForeignKey, UniqueConstraint
 
 from ..types import Tag, Category
 
@@ -19,6 +20,7 @@ class User(Model):
     user_target: Mapped[dict] = mapped_column(JSON().with_variant(JSONB, "postgresql"))
 
     subscribes: Mapped[list["Subscribe"]] = relationship(back_populates="user")
+    cookies: Mapped[list["Cookie"]] = relationship(back_populates="user")
 
     @property
     def saa_target(self) -> PlatformTarget:
@@ -36,6 +38,7 @@ class Target(Model):
 
     subscribes: Mapped[list["Subscribe"]] = relationship(back_populates="target")
     time_weight: Mapped[list["ScheduleTimeWeight"]] = relationship(back_populates="target")
+    cookies: Mapped[list["CookieTarget"]] = relationship(back_populates="target")
 
 
 class ScheduleTimeWeight(Model):
@@ -66,3 +69,25 @@ class Subscribe(Model):
 
     target: Mapped[Target] = relationship(back_populates="subscribes")
     user: Mapped[User] = relationship(back_populates="subscribes")
+
+
+class Cookie(Model):
+    id: Mapped[int] = mapped_column(primary_key=True)
+    user_id: Mapped[int] = mapped_column(ForeignKey("nonebot_bison_user.id"))
+    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))
+    status: Mapped[str] = mapped_column(String(20), default="")
+    tags: Mapped[dict[str, Any]] = mapped_column(JSON().with_variant(JSONB, "postgresql"), default={})
+
+    user: Mapped[User] = relationship(back_populates="cookies")
+    targets: Mapped[list["CookieTarget"]] = relationship(back_populates="cookie")
+
+
+class CookieTarget(Model):
+    id: Mapped[int] = mapped_column(primary_key=True)
+    target_id: Mapped[int] = mapped_column(ForeignKey("nonebot_bison_target.id", ondelete="CASCADE"))
+    cookie_id: Mapped[int] = mapped_column(ForeignKey("nonebot_bison_cookie.id", ondelete="CASCADE"))
+
+    target: Mapped[Target] = relationship(back_populates="cookies")
+    cookie: Mapped[Cookie] = relationship(back_populates="targets")

nonebot_bison/config/migrations/bdec8925b540_add_cookie.py

@@ -0,0 +1,63 @@
+"""empty message
+
+Revision ID: bdec8925b540
+Revises: f9baef347cc8
+Create Date: 2024-08-22 20:53:08.850051
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import Text
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "bdec8925b540"
+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("user_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("tags", sa.JSON().with_variant(postgresql.JSONB(astext_type=Text()), "postgresql"), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["user_id"], ["nonebot_bison_user.id"], name=op.f("fk_nonebot_bison_cookie_user_id_nonebot_bison_user")
+        ),
+        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 ###

tests/config/test_cookie.py

@@ -0,0 +1,107 @@
+import datetime
+
+from nonebug import App
+
+
+async def test_get_platform_target(app: App, init_scheduler):
+    from nonebot_plugin_saa import TargetQQGroup
+
+    from nonebot_bison.config.db_config import config
+    from nonebot_bison.types import Target as T_Target
+
+    await config.add_subscribe(
+        TargetQQGroup(group_id=123),
+        target=T_Target("weibo_id"),
+        target_name="weibo_name",
+        platform_name="weibo",
+        cats=[],
+        tags=[],
+    )
+    # await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie")
+    # cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123))
+    #
+    # res = await config.get_platform_target("weibo")
+    # assert len(res) == 2
+    # await config.del_subscribe(TargetQQGroup(group_id=123), T_Target("weibo_id1"), "weibo")
+    # res = await config.get_platform_target("weibo")
+    # assert len(res) == 2
+    # await config.del_subscribe(TargetQQGroup(group_id=123), T_Target("weibo_id"), "weibo")
+    # res = await config.get_platform_target("weibo")
+    # assert len(res) == 1
+    #
+    # async with AsyncSession(get_engine()) as sess:
+    #     res = await sess.scalars(select(Target).where(Target.platform_name == "weibo"))
+    #     assert len(res.all()) == 2
+    # await config.get_cookie_by_user(TargetQQGroup(group_id=123))
+
+
+async def test_cookie_by_user(app: App, init_scheduler):
+    from nonebot_plugin_saa import TargetQQGroup
+
+    from nonebot_bison.config.db_config import config
+    from nonebot_bison.types import Target as T_Target
+
+    await config.add_subscribe(
+        TargetQQGroup(group_id=123),
+        target=T_Target("weibo_id"),
+        target_name="weibo_name",
+        platform_name="weibo",
+        cats=[],
+        tags=[],
+    )
+
+    await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie")
+
+    cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123))
+    cookie = cookies[0]
+    assert len(cookies) == 1
+    assert cookie.content == "cookie"
+    assert cookie.platform_name == "weibo"
+    cookie.last_usage = 0
+    assert cookie.status == ""
+    assert cookie.tags == {}
+    cookie.content = "cookie1"
+    cookie.last_usage = datetime.datetime(2024, 8, 22, 0, 0, 0)
+    cookie.status = "status1"
+    cookie.tags = {"tag1": "value1"}
+    await config.update_cookie(cookie)
+    cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123))
+
+    assert len(cookies) == 1
+    assert cookies[0].content == cookie.content
+    assert cookies[0].last_usage == cookie.last_usage
+    assert cookies[0].status == cookie.status
+    assert cookies[0].tags == cookie.tags
+
+    await config.delete_cookie(cookies[0].id)
+    cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123))
+    assert len(cookies) == 0
+
+
+async def test_cookie_target_by_target(app: App, init_scheduler):
+    from nonebot_plugin_saa import TargetQQGroup
+
+    from nonebot_bison.config.db_config import config
+    from nonebot_bison.types import Target as T_Target
+
+    await config.add_subscribe(
+        TargetQQGroup(group_id=123),
+        target=T_Target("weibo_id"),
+        target_name="weibo_name",
+        platform_name="weibo",
+        cats=[],
+        tags=[],
+    )
+
+    id = await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie")
+
+    await config.add_cookie_target(T_Target("weibo_id"), "weibo", id)
+
+    cookies = await config.get_cookie_by_target(T_Target("weibo_id"), "weibo")
+    assert len(cookies) == 1
+    assert cookies[0].content == "cookie"
+    assert cookies[0].platform_name == "weibo"
+
+    await config.delete_cookie_target(T_Target("weibo_id"), "weibo", id)
+    cookies = await config.get_cookie_by_target(T_Target("weibo_id"), "weibo")
+    assert len(cookies) == 0