From 7913f7485a8ca9d5d15cfeb0cb9f2b55b916dc5c Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Thu, 22 Aug 2024 20:55:39 +0800 Subject: [PATCH] =?UTF-8?q?:sparkles:=20=E6=B7=BB=E5=8A=A0cookie=E7=9B=B8?= =?UTF-8?q?=E5=85=B3=E7=9A=84=E6=95=B0=E6=8D=AE=E5=BA=93=E8=A1=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 71 +++++++++++- nonebot_bison/config/db_model.py | 27 ++++- .../migrations/bdec8925b540_add_cookie.py | 63 +++++++++++ nonebot_bison/sub_manager/add_cookie.py | 0 tests/config/__init__.py | 0 tests/config/test_cookie.py | 107 ++++++++++++++++++ 6 files changed, 266 insertions(+), 2 deletions(-) create mode 100644 nonebot_bison/config/migrations/bdec8925b540_add_cookie.py create mode 100644 nonebot_bison/sub_manager/add_cookie.py create mode 100644 tests/config/__init__.py create mode 100644 tests/config/test_cookie.py diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 157b1ef..444d9d9 100644 --- a/nonebot_bison/config/db_config.py +++ b/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() diff --git a/nonebot_bison/config/db_model.py b/nonebot_bison/config/db_model.py index 849094d..603ae21 100644 --- a/nonebot_bison/config/db_model.py +++ b/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") diff --git a/nonebot_bison/config/migrations/bdec8925b540_add_cookie.py b/nonebot_bison/config/migrations/bdec8925b540_add_cookie.py new file mode 100644 index 0000000..ac650f3 --- /dev/null +++ b/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 ### diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/config/__init__.py b/tests/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/config/test_cookie.py b/tests/config/test_cookie.py new file mode 100644 index 0000000..4515078 --- /dev/null +++ b/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