From a3cbfefaa1b618c49cbe3e8d56069b43e7bc9407 Mon Sep 17 00:00:00 2001 From: HibiKier <775757368@qq.com> Date: Mon, 22 Dec 2025 10:16:02 +0800 Subject: [PATCH] feat: add TEMP cache type and enhance user and bot management with improved error handling and caching mechanisms --- zhenxun/builtin_plugins/init/__init_cache.py | 2 +- .../superuser/bot_manage/__init__.py | 15 +++- zhenxun/models/ban_console.py | 80 +++++++++--------- zhenxun/models/user_console.py | 82 ++++++++++++------- zhenxun/services/cache/__init__.py | 2 +- zhenxun/utils/enum.py | 2 + 6 files changed, 108 insertions(+), 75 deletions(-) diff --git a/zhenxun/builtin_plugins/init/__init_cache.py b/zhenxun/builtin_plugins/init/__init_cache.py index 8938c0a0..0d1eed5d 100644 --- a/zhenxun/builtin_plugins/init/__init_cache.py +++ b/zhenxun/builtin_plugins/init/__init_cache.py @@ -33,7 +33,7 @@ def register_cache_types(): CacheType.LEVEL, LevelUser, key_format="{user_id}_{group_id}" ) CacheRegistry.register(CacheType.BAN, BanConsole, key_format="{user_id}_{group_id}") - + CacheRegistry.register(CacheType.TEMP, None, 3600) if cache_config.cache_mode == CacheMode.NONE: logger.info("缓存功能已禁用，将直接从数据库获取数据") else: diff --git a/zhenxun/builtin_plugins/superuser/bot_manage/__init__.py b/zhenxun/builtin_plugins/superuser/bot_manage/__init__.py index f7923b53..8bbdb744 100644 --- a/zhenxun/builtin_plugins/superuser/bot_manage/__init__.py +++ b/zhenxun/builtin_plugins/superuser/bot_manage/__init__.py @@ -3,6 +3,7 @@ from typing import cast import nonebot from nonebot.adapters import Bot from nonebot.plugin import PluginMetadata +from tortoise.exceptions import IntegrityError from zhenxun.configs.utils import PluginExtraData from zhenxun.models.bot_console import BotConsole @@ -72,9 +73,17 @@ async def init_bot_console(bot: Bot): list[str], await TaskInfo.filter(status=True).values_list("module", flat=True) ) platform = PlatformUtils.get_platform(bot) - bot_data, created = await BotConsole.get_or_create( - bot_id=bot.self_id, platform=platform - ) + + try: + bot_data = await BotConsole.create( + bot_id=bot.self_id, + platform=platform, + ) + created = True + + except IntegrityError: + bot_data = await BotConsole.get(bot_id=bot.self_id) + created = False if not created: task_list = await _filter_blocked_items( diff --git a/zhenxun/models/ban_console.py b/zhenxun/models/ban_console.py index 4fec9608..90127f45 100644 --- a/zhenxun/models/ban_console.py +++ b/zhenxun/models/ban_console.py @@ -1,3 +1,4 @@ +import asyncio import time from typing import ClassVar from typing_extensions import Self @@ -28,6 +29,7 @@ class BanConsole(Model): """ban时长""" operator = fields.CharField(255) """使用Ban命令的用户""" + _inflight: ClassVar[dict[tuple[str | None, str | None], asyncio.Future]] = {} class Meta: # pyright: ignore [reportIncompatibleVariableOverride] table = "ban_console" @@ -44,29 +46,38 @@ class BanConsole(Model): @classmethod async def _get_data(cls, user_id: str | None, group_id: str | None) -> Self | None: - """获取数据 - - 参数: - user_id: 用户id - group_id: 群组id - - 异常: - UserAndGroupIsNone: 用户id和群组id都为空 - - 返回: - Self | None: Self - """ if not user_id and not group_id: raise UserAndGroupIsNone() - dao = DataAccess(cls) - if user_id: - return ( - await dao.safe_get_or_none(user_id=user_id, group_id=group_id) - if group_id - else await dao.safe_get_or_none(user_id=user_id, group_id__isnull=True) - ) - else: - return await dao.safe_get_or_none(user_id="", group_id=group_id) + + key = (user_id, group_id) + future = cls._inflight.get(key) + if future: + return await future + + loop = asyncio.get_running_loop() + future = loop.create_future() + cls._inflight[key] = future + + try: + dao = DataAccess(cls) + if user_id: + result = ( + await dao.safe_get_or_none(user_id=user_id, group_id=group_id) + if group_id + else await dao.safe_get_or_none( + user_id=user_id, group_id__isnull=True + ) + ) + else: + result = await dao.safe_get_or_none(user_id="", group_id=group_id) + + future.set_result(result) + return result + except Exception as e: + future.set_exception(e) + raise + finally: + cls._inflight.pop(key, None) @classmethod async def check_ban_level( @@ -143,30 +154,21 @@ class BanConsole(Model): duration: int, operator: str | None = None, ): - """ban掉目标用户 - - 参数: - user_id: 用户id - group_id: 群组id - ban_level: 使用命令者的权限等级 - duration: 时长，分钟，-1时为永久 - operator: 操作者id - """ logger.debug( f"封禁用户/群组，等级:{ban_level}，时长: {duration}", target=f"{group_id}:{user_id}", ) - target = await cls._get_data(user_id, group_id) - if target: - await cls.unban(user_id, group_id) - await cls.create( + + await cls.update_or_create( user_id=user_id, group_id=group_id, - ban_level=ban_level, - ban_time=int(time.time()), - ban_reason=reason, - duration=duration, - operator=operator or 0, + defaults={ + "ban_level": ban_level, + "ban_time": int(time.time()), + "ban_reason": reason, + "duration": duration, + "operator": operator or 0, + }, ) @classmethod diff --git a/zhenxun/models/user_console.py b/zhenxun/models/user_console.py index 27cd582f..d90e9c0e 100644 --- a/zhenxun/models/user_console.py +++ b/zhenxun/models/user_console.py @@ -1,6 +1,8 @@ from tortoise import fields +from tortoise.exceptions import IntegrityError from zhenxun.models.goods_info import GoodsInfo +from zhenxun.services.cache import CacheRoot from zhenxun.services.db_context import Model from zhenxun.utils.enum import CacheType, GoldHandle from zhenxun.utils.exception import GoodsNotFound, InsufficientGold @@ -38,24 +40,23 @@ class UserConsole(Model): @classmethod async def get_user(cls, user_id: str, platform: str | None = None) -> "UserConsole": - """获取用户 + if user := await cls.get_or_none(user_id=user_id): + return user - 参数: - user_id: 用户id - platform: 平台. - - 返回: - UserConsole: UserConsole - """ - if not await cls.exists(user_id=user_id): - await cls.create( - user_id=user_id, platform=platform, uid=await cls.get_new_uid() + try: + user, created = await cls.get_or_create( + user_id=user_id, + defaults={ + "platform": platform, + "uid": 0, + }, ) - # user, _ = await UserConsole.get_or_create( - # user_id=user_id, - # defaults={"platform": platform, "uid": await cls.get_new_uid()}, - # ) - return await cls.get(user_id=user_id) + if created: + user.uid = await cls.get_new_uid() + await user.save(update_fields=["uid"]) + return user + except IntegrityError: + return await cls.get(user_id=user_id) @classmethod async def get_new_uid(cls) -> int: @@ -64,9 +65,18 @@ class UserConsole(Model): 返回: int: 最新uid """ - if user := await cls.annotate().order_by("-uid").first(): - return user.uid + 1 - return 1 + uid: int | None = await CacheRoot.get(CacheType.TEMP, "USER_CONSOLE_UID") + if uid is None: + data: list[int] = ( # pyright: ignore[reportAssignmentType] + await cls.annotate() + .order_by("-uid") + .limit(1) + .values_list("uid", flat=True) + ) + uid = data[0] if data else 0 + uid = uid + 1 + await CacheRoot.set(CacheType.TEMP, "USER_CONSOLE_UID", uid) + return uid @classmethod async def add_gold( @@ -80,12 +90,14 @@ class UserConsole(Model): source: 来源 platform: 平台. """ - user, _ = await cls.get_or_create( + user, created = await cls.get_or_create( user_id=user_id, - defaults={"platform": platform, "uid": await cls.get_new_uid()}, + defaults={"platform": platform, "uid": 0}, ) + if created: + user.uid = await cls.get_new_uid() user.gold += gold - await user.save(update_fields=["gold"]) + await user.save(update_fields=["gold", "uid"]) await UserGoldLog.create( user_id=user_id, gold=gold, handle=GoldHandle.GET, source=source ) @@ -111,14 +123,16 @@ class UserConsole(Model): 异常: InsufficientGold: 金币不足 """ - user, _ = await cls.get_or_create( + user, created = await cls.get_or_create( user_id=user_id, - defaults={"platform": platform, "uid": await cls.get_new_uid()}, + defaults={"platform": platform, "uid": 0}, ) + if created: + user.uid = await cls.get_new_uid() if user.gold < gold: raise InsufficientGold() user.gold -= gold - await user.save(update_fields=["gold"]) + await user.save(update_fields=["gold", "uid"]) await UserGoldLog.create( user_id=user_id, gold=gold, handle=handle, source=plugin_module ) @@ -135,14 +149,16 @@ class UserConsole(Model): num: 道具数量. platform: 平台. """ - user, _ = await cls.get_or_create( + user, created = await cls.get_or_create( user_id=user_id, - defaults={"platform": platform, "uid": await cls.get_new_uid()}, + defaults={"platform": platform, "uid": 0}, ) + if created: + user.uid = await cls.get_new_uid() if goods_uuid not in user.props: user.props[goods_uuid] = 0 user.props[goods_uuid] += num - await user.save(update_fields=["props"]) + await user.save(update_fields=["props", "uid"]) @classmethod async def add_props_by_name( @@ -172,17 +188,21 @@ class UserConsole(Model): num: 道具数量. platform: 平台. """ - user, _ = await cls.get_or_create( + user, created = await cls.get_or_create( user_id=user_id, - defaults={"platform": platform, "uid": await cls.get_new_uid()}, + defaults={"platform": platform, "uid": 0}, ) + if created: + user.uid = await cls.get_new_uid() if goods_uuid not in user.props or user.props[goods_uuid] < num: + if created: + await user.save(update_fields=["uid"]) raise GoodsNotFound("未找到商品或道具数量不足...") user.props[goods_uuid] -= num if user.props[goods_uuid] <= 0: del user.props[goods_uuid] - await user.save(update_fields=["props"]) + await user.save(update_fields=["props", "uid"]) @classmethod async def use_props_by_name( diff --git a/zhenxun/services/cache/__init__.py b/zhenxun/services/cache/__init__.py index aba34539..ca448cc1 100644 --- a/zhenxun/services/cache/__init__.py +++ b/zhenxun/services/cache/__init__.py @@ -708,7 +708,7 @@ class CacheManager: if self._cache_backend: try: await self._cache_backend.close() # type: ignore - except (AttributeError, Exception) as e: + except Exception as e: logger.warning(f"关闭缓存连接失败: {e}", LOG_COMMAND) self._cache_backend = None diff --git a/zhenxun/utils/enum.py b/zhenxun/utils/enum.py index 0412ee4a..c8e1023e 100644 --- a/zhenxun/utils/enum.py +++ b/zhenxun/utils/enum.py @@ -65,6 +65,8 @@ class CacheType(StrEnum): """用户权限""" LIMIT = "GLOBAL_LIMIT" """插件限制""" + TEMP = "TEMP" + """临时缓存""" class DbLockType(StrEnum):