feat✨: 数据迁移

zhenxun/builtin_plugins/__init__.py

@@ -1,6 +1,13 @@
 import os
 
 from nonebot import require
+from tortoise import Tortoise
+
+from zhenxun.models.goods_info import GoodsInfo
+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
 
 require("nonebot_plugin_apscheduler")
 require("nonebot_plugin_alconna")
@@ -13,7 +20,105 @@ enable_auto_select_bot()
 from pathlib import Path
 
 import nonebot
+import ujson as json
 
 path = Path(__file__).parent / "platform"
 for d in os.listdir(path):
     nonebot.load_plugins(str((path / d).resolve()))
+
+
+driver = nonebot.get_driver()
+
+flag = True
+
+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_bot_connect
+async def _():
+    global flag
+    await shop_register.load_register()
+    if (
+        flag
+        and not await UserConsole.annotate().count()
+        and not await SignUser.annotate().count()
+    ):
+        flag = False
+        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 = []
+        uid = await UserConsole.get_new_uid()
+        for old_sign in old_sign_list:
+            sign_id_list.append(old_sign["user_id"])
+            old_bag = [b for b in old_bag_list if b["user_id"] == old_sign["user_id"]]
+            if old_bag:
+                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=uid,
+                        props=props,
+                        gold=old_bag["gold"],
+                    )
+                )
+            else:
+                create_list.append(
+                    UserConsole(user_id=old_sign["user_id"], platform="qq", uid=uid)
+                )
+            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"])
+            if not user_console:
+                user_console = 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("迁移签到数据完成!")

zhenxun/builtin_plugins/hooks/__init__.py

@@ -40,4 +40,4 @@ Config.add_plugin_config(
     type=int,
 )
 
-nonebot.load_plugins(str(Path(__file__).parent.resolve()))
+# nonebot.load_plugins(str(Path(__file__).parent.resolve()))

zhenxun/builtin_plugins/init/init_plugin.py

@@ -39,6 +39,7 @@ async def _handle_setting(
         setting = extra_data.setting or PluginSetting()
         if metadata.type == "library":
             extra_data.plugin_type = PluginType.HIDDEN
+            extra_data.menu_type = ""
         plugin_list.append(
             PluginInfo(
                 module=plugin.name,

zhenxun/builtin_plugins/scripts.py

@@ -1,7 +1,6 @@
 from asyncio.exceptions import TimeoutError
 
 import nonebot
-import ujson as json
 from nonebot.drivers import Driver
 from nonebot_plugin_apscheduler import scheduler
 
@@ -9,6 +8,12 @@ from zhenxun.configs.path_config import TEXT_PATH
 from zhenxun.services.log import logger
 from zhenxun.utils.http_utils import AsyncHttpx
 
+try:
+    import ujson as json
+except ModuleNotFoundError:
+    import json
+
+
 driver: Driver = nonebot.get_driver()
 
 
@@ -40,13 +45,13 @@ async def update_city():
                     data[provinces_data[province]].append(city_data[city])
             with open(china_city, "w", encoding="utf8") as f:
                 json.dump(data, f, indent=4, ensure_ascii=False)
-            logger.info("自动更新城市列表完成...")
+            logger.info("自动更新城市列表完成.....")
         except TimeoutError as e:
             logger.warning("自动更新城市列表超时...", e=e)
         except ValueError as e:
-            logger.warning("自动城市列表失败...", e=e)
+            logger.warning("自动城市列表失败.....", e=e)
         except Exception as e:
-            logger.error(f"自动城市列表未知错误...", e=e)
+            logger.error(f"自动城市列表未知错误", e=e)
 
 
 # 自动更新城市列表

zhenxun/builtin_plugins/shop/_data_source.py

@@ -1,5 +1,8 @@
 import time
-from typing import Dict
+from typing import Any, Callable, Dict
+
+from nonebot.adapters import Event
+from pydantic import BaseModel, create_model
 
 from zhenxun.configs.path_config import IMAGE_PATH
 from zhenxun.models.goods_info import GoodsInfo
@@ -13,8 +16,80 @@ from zhenxun.utils.image_utils import BuildImage, ImageTemplate, text2image
 ICON_PATH = IMAGE_PATH / "shop_icon"
 
 
+class Goods(BaseModel):
+
+    before_handle: list[Callable] = []
+    after_handle: list[Callable] = []
+    func: Callable | None = None
+    params: Any | None = None
+    send_success_msg: bool = True
+    max_num_limit: int = 1
+    model: Any | None = None
+
+
+class ShopParam(BaseModel):
+
+    goods_name: str
+    """商品名称"""
+    user_id: int
+    """用户id"""
+    group_id: int
+    """群聊id"""
+    bot: Any
+    """bot"""
+    event: Event
+    """event"""
+    num: int
+    """道具单次使用数量"""
+    message: str
+    """message"""
+    text: str
+    """text"""
+    send_success_msg: bool = True
+    """是否发送使用成功信息"""
+    max_num_limit: int = 1
+    """单次使用最大次数"""
+
+
 class ShopManage:
 
+    uuid2goods: Dict[str, Goods] = {}
+
+    @classmethod
+    async def register_use(
+        cls,
+        uuid: str,
+        func: Callable,
+        send_success_msg: bool = True,
+        max_num_limit: int = 1,
+        before_handle: list[Callable] = [],
+        after_handle: list[Callable] = [],
+        **kwargs,
+    ):
+        """注册使用方法
+
+        参数:
+            uuid: uuid
+            func: 使用函数
+            send_success_msg: 使用成功时发送消息.
+            max_num_limit: 单次最大使用限制.
+            before_handle: 使用前函数.
+            after_handle: 使用后函数.
+
+        异常:
+            ValueError: 该商品使用函数已被注册！
+        """
+        if uuid in cls.uuid2goods:
+            raise ValueError("该商品使用函数已被注册！")
+        kwargs["send_success_msg"] = send_success_msg
+        kwargs["max_num_limit"] = max_num_limit
+        cls.uuid2func = Goods(
+            model=create_model(f"{uuid}_model", __base__=ShopParam, **kwargs),
+            params=kwargs,
+            before_handle=before_handle,
+            after_handle=after_handle,
+        )
+
     @classmethod
     async def buy_prop(
         cls, user_id: str, name: str, num: int = 1, platform: str | None = None