From 1cd778c2e0506208b5520a755917c65f3692ea13 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sat, 31 Aug 2024 23:07:43 +0800 Subject: [PATCH] =?UTF-8?q?:recycle:=20cookie=20=E7=BB=84=E4=BB=B6?= =?UTF-8?q?=E4=B8=8D=E5=86=8D=E4=B8=8E=20user=20=E5=85=B3=E8=81=94?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 28 ++------- nonebot_bison/config/db_model.py | 3 - .../migrations/590dc2911ea7_add_cookie.py | 59 +++++++++++++++++++ nonebot_bison/sub_manager/add_cookie.py | 9 +-- .../sub_manager/add_cookie_target.py | 21 +++---- nonebot_bison/sub_manager/utils.py | 15 ++++- 6 files changed, 85 insertions(+), 50 deletions(-) create mode 100644 nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 97cc0db..909d6f9 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -259,14 +259,9 @@ class DBConfig: ) return res - async def get_cookie( - self, user: PlatformTarget = None, platform_name: str = None, target: T_Target = None - ) -> list[Cookie]: + async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> list[Cookie]: async with create_session() as sess: - query = select(Cookie).distinct().join(User) - if user: - user_id = await sess.scalar(select(User.id).where(User.user_target == model_dump(user))) - query = query.where(Cookie.user_id == user_id) + query = select(Cookie).distinct() if platform_name: query = query.where(Cookie.platform_name == platform_name) query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets)) @@ -277,24 +272,9 @@ class DBConfig: res = [cookie for cookie in res if cookie.id in ids] return res - async def get_cookie_by_user_and_platform(self, user: PlatformTarget, platform_name: str) -> list[Cookie]: + async def add_cookie(self, platform_name: str, content: str) -> int: async with create_session() as sess: - res = await sess.scalar( - select(User) - .where(User.user_target == model_dump(user)) - .join(Cookie) - .where(Cookie.platform_name == platform_name) - .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) + cookie = Cookie(platform_name=platform_name, content=content) sess.add(cookie) await sess.commit() await sess.refresh(cookie) diff --git a/nonebot_bison/config/db_model.py b/nonebot_bison/config/db_model.py index 603ae21..2c74721 100644 --- a/nonebot_bison/config/db_model.py +++ b/nonebot_bison/config/db_model.py @@ -20,7 +20,6 @@ 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: @@ -73,14 +72,12 @@ class Subscribe(Model): 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") diff --git a/nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py b/nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py new file mode 100644 index 0000000..c8fbd1d --- /dev/null +++ b/nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py @@ -0,0 +1,59 @@ +"""empty message + +Revision ID: 590dc2911ea7 +Revises: f9baef347cc8 +Create Date: 2024-08-31 23:06:02.123932 + +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy import Text +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "590dc2911ea7" +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("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 ### diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index f6ba356..8702a3c 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -1,21 +1,18 @@ from nonebot.typing import T_State from nonebot.matcher import Matcher from nonebot.params import Arg, ArgPlainText -from nonebot_plugin_saa import PlatformTarget from nonebot.adapters import Message, MessageTemplate from ..types import Target from ..config import config from ..platform import platform_manager from ..apis import check_sub_target_cookie -from .utils import common_platform, ensure_user_info, gen_handle_cancel +from .utils import common_platform, gen_handle_cancel def do_add_cookie(add_cookie: type[Matcher]): handle_cancel = gen_handle_cancel(add_cookie, "已中止添加cookie") - add_cookie.handle()(ensure_user_info(add_cookie)) - @add_cookie.handle() async def init_promote(state: T_State): state["_prompt"] = ( @@ -56,8 +53,8 @@ def do_add_cookie(add_cookie: type[Matcher]): state["name"] = await check_sub_target_cookie(state["platform"], Target(""), cookie_text) @add_cookie.handle() - async def add_cookie_process(state: T_State, user: PlatformTarget = Arg("target_user_info")): - await config.add_cookie(user, state["platform"], state["cookie"]) + async def add_cookie_process(state: T_State): + await config.add_cookie(state["platform"], state["cookie"]) await add_cookie.finish( f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}" + "\n请使用“关联cookie”为 Cookie 关联订阅" ) diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py index 7244736..5a1a4d0 100644 --- a/nonebot_bison/sub_manager/add_cookie_target.py +++ b/nonebot_bison/sub_manager/add_cookie_target.py @@ -1,25 +1,21 @@ from nonebot.typing import T_State from nonebot.matcher import Matcher -from nonebot.params import Arg, ArgPlainText +from nonebot.params import ArgPlainText +from nonebot_plugin_saa import MessageFactory from nonebot.internal.adapter import MessageTemplate -from nonebot_plugin_saa import MessageFactory, PlatformTarget from ..config import config from ..utils import parse_text from ..apis import get_cookie_friendly_name -from .utils import ensure_user_info, gen_handle_cancel, generate_sub_list_text +from .utils import gen_handle_cancel, generate_sub_list_text def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): handle_cancel = gen_handle_cancel(add_cookie_target_matcher, "已中止关联 cookie") - add_cookie_target_matcher.handle()(ensure_user_info(add_cookie_target_matcher)) - @add_cookie_target_matcher.handle() - async def init_promote(state: T_State, user_info: PlatformTarget = Arg("target_user_info")): - res = await generate_sub_list_text( - add_cookie_target_matcher, state, user_info, is_index=True, is_show_cookie=True - ) + async def init_promote(state: T_State): + res = await generate_sub_list_text(add_cookie_target_matcher, state, is_index=True, is_show_cookie=True) res += "请输入要关联 cookie 的订阅的序号\n输入'取消'中止" await MessageFactory(await parse_text(res)).send() @@ -34,11 +30,8 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): @add_cookie_target_matcher.handle() async def init_promote_cookie(state: T_State): - cookies = await config.get_cookie( - user=state["target_user_info"], platform_name=state["target"]["platform_name"] - ) + cookies = await config.get_cookie(platform_name=state["target"]["platform_name"]) associated_cookies = await config.get_cookie( - user=state["target_user_info"], target=state["target"]["target"], platform_name=state["target"]["platform_name"], ) @@ -62,7 +55,7 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): await add_cookie_target_matcher.reject("序号错误") @add_cookie_target_matcher.handle() - async def add_cookie_target_process(state: T_State, user: PlatformTarget = Arg("target_user_info")): + async def add_cookie_target_process(state: T_State): await config.add_cookie_target(state["target"]["target"], state["target"]["platform_name"], state["cookie"].id) await add_cookie_target_matcher.finish( f"已关联 Cookie: {await get_cookie_friendly_name(state['cookie'])} " diff --git a/nonebot_bison/sub_manager/utils.py b/nonebot_bison/sub_manager/utils.py index 71a06e3..4b7dd6f 100644 --- a/nonebot_bison/sub_manager/utils.py +++ b/nonebot_bison/sub_manager/utils.py @@ -1,5 +1,7 @@ import contextlib from typing import Annotated +from itertools import groupby +from operator import attrgetter from nonebot.rule import Rule from nonebot.adapters import Event @@ -66,9 +68,16 @@ def admin_permission(): async def generate_sub_list_text( - matcher: type[Matcher], state: T_State, user_info: PlatformTarget, is_index=False, is_show_cookie=False + matcher: type[Matcher], state: T_State, user_info: PlatformTarget = None, is_index=False, is_show_cookie=False ): - sub_list = await config.list_subscribe(user_info) + if user_info: + sub_list = await config.list_subscribe(user_info) + else: + sub_list = await config.list_subs_with_all_info() + sub_list = [ + next(group) + for key, group in groupby(sorted(sub_list, key=attrgetter("target_id")), key=attrgetter("target_id")) + ] if not sub_list: await matcher.finish("暂无已订阅账号\n请使用“添加订阅”命令添加订阅") res = "订阅的帐号为：\n" @@ -88,7 +97,7 @@ async def generate_sub_list_text( res += " {}".format(", ".join(sub.tags)) + "\n" if is_show_cookie: target_cookies = await config.get_cookie( - user=user_info, target=sub.target.target, platform_name=sub.target.platform_name + target=sub.target.target, platform_name=sub.target.platform_name ) if target_cookies: res += " 关联的 Cookie：\n"