From 990012520ac335a05c1ef2e5794ee00bd2242984 Mon Sep 17 00:00:00 2001 From: omg-xtao <100690902+omg-xtao@users.noreply.github.com> Date: Mon, 18 Nov 2024 21:38:50 +0800 Subject: [PATCH] :sparkles: Support mult user task service --- .../versions/a9c8cde17cd8_tasks_player_id.py | 82 +++++++++++++++++++ gram_core | 2 +- plugins/tools/daily_note.py | 52 +++++++----- plugins/zzz/daily_note_tasks.py | 54 ++++++------ plugins/zzz/sign.py | 32 ++++---- 5 files changed, 162 insertions(+), 60 deletions(-) create mode 100644 alembic/versions/a9c8cde17cd8_tasks_player_id.py diff --git a/alembic/versions/a9c8cde17cd8_tasks_player_id.py b/alembic/versions/a9c8cde17cd8_tasks_player_id.py new file mode 100644 index 0000000..0c16212 --- /dev/null +++ b/alembic/versions/a9c8cde17cd8_tasks_player_id.py @@ -0,0 +1,82 @@ +"""tasks_player_id + +Revision ID: a9c8cde17cd8 +Revises: 1220c5c80757 +Create Date: 2024-11-18 15:47:51.218077 + +""" + +import logging + +from alembic import op +import sqlalchemy as sa +from sqlalchemy import text +from sqlalchemy.dialects import mysql +from sqlalchemy.exc import NoSuchTableError + +# revision identifiers, used by Alembic. +revision = "a9c8cde17cd8" +down_revision = "1220c5c80757" +branch_labels = None +depends_on = None + +logger = logging.getLogger(__name__) + + +def update_player_id(): + connection = op.get_bind() + try: + statement = "SELECT user_id FROM task;" + task_table_data = connection.execute(text(statement)) + need_check_user_id = [] + if task_table_data is not None: + for row in task_table_data: + need_check_user_id.append(row[0]) + need_check_user_id = list(set(need_check_user_id)) + except NoSuchTableError: + logger.warning("Table 'task' doesn't exist") + return # should not happen + try: + statement = "SELECT user_id, player_id, is_chosen FROM players;" + players_table_data = connection.execute(text(statement)) + player_id_map = {} + if players_table_data is not None: + for row in players_table_data: + if not row[2]: + continue + uid, pid = row[0], row[1] + if uid not in player_id_map: + player_id_map[uid] = pid + except NoSuchTableError: + logger.warning("Table 'players' doesn't exist") + return # should not happen + + update = "UPDATE task SET player_id=:player_id WHERE user_id=:user_id;" + + for uid in need_check_user_id: + player_id = player_id_map.get(uid, None) + if player_id is None: + logger.warning("user_id %s doesn't exist player", uid) + continue + try: + with op.get_context().autocommit_block(): + connection.execute(text(update), dict(player_id=player_id, user_id=uid)) + except Exception as exc: # pylint: disable=W0703 + logger.error("Process sign->task Exception", exc_info=exc) # pylint: disable=W0703 + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("task", sa.Column("player_id", sa.BigInteger(), nullable=False)) + op.alter_column("task", "user_id", existing_type=mysql.BIGINT(), nullable=False) + op.create_index(op.f("ix_task_player_id"), "task", ["player_id"], unique=False) + update_player_id() + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index(op.f("ix_task_player_id"), table_name="task") + op.alter_column("task", "user_id", existing_type=mysql.BIGINT(), nullable=True) + op.drop_column("task", "player_id") + # ### end Alembic commands ### diff --git a/gram_core b/gram_core index 7129322..112b2e9 160000 --- a/gram_core +++ b/gram_core @@ -1 +1 @@ -Subproject commit 7129322f6dd9ae82880807a0a5a8579a5f150369 +Subproject commit 112b2e92d8492df17dbae23024fa805e3510a56e diff --git a/plugins/tools/daily_note.py b/plugins/tools/daily_note.py index 071e62e..339964b 100644 --- a/plugins/tools/daily_note.py +++ b/plugins/tools/daily_note.py @@ -52,6 +52,9 @@ class DailyData(TaskDataBase): class WebAppData(BaseModel): + user_id: int + player_id: int + resin: Optional[ResinData] expedition: Optional[ExpeditionData] daily: Optional[DailyData] @@ -61,11 +64,13 @@ class DailyNoteTaskUser: def __init__( self, user_id: int, + player_id: int, resin_db: Optional[TaskUser] = None, expedition_db: Optional[TaskUser] = None, daily_db: Optional[TaskUser] = None, ): self.user_id = user_id + self.player_id = player_id self.resin_db = resin_db self.expedition_db = expedition_db self.daily_db = daily_db @@ -107,6 +112,8 @@ class DailyNoteTaskUser: return base64.b64encode( ( WebAppData( + user_id=self.user_id, + player_id=self.player_id, resin=self.set_model_noticed(self.resin) if self.resin else None, expedition=self.set_model_noticed(self.expedition) if self.expedition else None, daily=self.set_model_noticed(self.daily) if self.daily else None, @@ -136,12 +143,13 @@ class DailyNoteSystem(Plugin): self.expedition_service = expedition_service self.daily_service = daily_service - async def get_single_task_user(self, user_id: int) -> DailyNoteTaskUser: - resin_db = await self.resin_service.get_by_user_id(user_id) - expedition_db = await self.expedition_service.get_by_user_id(user_id) - daily_db = await self.daily_service.get_by_user_id(user_id) + async def get_single_task_user(self, user_id: int, player_id: int) -> DailyNoteTaskUser: + resin_db = await self.resin_service.get_by_user_id(user_id, player_id) + expedition_db = await self.expedition_service.get_by_user_id(user_id, player_id) + daily_db = await self.daily_service.get_by_user_id(user_id, player_id) return DailyNoteTaskUser( user_id=user_id, + player_id=player_id, resin_db=resin_db, expedition_db=expedition_db, daily_db=daily_db, @@ -212,20 +220,17 @@ class DailyNoteSystem(Plugin): expedition_list = await self.expedition_service.get_all() daily_list = await self.daily_service.get_all() user_list = set() - for i in resin_list: - user_list.add(i.user_id) - for i in expedition_list: - user_list.add(i.user_id) - for i in daily_list: - user_list.add(i.user_id) + for i in resin_list + expedition_list + daily_list: + user_list.add((i.user_id, i.player_id)) return [ DailyNoteTaskUser( user_id=i, - resin_db=next((x for x in resin_list if x.user_id == i), None), - expedition_db=next((x for x in expedition_list if x.user_id == i), None), - daily_db=next((x for x in daily_list if x.user_id == i), None), + player_id=p, + resin_db=next((x for x in resin_list if x.user_id == i and x.player_id == p), None), + expedition_db=next((x for x in expedition_list if x.user_id == i and x.player_id == p), None), + daily_db=next((x for x in daily_list if x.user_id == i and x.player_id == p), None), ) - for i in user_list + for i, p in user_list ] async def remove_task_user(self, user: DailyNoteTaskUser): @@ -265,11 +270,12 @@ class DailyNoteSystem(Plugin): return need_verify async def import_web_config_resin(self, user: DailyNoteTaskUser, web_config: WebAppData): - user_id = user.user_id + user_id, player_id = user.user_id, user.player_id if web_config.resin.noticed: if not user.resin_db: resin = self.resin_service.create( user_id, + player_id, user_id, status=TaskStatusEnum.STATUS_SUCCESS, data=ResinData(notice_num=web_config.resin.notice_num).dict(), @@ -286,11 +292,12 @@ class DailyNoteSystem(Plugin): user.resin = None async def import_web_config_expedition(self, user: DailyNoteTaskUser, web_config: WebAppData): - user_id = user.user_id + user_id, player_id = user.user_id, user.player_id if web_config.expedition.noticed: if not user.expedition_db: expedition = self.expedition_service.create( user_id, + player_id, user_id, status=TaskStatusEnum.STATUS_SUCCESS, data=ExpeditionData().dict(), @@ -306,11 +313,12 @@ class DailyNoteSystem(Plugin): user.expedition = None async def import_web_config_daily(self, user: DailyNoteTaskUser, web_config: WebAppData): - user_id = user.user_id + user_id, player_id = user.user_id, user.player_id if web_config.daily.noticed: if not user.daily_db: daily = self.daily_service.create( user_id, + player_id, user_id, status=TaskStatusEnum.STATUS_SUCCESS, data=DailyData(notice_hour=web_config.daily.notice_hour).dict(), @@ -325,8 +333,9 @@ class DailyNoteSystem(Plugin): user.daily_db = None user.daily = None - async def import_web_config(self, user_id: int, web_config: WebAppData): - user = await self.get_single_task_user(user_id) + async def import_web_config(self, web_config: WebAppData): + user_id, player_id = web_config.user_id, web_config.player_id + user = await self.get_single_task_user(user_id, player_id) if web_config.resin: await self.import_web_config_resin(user, web_config) if web_config.expedition: @@ -346,9 +355,10 @@ class DailyNoteSystem(Plugin): if task_db.status not in include_status: continue user_id = task_db.user_id - logger.debug("自动便签提醒 - 请求便签信息 user_id[%s]", user_id) + player_id = task_db.player_id + logger.debug("自动便签提醒 - 请求便签信息 user_id[%s] player_id[%s]", user_id, player_id) try: - async with self.genshin_helper.genshin(user_id) as client: + async with self.genshin_helper.genshin(user_id, player_id=player_id) as client: text = await self.start_get_notes(client, task_db) except InvalidCookies: text = "自动便签提醒执行失败，Cookie无效" diff --git a/plugins/zzz/daily_note_tasks.py b/plugins/zzz/daily_note_tasks.py index e50a3a6..8907e32 100644 --- a/plugins/zzz/daily_note_tasks.py +++ b/plugins/zzz/daily_note_tasks.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Union, Optional from pydantic import ValidationError from simnet import Region @@ -13,7 +13,7 @@ from core.services.cookies.services import CookiesService from core.services.players.services import PlayersService, PlayerInfoService from plugins.app.webapp import WebApp from plugins.tools.daily_note import DailyNoteSystem, WebAppData -from plugins.tools.genshin import GenshinHelper, CookiesNotFoundError, PlayerNotFoundError +from plugins.tools.genshin import GenshinHelper from utils.log import logger if TYPE_CHECKING: @@ -46,44 +46,50 @@ class DailyNoteTasksPlugin(Plugin.Conversation): @handler.command(command="daily_note_tasks", filters=filters.ChatType.PRIVATE, block=False) @handler.command(command="start", filters=filters.Regex(r"daily_note_tasks$"), block=False) async def command_start(self, update: Update, _: CallbackContext) -> int: - user = update.effective_user + user_id = await self.get_real_user_id(update) + uid, offset = self.get_real_uid_or_offset(update) message = update.effective_message - logger.info("用户 %s[%s] 设置自动便签提醒命令请求", user.full_name, user.id) - text = await self.check_genshin_user(user.id, False) - if text != "ok": + self.log_user(update, logger.info, "设置自动便签提醒命令请求") + text = await self.check_genshin_user(user_id, uid, offset, False) + if isinstance(text, str): await message.reply_text(text, reply_markup=ReplyKeyboardRemove()) return ConversationHandler.END - note_user = await self.note_system.get_single_task_user(user.id) + note_user = await self.note_system.get_single_task_user(user_id, text) url = f"{config.pass_challenge_user_web}/tasks4?command=tasks&bot_data={note_user.web_config}" - text = ( - f'你好 {user.mention_markdown_v2()} {escape_markdown("！请点击下方按钮，开始设置，或者回复退出取消操作")}' - ) + text = escape_markdown("请点击下方按钮，开始设置，或者回复退出取消操作") await message.reply_markdown_v2( text, - reply_markup=ReplyKeyboardMarkup.from_button( - KeyboardButton( - text="点我开始设置", - web_app=WebAppInfo(url=url), - ) + reply_markup=ReplyKeyboardMarkup( + [ + [ + KeyboardButton( + text="点我开始设置", + web_app=WebAppInfo(url=url), + ) + ], + [ + "退出", + ], + ] ), ) return SET_BY_WEB - async def check_genshin_user(self, user_id: int, request_note: bool) -> str: + async def check_genshin_user( + self, user_id: int, uid: int, offset: Optional[int], request_note: bool + ) -> Union[str, int]: try: - async with self.helper.genshin(user_id) as client: + async with self.helper.genshin(user_id, player_id=uid, offset=offset) as client: client: "ZZZClient" if request_note: await client.get_zzz_notes() - return "ok" + return client.player_id except ValueError: return "Cookies 缺少 stoken ，请尝试重新绑定账号。" except DataNotPublic: return "查询失败惹，可能是便签功能被禁用了？请尝试通过米游社或者 hoyolab 获取一次便签信息后重试。" except SimnetBadRequest as e: return f"获取便签失败，可能遇到验证码风控，请尝试重新绑定账号。{e}" - except (CookiesNotFoundError, PlayerNotFoundError): - return config.notice.user_not_found @conversation.state(state=SET_BY_WEB) @handler.message(filters=filters.TEXT & ~filters.COMMAND, block=False) @@ -107,17 +113,19 @@ class DailyNoteTasksPlugin(Plugin.Conversation): if result.path == "tasks": try: validate = WebAppData(**result.data) + if validate.user_id != user.id: + raise ValidationError(None, None) except ValidationError: await message.reply_text( "数据错误\n电量提醒数值必须在 60 ~ 200 之间\n每日任务提醒时间必须在 0 ~ 23 之间", reply_markup=ReplyKeyboardRemove(), ) return ConversationHandler.END - need_note = await self.check_genshin_user(user.id, True) - if need_note != "ok": + need_note = await self.check_genshin_user(validate.user_id, validate.player_id, None, True) + if isinstance(need_note, str): await message.reply_text(need_note, reply_markup=ReplyKeyboardRemove()) return ConversationHandler.END - await self.note_system.import_web_config(user.id, validate) + await self.note_system.import_web_config(validate) await message.reply_text("修改设置成功", reply_markup=ReplyKeyboardRemove()) else: logger.warning( diff --git a/plugins/zzz/sign.py b/plugins/zzz/sign.py index 56fa6af..2fcd547 100644 --- a/plugins/zzz/sign.py +++ b/plugins/zzz/sign.py @@ -40,37 +40,38 @@ class Sign(Plugin): self.players_service = player self.cookies_service = cookies - async def _process_auto_sign(self, user_id: int, chat_id: int, method: str) -> str: - player = await self.players_service.get_player(user_id) + async def _process_auto_sign(self, user_id: int, player_id: int, offset: int, chat_id: int, method: str) -> str: + player = await self.players_service.get_player(user_id, player_id=player_id, offset=offset) if player is None: return config.notice.user_not_found cookie_model = await self.cookies_service.get(player.user_id, player.account_id, player.region) if cookie_model is None: return config.notice.user_not_found - user: SignUser = await self.sign_service.get_by_user_id(user_id) + user: SignUser = await self.sign_service.get_by_user_id(user_id, player.player_id) if user: if method == "关闭": await self.sign_service.remove(user) - return "关闭自动签到成功" + return f"UID {player.player_id} 关闭自动签到成功" if method == "开启": if user.chat_id == chat_id: - return "自动签到已经开启过了" + return f"UID {player.player_id} 自动签到已经开启过了" user.chat_id = chat_id user.status = TaskStatusEnum.STATUS_SUCCESS await self.sign_service.update(user) - return "修改自动签到通知对话成功" + return f"UID {player.player_id} 修改自动签到通知对话成功" elif method == "关闭": - return "您还没有开启自动签到" + return f"UID {player.player_id} 您还没有开启自动签到" elif method == "开启": - user = self.sign_service.create(user_id, chat_id, TaskStatusEnum.STATUS_SUCCESS) + user = self.sign_service.create(user_id, player.player_id, chat_id, TaskStatusEnum.STATUS_SUCCESS) await self.sign_service.add(user) - return "开启自动签到成功" + return f"UID {player.player_id} 开启自动签到成功" @handler.command(command="sign", cookie=True, block=False) @handler.message(filters=filters.Regex("^每日签到(.*)"), cookie=True, block=False) @handler.command(command="start", filters=filters.Regex("sign$"), block=False) async def command_start(self, update: Update, context: CallbackContext) -> None: user_id = await self.get_real_user_id(update) + uid, offset = self.get_real_uid_or_offset(update) message = update.effective_message args = self.get_args(context) validate: Optional[str] = None @@ -78,12 +79,12 @@ class Sign(Plugin): msg = None if args[0] == "开启自动签到": if await self.user_admin_service.is_admin(user_id): - msg = await self._process_auto_sign(user_id, message.chat_id, "开启") + msg = await self._process_auto_sign(user_id, uid, offset, message.chat_id, "开启") else: - msg = await self._process_auto_sign(user_id, user_id, "开启") + msg = await self._process_auto_sign(user_id, uid, offset, user_id, "开启") elif args[0] == "关闭自动签到": - msg = await self._process_auto_sign(user_id, message.chat_id, "关闭") - elif args[0] != "sign": + msg = await self._process_auto_sign(user_id, uid, offset, message.chat_id, "关闭") + elif args[0] != "sign" and not args[0].startswith("@"): validate = args[0] if msg: self.log_user(update, logger.info, "自动签到命令请求 || 参数 %s", args[0]) @@ -96,7 +97,7 @@ class Sign(Plugin): if filters.ChatType.GROUPS.filter(message): self.add_delete_message_job(message) try: - async with self.genshin_helper.genshin(user_id) as client: + async with self.genshin_helper.genshin(user_id, player_id=uid, offset=offset) as client: await message.reply_chat_action(ChatAction.TYPING) _, challenge = await self.sign_system.get_challenge(client.player_id) if validate: @@ -132,13 +133,14 @@ class Sign(Plugin): @handler.command(command="start", filters=filters.Regex(r" challenge_sign_(.*)"), block=False) async def command_challenge(self, update: Update, context: CallbackContext) -> None: user = update.effective_user + uid, offset = self.get_real_uid_or_offset(update) message = update.effective_message args = context.args _data = args[0].split("_") validate = _data[2] logger.info("用户 %s[%s] 通过start命令 进入签到流程", user.full_name, user.id) try: - async with self.genshin_helper.genshin(user.id) as client: + async with self.genshin_helper.genshin(user.id, player_id=uid, offset=offset) as client: await message.reply_chat_action(ChatAction.TYPING) _, challenge = await self.sign_system.get_challenge(client.player_id) if not challenge: