zhenxun_bot/zhenxun/builtin_plugins/__init__.py

from datetime import datetime
import uuid

import nonebot
from nonebot import require
from nonebot.adapters import Bot
from nonebot.drivers import Driver
from tortoise import Tortoise
from tortoise.exceptions import OperationalError
import ujson as json

require("nonebot_plugin_apscheduler")
require("nonebot_plugin_alconna")
require("nonebot_plugin_session")
require("nonebot_plugin_userinfo")
require("nonebot_plugin_htmlrender")
# require("nonebot_plugin_uninfo")

from zhenxun.models.bot_connect_log import BotConnectLog
from zhenxun.models.bot_console import BotConsole
from zhenxun.models.goods_info import GoodsInfo
from zhenxun.models.group_member_info import GroupInfoUser
from zhenxun.models.sign_user import SignUser
from zhenxun.models.user_console import UserConsole
from zhenxun.services.log import logger
from zhenxun.utils.decorator.shop import shop_register
from zhenxun.utils.platform import PlatformUtils

driver: Driver = nonebot.get_driver()


@driver.on_bot_connect
async def _(bot: Bot):
    logger.debug(f"Bot: {bot.self_id} 建立连接...")
    await BotConnectLog.create(
        bot_id=bot.self_id, platform=bot.adapter, connect_time=datetime.now(), type=1
    )
    if not await BotConsole.exists(bot_id=bot.self_id):
        await BotConsole.create(
            bot_id=bot.self_id, platform=PlatformUtils.get_platform(bot)
        )


@driver.on_bot_disconnect
async def _(bot: Bot):
    logger.debug(f"Bot: {bot.self_id} 断开连接...")
    await BotConnectLog.create(
        bot_id=bot.self_id, platform=bot.adapter, connect_time=datetime.now(), type=0
    )


SIGN_SQL = """
select distinct on("user_id") t1.user_id, t1.checkin_count, t1.add_probability,
t1.specify_probability, t1.impression
from public.sign_group_users t1
  join (
    select user_id, max(t2.impression) as max_impression
    from public.sign_group_users t2
    group by user_id
  ) t on t.user_id = t1.user_id and t.max_impression = t1.impression
"""

BAG_SQL = """
select t1.user_id, t1.gold, t1.property
from public.bag_users t1
  join (
    select user_id, max(t2.gold) as max_gold
    from public.bag_users t2
    group by user_id
  ) t on t.user_id = t1.user_id and t.max_gold = t1.gold
"""


@driver.on_startup
async def _():
    """签到与用户的数据迁移"""
    if goods_list := await GoodsInfo.filter(uuid__isnull=True).all():
        for goods in goods_list:
            goods.uuid = uuid.uuid1()  # type: ignore
        await GoodsInfo.bulk_update(goods_list, ["uuid"], 10)
    await shop_register.load_register()
    if (
        not await UserConsole.annotate().count()
        and not await SignUser.annotate().count()
    ):
        try:
            group_user = []
            try:
                group_user = await GroupInfoUser.filter(uid__isnull=False).all()
            except Exception as e:
                logger.warning("获取GroupInfoUser数据uid失败...", e=e)
            user2uid = {u.user_id: u.uid for u in group_user}
            db = Tortoise.get_connection("default")
            old_sign_list = await db.execute_query_dict(SIGN_SQL)
            old_bag_list = await db.execute_query_dict(BAG_SQL)
            goods = {
                g["goods_name"]: g["uuid"]
                for g in await GoodsInfo.annotate().values("goods_name", "uuid")
            }
            create_list = []
            sign_id_list = []
            max_uid = max(user2uid.values()) + 1 if user2uid else 0
            for old_sign in old_sign_list:
                sign_id_list.append(old_sign["user_id"])
                if old_bag := [
                    b for b in old_bag_list if b["user_id"] == old_sign["user_id"]
                ]:
                    old_bag = old_bag[0]
                    property = json.loads(old_bag["property"])
                    props = {}
                    if property:
                        for name, num in property.items():
                            if name in goods:
                                props[goods[name]] = num
                    create_list.append(
                        UserConsole(
                            user_id=old_sign["user_id"],
                            platform="qq",
                            uid=user2uid.get(old_sign["user_id"]) or max_uid,
                            props=props,
                            gold=old_bag["gold"],
                        )
                    )
                    if not user2uid.get(old_sign["user_id"]):
                        max_uid += 1
                else:
                    create_list.append(
                        UserConsole(
                            user_id=old_sign["user_id"], platform="qq", uid=max_uid
                        )
                    )
                    max_uid += 1
            if create_list:
                logger.info("开始迁移用户数据...")
                await UserConsole.bulk_create(create_list, 10)
                logger.info("迁移用户数据完成!")
            create_list.clear()
            uc_dict = {u.user_id: u for u in await UserConsole.all()}
            for old_sign in old_sign_list:
                user_console = uc_dict.get(
                    old_sign["user_id"]
                ) or await UserConsole.get_user(old_sign["user_id"], "qq")
                create_list.append(
                    SignUser(
                        user_id=old_sign["user_id"],
                        user_console=user_console,
                        platform="qq",
                        sign_count=old_sign["checkin_count"],
                        impression=old_sign["impression"],
                        add_probability=old_sign["add_probability"],
                        specify_probability=old_sign["specify_probability"],
                    )
                )
            if create_list:
                logger.info("开始迁移签到数据...")
                await SignUser.bulk_create(create_list, 10)
                logger.info("迁移签到数据完成!")
        except OperationalError as e:
            logger.warning("数据迁移", e=e)