from tortoise import BaseDBAsyncClient, Tortoise, fields from tortoise.exceptions import IntegrityError from zhenxun.configs.config import BotConfig from zhenxun.models.goods_info import GoodsInfo from zhenxun.services.db_context import Model from zhenxun.services.log import logger from zhenxun.utils.enum import CacheType, GoldHandle from zhenxun.utils.exception import GoodsNotFound, InsufficientGold from .user_gold_log import UserGoldLog class UserConsole(Model): id = fields.IntField(pk=True, generated=True, auto_increment=True) """自增id""" user_id = fields.CharField(255, unique=True, description="用户id") """用户id""" uid = fields.IntField(description="UID", unique=True) """UID,用户可修改""" gold = fields.IntField(default=100, description="金币数量") """金币数量""" sign = fields.ReverseRelation["SignUser"] # type: ignore """好感度""" props: dict[str, int] = fields.JSONField(default={}) # type: ignore """道具""" platform = fields.CharField(255, null=True, description="平台") """平台""" create_time = fields.DatetimeField(auto_now_add=True, description="创建时间") """创建时间""" class Meta: # pyright: ignore [reportIncompatibleVariableOverride] table = "user_console" table_description = "用户数据表" indexes = [("user_id",), ("uid",)] # noqa: RUF012 cache_type = CacheType.USERS """缓存类型""" cache_key_field = "user_id" """缓存键字段""" @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 # 使用数据库序列获取 uid,原子操作无竞争 uid = await cls._next_uid_from_sequence() try: return await cls.create(user_id=user_id, uid=uid, platform=platform) except IntegrityError: # user_id 冲突(并发创建同一用户) if user := await cls.get_or_none(user_id=user_id): return user # uid 冲突(极罕见,用户手动修改了 uid),重试 for _ in range(3): try: uid = await cls._next_uid_from_sequence() return await cls.create(user_id=user_id, uid=uid, platform=platform) except IntegrityError: if user := await cls.get_or_none(user_id=user_id): return user raise @classmethod async def _next_uid_from_sequence(cls) -> int: """获取下一个 UID(原子操作,支持 PostgreSQL/MySQL/SQLite)""" conn = Tortoise.get_connection("default") db_type = BotConfig.get_sql_type() try: if db_type == "postgresql": return await cls._next_uid_postgresql(conn) elif db_type == "mysql": return await cls._next_uid_mysql(conn) else: # sqlite return await cls._next_uid_sqlite(conn) except Exception as e: logger.debug(f"序列获取失败,使用备用方案: {e}") return await cls._get_max_uid() + 1 @classmethod async def _next_uid_postgresql(cls, conn: BaseDBAsyncClient) -> int: """PostgreSQL: 使用序列""" result = await conn.execute_query_dict( "SELECT nextval('user_console_uid_seq') as uid" ) return result[0]["uid"] @classmethod async def _next_uid_mysql(cls, conn: BaseDBAsyncClient) -> int: """MySQL: 使用序列表实现原子自增""" # 原子更新并获取新值 await conn.execute_query( """ INSERT INTO user_console_sequence (id, current_value) VALUES (1, 1) ON DUPLICATE KEY UPDATE current_value = current_value + 1 """ ) result = await conn.execute_query_dict( "SELECT current_value as uid FROM user_console_sequence WHERE id = 1" ) return result[0]["uid"] @classmethod async def _next_uid_sqlite(cls, conn: BaseDBAsyncClient) -> int: """SQLite: 使用序列表实现原子自增""" # SQLite 使用 INSERT OR REPLACE 实现原子操作 await conn.execute_query( """ INSERT OR REPLACE INTO user_console_sequence (id, current_value) VALUES (1, COALESCE( (SELECT current_value + 1 FROM user_console_sequence WHERE id = 1), (SELECT COALESCE(MAX(uid), 0) + 1 FROM user_console) )) """ ) result = await conn.execute_query_dict( "SELECT current_value as uid FROM user_console_sequence WHERE id = 1" ) return result[0]["uid"] @classmethod async def _get_max_uid(cls) -> int: """获取当前最大 uid(备用方案)""" data: list[int] = ( # pyright: ignore[reportAssignmentType] await cls.annotate().order_by("-uid").limit(1).values_list("uid", flat=True) ) return data[0] if data else 0 @classmethod async def get_user_count(cls) -> int: """获取用户总数 返回: int: 用户总数 """ return await cls.all().count() @classmethod async def add_gold( cls, user_id: str, gold: int, source: str, platform: str | None = None ): """添加金币 参数: user_id: 用户id gold: 金币 source: 来源 platform: 平台. """ user = await cls.get_user(user_id, platform) user.gold += gold await user.save(update_fields=["gold"]) await UserGoldLog.create( user_id=user_id, gold=gold, handle=GoldHandle.GET, source=source ) @classmethod async def reduce_gold( cls, user_id: str, gold: int, handle: GoldHandle, plugin_module: str, platform: str | None = None, ): """消耗金币 参数: user_id: 用户id gold: 金币 handle: 金币处理 plugin_name: 插件模块 platform: 平台. 异常: InsufficientGold: 金币不足 """ user = await cls.get_user(user_id, platform) if user.gold < gold: raise InsufficientGold() user.gold -= gold await user.save(update_fields=["gold"]) await UserGoldLog.create( user_id=user_id, gold=gold, handle=handle, source=plugin_module ) @classmethod async def add_props( cls, user_id: str, goods_uuid: str, num: int = 1, platform: str | None = None ): """添加道具 参数: user_id: 用户id goods_uuid: 道具uuid num: 道具数量. platform: 平台. """ user = await cls.get_user(user_id, platform) if goods_uuid not in user.props: user.props[goods_uuid] = 0 user.props[goods_uuid] += num await user.save(update_fields=["props"]) @classmethod async def add_props_by_name( cls, user_id: str, name: str, num: int = 1, platform: str | None = None ): """根据名称添加道具 参数: user_id: 用户id name: 道具名称 num: 道具数量. platform: 平台. """ if goods := await GoodsInfo.get_or_none(goods_name=name): return await cls.add_props(user_id, goods.uuid, num, platform) raise GoodsNotFound("未找到商品...") @classmethod async def use_props( cls, user_id: str, goods_uuid: str, num: int = 1, platform: str | None = None ): """添加道具 参数: user_id: 用户id goods_uuid: 道具uuid num: 道具数量. platform: 平台. """ user = await cls.get_user(user_id, platform) if goods_uuid not in user.props or user.props[goods_uuid] < num: raise GoodsNotFound("未找到商品或道具数量不足...") user.props[goods_uuid] -= num if user.props[goods_uuid] <= 0: del user.props[goods_uuid] await user.save(update_fields=["props"]) @classmethod async def use_props_by_name( cls, user_id: str, name: str, num: int = 1, platform: str | None = None ): """根据名称添加道具 参数: user_id: 用户id name: 道具名称 num: 道具数量. platform: 平台. """ if goods := await GoodsInfo.get_or_none(goods_name=name): return await cls.use_props(user_id, goods.uuid, num, platform) raise GoodsNotFound("未找到商品...") @classmethod async def _run_script(cls): """初始化脚本,根据数据库类型创建序列/表""" db_type = BotConfig.get_sql_type() # 通用索引 scripts = [ "CREATE INDEX IF NOT EXISTS idx_user_console_user_id " "ON user_console(user_id);", "CREATE INDEX IF NOT EXISTS idx_user_console_uid ON user_console(uid);", ] # 根据数据库类型添加序列初始化脚本 if db_type == "postgresql": scripts.append( """ DO $$ BEGIN IF NOT EXISTS ( SELECT 1 FROM pg_sequences WHERE schemaname = 'public' AND sequencename = 'user_console_uid_seq' ) THEN CREATE SEQUENCE user_console_uid_seq; PERFORM setval( 'user_console_uid_seq', COALESCE((SELECT MAX(uid) FROM user_console), 0) + 1, false ); END IF; END $$; """ ) elif db_type == "mysql": # MySQL: 创建序列表 scripts.extend( [ """ CREATE TABLE IF NOT EXISTS user_console_sequence ( id INT PRIMARY KEY, current_value BIGINT NOT NULL DEFAULT 0 ); """, """ INSERT IGNORE INTO user_console_sequence (id, current_value) SELECT 1, COALESCE(MAX(uid), 0) FROM user_console; """, ] ) else: # sqlite # SQLite: 创建序列表 scripts.extend( [ """ CREATE TABLE IF NOT EXISTS user_console_sequence ( id INTEGER PRIMARY KEY, current_value INTEGER NOT NULL DEFAULT 0 ); """, """ INSERT OR IGNORE INTO user_console_sequence (id, current_value) SELECT 1, COALESCE(MAX(uid), 0) FROM user_console; """, ] ) return scripts