mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-09-29 00:32:06 +08:00
329 lines
11 KiB
Python
329 lines
11 KiB
Python
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
|