mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-10-08 05:10:01 +08:00
refactor: enhance UserConsole UID management and remove mute plugin
This commit is contained in:
HibiKier
+156
-28
@@ -1,12 +1,11 @@
|
||||
from typing import ClassVar
|
||||
|
||||
from tortoise import fields
|
||||
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.cache import CacheRoot
|
||||
from zhenxun.services.db_context import Model
|
||||
from zhenxun.utils.enum import CacheType, DbLockType, GoldHandle
|
||||
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
|
||||
@@ -18,7 +17,7 @@ class UserConsole(Model):
|
||||
user_id = fields.CharField(255, unique=True, description="用户id")
|
||||
"""用户id"""
|
||||
uid = fields.IntField(description="UID", unique=True)
|
||||
"""UID"""
|
||||
"""UID，用户可修改"""
|
||||
gold = fields.IntField(default=100, description="金币数量")
|
||||
"""金币数量"""
|
||||
sign = fields.ReverseRelation["SignUser"] # type: ignore
|
||||
@@ -39,39 +38,107 @@ class UserConsole(Model):
|
||||
"""缓存类型"""
|
||||
cache_key_field = "user_id"
|
||||
"""缓存键字段"""
|
||||
lock_fields: ClassVar[dict[DbLockType, str]] = {DbLockType.CREATE: "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:
|
||||
uid = await cls.get_new_uid()
|
||||
return await cls.create(user_id=user_id, uid=uid, platform=platform)
|
||||
except IntegrityError:
|
||||
# 并发竞争下，依赖 user_id 唯一约束兜底查询
|
||||
return await cls.get(user_id=user_id)
|
||||
# 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 get_new_uid(cls) -> int:
|
||||
"""获取最新uid
|
||||
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: 最新uid
|
||||
int: 用户总数
|
||||
"""
|
||||
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 or 1
|
||||
return await cls.all().count()
|
||||
|
||||
@classmethod
|
||||
async def add_gold(
|
||||
@@ -194,7 +261,68 @@ class UserConsole(Model):
|
||||
|
||||
@classmethod
|
||||
async def _run_script(cls):
|
||||
return [
|
||||
"CREATE INDEX idx_user_console_user_id ON user_console(user_id);",
|
||||
"CREATE INDEX idx_user_console_uid ON user_console(uid);",
|
||||
"""初始化脚本，根据数据库类型创建序列/表"""
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user