PamGram/plugins/admin/migrate.py

from functools import partial
from typing import Dict, List, TYPE_CHECKING

from telegram import InlineKeyboardButton, InlineKeyboardMarkup

from core.plugin import Plugin, handler
from gram_core.plugin.methods.migrate_data import IMigrateData, MigrateDataException
from gram_core.services.players import PlayersService
from utils.log import logger

if TYPE_CHECKING:
    from telegram import Update
    from telegram.ext import ContextTypes


class MigrateAdmin(Plugin):
    def __init__(self, players_service: PlayersService):
        self.players_service = players_service
        self.cache_data: Dict[int, List[IMigrateData]] = {}
        self.wait_time = 60

    async def _add_pop_cache_job(self, user_id: int) -> None:
        if user_id in self.cache_data:
            del self.cache_data[user_id]

    def add_pop_cache_job(self, user_id: int) -> None:
        job_queue = self.application.job_queue
        if job_queue is None:
            raise RuntimeError
        job_queue.run_once(
            callback=partial(self._add_pop_cache_job, user_id=user_id), when=60, name=f"{user_id}|migrate_pop_cache"
        )

    def cancel_pop_cache_job(self, user_id: int) -> None:
        job_queue = self.application.job_queue
        if job_queue is None:
            raise RuntimeError
        if job := job_queue.get_jobs_by_name(f"{user_id}|migrate_pop_cache"):
            job[0].schedule_removal()

    @handler.command(command="migrate_admin", block=False, admin=True)
    async def migrate_admin_command(self, update: "Update", context: "ContextTypes.DEFAULT_TYPE"):
        message = update.effective_message
        args = self.get_args(context)
        logger.info("管理员 %s[%s] migrate_admin 命令请求", message.from_user.full_name, message.from_user.id)
        if (not args) or len(args) < 2:
            await message.reply_text("参数错误，请指定新旧用户 id ！")
            return
        try:
            old_user_id, new_user_id = int(args[0]), int(args[1])
        except ValueError:
            await message.reply_text("参数错误，请指定新旧用户 id ！")
            return
        if old_user_id in self.cache_data:
            await message.reply_text("该用户正在迁移数据中，请稍后再试！")
            return
        data = []
        players = await self.players_service.get_all_by_user_id(old_user_id)
        for _, instance in self.application.managers.plugins_map.items():
            if _data := await instance.get_migrate_data(old_user_id, new_user_id, players):
                data.append(_data)
        if not data:
            await message.reply_text("没有需要迁移的数据！")
            return
        text = "确定迁移以下数据？\n\n"
        for d in data:
            text += f"- {await d.migrate_data_msg()}\n"
        self.cache_data[old_user_id] = data
        buttons = [
            [
                InlineKeyboardButton(
                    "确定迁移",
                    callback_data=f"migrate_admin|{old_user_id}",
                )
            ],
        ]
        await message.reply_text(text, reply_markup=InlineKeyboardMarkup(buttons))
        self.add_pop_cache_job(old_user_id)

    async def try_migrate_data(self, user_id: int) -> str:
        text = []
        for d in self.cache_data[user_id]:
            try:
                logger.info("开始迁移数据 class[%s]", d.__class__.__name__)
                await d.migrate_data()
                logger.info("迁移数据成功 class[%s]", d.__class__.__name__)
            except MigrateDataException as e:
                text.append(e.msg)
            except Exception as e:
                logger.exception("迁移数据失败，未知错误！ class[%s]", d.__class__.__name__, exc_info=e)
                text.append("迁移部分数据出现未知错误，请联系管理员！")
        if text:
            return "- " + "\n- ".join(text)

    @handler.callback_query(pattern=r"^migrate_admin\|", block=False)
    async def callback_query_migrate_admin(self, update: "Update", _: "ContextTypes.DEFAULT_TYPE") -> None:
        callback_query = update.callback_query
        user = callback_query.from_user
        message = callback_query.message
        logger.info("管理员 %s[%s] migrate_admin callback 请求", user.full_name, user.id)

        async def get_migrate_admin_callback(callback_query_data: str) -> int:
            _data = callback_query_data.split("|")
            _old_user_id = int(_data[1])
            logger.debug("callback_query_data函数返回 old_user_id[%s]", _old_user_id)
            return _old_user_id

        old_user_id = await get_migrate_admin_callback(callback_query.data)
        if old_user_id not in self.cache_data:
            await callback_query.answer("请求已过期，请重新发起请求！", show_alert=True)
            self.add_delete_message_job(message, delay=5)
            return
        self.cancel_pop_cache_job(old_user_id)
        await message.edit_text("正在迁移数据，请稍后...", reply_markup=None)
        try:
            text = await self.try_migrate_data(old_user_id)
        finally:
            await self._add_pop_cache_job(old_user_id)
        if text:
            await message.edit_text(f"迁移部分数据失败！\n\n{text}")
            return
        await message.edit_text("迁移数据成功！")