mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-10-12 23:40:03 +08:00
✨ 引入缓存机制 (#1889)
* 添加全局cache * ✨ 构建缓存，hook使用缓存 * ✨ 新增数据库Model方法监控 * ✨ 数据库添加semaphore锁 * 🩹 优化webapi返回数据 * ✨ 添加增量缓存与缓存过期 * 🎨 优化检测代码结构 * ⚡ 优化hook权限检测性能 * 🐛 添加新异常判断跳过权限检测 * ✨ 添加插件limit缓存 * 🎨 代码格式优化 * 🐛 修复代码导入 * 🐛 修复刷新时检查 * 👽 Rename exception for missing database URL in initialization * ♿ Update default database URL to SQLite in configuration * 🔧 Update tortoise-orm and aiocache dependencies restrictions; add optional redis and asyncpg support * 🐛 修复ban检测 * 🐛 修复所有插件关闭时缓存更新 * 🐛 尝试迁移至aiocache * 🐛 完善aiocache缓存 * ⚡ 代码性能优化 * 🐛 移除获取封禁缓存时的日志记录 * 🐛 修复缓存类型声明，优化封禁用户处理逻辑 * 🐛 优化LevelUser权限更新逻辑及数据库迁移 * ✨ cache支持redis连接 * 🚨 auto fix by pre-commit hooks * ⚡ ：增强获取群组的安全性和准确性。同时，优化了缓存管理中的相关逻辑，确保缓存操作的一致性。 * ✨ feat(auth_limit): 将插件初始化逻辑的启动装饰器更改为优先级管理器 * 🔧 修复日志记录级别 * 🔧 更新数据库连接字符串 * 🔧 更新数据库连接字符串为内存数据库，并优化权限检查逻辑 * ✨ feat(cache): 增加缓存功能配置项，并新增数据访问层以支持缓存逻辑 * ♻️ 重构cache * ✨ feat(cache): 增强缓存管理，新增缓存字典和缓存列表功能，支持过期时间管理 * 🔧 修复Notebook类中的viewport高度设置，将其从1000调整为10 * ✨ 更新插件管理逻辑，替换缓存服务为CacheRoot并优化缓存失效处理 * ✨ 更新RegisterConfig类中的type字段 * ✨ 修复清理重复记录逻辑，确保检查记录的id属性有效性 * ⚡ 超级无敌大优化，解决延迟与卡死问题 * ✨ 更新封禁功能，增加封禁时长参数和描述，优化插件信息返回结构 * ✨ 更新zhenxun_help.py中的viewport高度，将其从453调整为10，以优化页面显示效果 * ✨ 优化插件分类逻辑，增加插件ID排序，并更新插件信息返回结构 --------- Co-authored-by: BalconyJH <balconyjh@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
HibiKier
co-authored by
BalconyJH
pre-commit-ci[bot]
BalconyJH
pre-commit-ci[bot]
parent
6283c3d13d
commit
8649aaaa54
Vendored
+1065
File diff suppressed because it is too large
Load Diff
+452
@@ -0,0 +1,452 @@
|
||||
from dataclasses import dataclass
|
||||
import time
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheData(Generic[T]):
|
||||
"""缓存数据类，存储数据和过期时间"""
|
||||
|
||||
value: T
|
||||
expire_time: float = 0 # 0表示永不过期
|
||||
|
||||
|
||||
class CacheDict:
|
||||
"""缓存字典类，提供类似普通字典的接口，数据只存储在内存中"""
|
||||
|
||||
def __init__(self, name: str, expire: int = 0):
|
||||
"""初始化缓存字典
|
||||
|
||||
参数:
|
||||
name: 字典名称
|
||||
expire: 过期时间（秒），默认为0表示永不过期
|
||||
"""
|
||||
self.name = name.upper()
|
||||
self.expire = expire
|
||||
self._data: dict[str, CacheData[Any]] = {}
|
||||
|
||||
def __getitem__(self, key: str) -> Any:
|
||||
"""获取字典项
|
||||
|
||||
参数:
|
||||
key: 字典键
|
||||
|
||||
返回:
|
||||
Any: 字典值
|
||||
"""
|
||||
data = self._data.get(key)
|
||||
if data is None:
|
||||
return None
|
||||
|
||||
# 检查是否过期
|
||||
if data.expire_time > 0 and data.expire_time < time.time():
|
||||
del self._data[key]
|
||||
return None
|
||||
|
||||
return data.value
|
||||
|
||||
def __setitem__(self, key: str, value: Any) -> None:
|
||||
"""设置字典项
|
||||
|
||||
参数:
|
||||
key: 字典键
|
||||
value: 字典值
|
||||
"""
|
||||
# 计算过期时间
|
||||
expire_time = 0
|
||||
if self.expire > 0:
|
||||
expire_time = time.time() + self.expire
|
||||
|
||||
self._data[key] = CacheData(value=value, expire_time=expire_time)
|
||||
|
||||
def __delitem__(self, key: str) -> None:
|
||||
"""删除字典项
|
||||
|
||||
参数:
|
||||
key: 字典键
|
||||
"""
|
||||
if key in self._data:
|
||||
del self._data[key]
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
"""检查键是否存在
|
||||
|
||||
参数:
|
||||
key: 字典键
|
||||
|
||||
返回:
|
||||
bool: 是否存在
|
||||
"""
|
||||
if key not in self._data:
|
||||
return False
|
||||
|
||||
# 检查是否过期
|
||||
data = self._data[key]
|
||||
if data.expire_time > 0 and data.expire_time < time.time():
|
||||
del self._data[key]
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
"""获取字典项，如果不存在返回默认值
|
||||
|
||||
参数:
|
||||
key: 字典键
|
||||
default: 默认值
|
||||
|
||||
返回:
|
||||
Any: 字典值或默认值
|
||||
"""
|
||||
value = self[key]
|
||||
return default if value is None else value
|
||||
|
||||
def set(self, key: str, value: Any, expire: int | None = None) -> None:
|
||||
"""设置字典项
|
||||
|
||||
参数:
|
||||
key: 字典键
|
||||
value: 字典值
|
||||
expire: 过期时间（秒），为None时使用默认值
|
||||
"""
|
||||
# 计算过期时间
|
||||
expire_time = 0
|
||||
if expire is not None and expire > 0:
|
||||
expire_time = time.time() + expire
|
||||
elif self.expire > 0:
|
||||
expire_time = time.time() + self.expire
|
||||
|
||||
self._data[key] = CacheData(value=value, expire_time=expire_time)
|
||||
|
||||
def pop(self, key: str, default: Any = None) -> Any:
|
||||
"""删除并返回字典项
|
||||
|
||||
参数:
|
||||
key: 字典键
|
||||
default: 默认值
|
||||
|
||||
返回:
|
||||
Any: 字典值或默认值
|
||||
"""
|
||||
if key not in self._data:
|
||||
return default
|
||||
|
||||
data = self._data.pop(key)
|
||||
|
||||
# 检查是否过期
|
||||
if data.expire_time > 0 and data.expire_time < time.time():
|
||||
return default
|
||||
|
||||
return data.value
|
||||
|
||||
def clear(self) -> None:
|
||||
"""清空字典"""
|
||||
self._data.clear()
|
||||
|
||||
def keys(self) -> list[str]:
|
||||
"""获取所有键
|
||||
|
||||
返回:
|
||||
list[str]: 键列表
|
||||
"""
|
||||
# 清理过期的键
|
||||
self._clean_expired()
|
||||
return list(self._data.keys())
|
||||
|
||||
def values(self) -> list[Any]:
|
||||
"""获取所有值
|
||||
|
||||
返回:
|
||||
list[Any]: 值列表
|
||||
"""
|
||||
# 清理过期的键
|
||||
self._clean_expired()
|
||||
return [data.value for data in self._data.values()]
|
||||
|
||||
def items(self) -> list[tuple[str, Any]]:
|
||||
"""获取所有键值对
|
||||
|
||||
返回:
|
||||
list[tuple[str, Any]]: 键值对列表
|
||||
"""
|
||||
# 清理过期的键
|
||||
self._clean_expired()
|
||||
return [(key, data.value) for key, data in self._data.items()]
|
||||
|
||||
def _clean_expired(self) -> None:
|
||||
"""清理过期的键"""
|
||||
now = time.time()
|
||||
expired_keys = [
|
||||
key
|
||||
for key, data in self._data.items()
|
||||
if data.expire_time > 0 and data.expire_time < now
|
||||
]
|
||||
for key in expired_keys:
|
||||
del self._data[key]
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""获取字典长度
|
||||
|
||||
返回:
|
||||
int: 字典长度
|
||||
"""
|
||||
# 清理过期的键
|
||||
self._clean_expired()
|
||||
return len(self._data)
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""字符串表示
|
||||
|
||||
返回:
|
||||
str: 字符串表示
|
||||
"""
|
||||
# 清理过期的键
|
||||
self._clean_expired()
|
||||
return f"CacheDict({self.name}, {len(self._data)} items)"
|
||||
|
||||
|
||||
class CacheList:
|
||||
"""缓存列表类，提供类似普通列表的接口，数据只存储在内存中"""
|
||||
|
||||
def __init__(self, name: str, expire: int = 0):
|
||||
"""初始化缓存列表
|
||||
|
||||
参数:
|
||||
name: 列表名称
|
||||
expire: 过期时间（秒），默认为0表示永不过期
|
||||
"""
|
||||
self.name = name.upper()
|
||||
self.expire = expire
|
||||
self._data: list[CacheData[Any]] = []
|
||||
self._expire_time = 0
|
||||
|
||||
# 如果设置了过期时间，计算整个列表的过期时间
|
||||
if self.expire > 0:
|
||||
self._expire_time = time.time() + self.expire
|
||||
|
||||
def __getitem__(self, index: int) -> Any:
|
||||
"""获取列表项
|
||||
|
||||
参数:
|
||||
index: 列表索引
|
||||
|
||||
返回:
|
||||
Any: 列表值
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
raise IndexError(f"列表索引 {index} 超出范围")
|
||||
|
||||
if 0 <= index < len(self._data):
|
||||
return self._data[index].value
|
||||
raise IndexError(f"列表索引 {index} 超出范围")
|
||||
|
||||
def __setitem__(self, index: int, value: Any) -> None:
|
||||
"""设置列表项
|
||||
|
||||
参数:
|
||||
index: 列表索引
|
||||
value: 列表值
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
|
||||
# 确保索引有效
|
||||
while len(self._data) <= index:
|
||||
self._data.append(CacheData(value=None))
|
||||
self._data[index] = CacheData(value=value)
|
||||
|
||||
# 更新过期时间
|
||||
self._update_expire_time()
|
||||
|
||||
def __delitem__(self, index: int) -> None:
|
||||
"""删除列表项
|
||||
|
||||
参数:
|
||||
index: 列表索引
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
raise IndexError(f"列表索引 {index} 超出范围")
|
||||
|
||||
if 0 <= index < len(self._data):
|
||||
del self._data[index]
|
||||
# 更新过期时间
|
||||
self._update_expire_time()
|
||||
else:
|
||||
raise IndexError(f"列表索引 {index} 超出范围")
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""获取列表长度
|
||||
|
||||
返回:
|
||||
int: 列表长度
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
return len(self._data)
|
||||
|
||||
def append(self, value: Any) -> None:
|
||||
"""添加列表项
|
||||
|
||||
参数:
|
||||
value: 列表值
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
|
||||
self._data.append(CacheData(value=value))
|
||||
|
||||
# 更新过期时间
|
||||
self._update_expire_time()
|
||||
|
||||
def extend(self, values: list[Any]) -> None:
|
||||
"""扩展列表
|
||||
|
||||
参数:
|
||||
values: 要添加的值列表
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
|
||||
self._data.extend([CacheData(value=v) for v in values])
|
||||
|
||||
# 更新过期时间
|
||||
self._update_expire_time()
|
||||
|
||||
def insert(self, index: int, value: Any) -> None:
|
||||
"""插入列表项
|
||||
|
||||
参数:
|
||||
index: 插入位置
|
||||
value: 列表值
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
|
||||
self._data.insert(index, CacheData(value=value))
|
||||
|
||||
# 更新过期时间
|
||||
self._update_expire_time()
|
||||
|
||||
def pop(self, index: int = -1) -> Any:
|
||||
"""删除并返回列表项
|
||||
|
||||
参数:
|
||||
index: 列表索引，默认为最后一项
|
||||
|
||||
返回:
|
||||
Any: 列表值
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
raise IndexError("从空列表中弹出")
|
||||
|
||||
if not self._data:
|
||||
raise IndexError("从空列表中弹出")
|
||||
|
||||
item = self._data.pop(index)
|
||||
|
||||
# 更新过期时间
|
||||
self._update_expire_time()
|
||||
|
||||
return item.value
|
||||
|
||||
def remove(self, value: Any) -> None:
|
||||
"""删除第一个匹配的列表项
|
||||
|
||||
参数:
|
||||
value: 要删除的值
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
raise ValueError(f"{value} 不在列表中")
|
||||
|
||||
# 查找匹配的项
|
||||
for i, item in enumerate(self._data):
|
||||
if item.value == value:
|
||||
del self._data[i]
|
||||
# 更新过期时间
|
||||
self._update_expire_time()
|
||||
return
|
||||
|
||||
raise ValueError(f"{value} 不在列表中")
|
||||
|
||||
def clear(self) -> None:
|
||||
"""清空列表"""
|
||||
self._data.clear()
|
||||
# 重置过期时间
|
||||
self._update_expire_time()
|
||||
|
||||
def index(self, value: Any, start: int = 0, end: int | None = None) -> int:
|
||||
"""查找值的索引
|
||||
|
||||
参数:
|
||||
value: 要查找的值
|
||||
start: 起始索引
|
||||
end: 结束索引
|
||||
|
||||
返回:
|
||||
int: 索引位置
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
raise ValueError(f"{value} 不在列表中")
|
||||
|
||||
end = end if end is not None else len(self._data)
|
||||
|
||||
for i in range(start, min(end, len(self._data))):
|
||||
if self._data[i].value == value:
|
||||
return i
|
||||
|
||||
raise ValueError(f"{value} 不在列表中")
|
||||
|
||||
def count(self, value: Any) -> int:
|
||||
"""计算值出现的次数
|
||||
|
||||
参数:
|
||||
value: 要计数的值
|
||||
|
||||
返回:
|
||||
int: 出现次数
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
return 0
|
||||
|
||||
return sum(1 for item in self._data if item.value == value)
|
||||
|
||||
def _is_expired(self) -> bool:
|
||||
"""检查整个列表是否过期"""
|
||||
return self._expire_time > 0 and self._expire_time < time.time()
|
||||
|
||||
def _update_expire_time(self) -> None:
|
||||
"""更新过期时间"""
|
||||
if self.expire > 0:
|
||||
self._expire_time = time.time() + self.expire
|
||||
else:
|
||||
self._expire_time = 0
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""字符串表示
|
||||
|
||||
返回:
|
||||
str: 字符串表示
|
||||
"""
|
||||
# 检查整个列表是否过期
|
||||
if self._is_expired():
|
||||
self.clear()
|
||||
return f"CacheList({self.name}, {len(self._data)} items)"
|
||||
Vendored
+35
@@ -0,0 +1,35 @@
|
||||
"""
|
||||
缓存系统配置
|
||||
"""
|
||||
|
||||
# 日志标识
|
||||
LOG_COMMAND = "CacheRoot"
|
||||
|
||||
# 默认缓存过期时间（秒）
|
||||
DEFAULT_EXPIRE = 600
|
||||
|
||||
# 缓存键前缀
|
||||
CACHE_KEY_PREFIX = "ZHENXUN"
|
||||
|
||||
# 缓存键分隔符
|
||||
CACHE_KEY_SEPARATOR = ":"
|
||||
|
||||
# 复合键分隔符（用于分隔tuple类型的cache_key_field）
|
||||
COMPOSITE_KEY_SEPARATOR = "_"
|
||||
|
||||
|
||||
# 缓存模式
|
||||
class CacheMode:
|
||||
# 内存缓存 - 使用内存存储缓存数据
|
||||
MEMORY = "MEMORY"
|
||||
# Redis缓存 - 使用Redis服务器存储缓存数据
|
||||
REDIS = "REDIS"
|
||||
# 不使用缓存 - 将使用ttl=0的内存缓存，相当于直接从数据库获取数据
|
||||
NONE = "NONE"
|
||||
|
||||
|
||||
SPECIAL_KEY_FORMATS = {
|
||||
"LEVEL": "{user_id}" + COMPOSITE_KEY_SEPARATOR + "{group_id}",
|
||||
"BAN": "{user_id}" + COMPOSITE_KEY_SEPARATOR + "{group_id}",
|
||||
"GROUPS": "{group_id}" + COMPOSITE_KEY_SEPARATOR + "{channel_id}",
|
||||
}
|
||||
@@ -0,0 +1,653 @@
|
||||
from typing import Any, ClassVar, Generic, TypeVar, cast
|
||||
|
||||
from zhenxun.services.cache import Cache, CacheRoot, cache_config
|
||||
from zhenxun.services.cache.config import COMPOSITE_KEY_SEPARATOR, CacheMode
|
||||
from zhenxun.services.db_context import Model, with_db_timeout
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
T = TypeVar("T", bound=Model)
|
||||
|
||||
|
||||
class DataAccess(Generic[T]):
|
||||
"""数据访问层，根据配置决定是否使用缓存
|
||||
|
||||
使用示例:
|
||||
```python
|
||||
from zhenxun.services import DataAccess
|
||||
from zhenxun.models.plugin_info import PluginInfo
|
||||
|
||||
# 创建数据访问对象
|
||||
plugin_dao = DataAccess(PluginInfo)
|
||||
|
||||
# 获取单个数据
|
||||
plugin = await plugin_dao.get(module="example_module")
|
||||
|
||||
# 获取所有数据
|
||||
all_plugins = await plugin_dao.all()
|
||||
|
||||
# 筛选数据
|
||||
enabled_plugins = await plugin_dao.filter(status=True)
|
||||
|
||||
# 创建数据
|
||||
new_plugin = await plugin_dao.create(
|
||||
module="new_module",
|
||||
name="新插件",
|
||||
status=True
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
# 添加缓存统计信息
|
||||
_cache_stats: ClassVar[dict] = {}
|
||||
# 空结果标记
|
||||
_NULL_RESULT = "__NULL_RESULT_PLACEHOLDER__"
|
||||
# 默认空结果缓存时间（秒）- 设置为5分钟，避免频繁查询数据库
|
||||
_NULL_RESULT_TTL = 300
|
||||
|
||||
@classmethod
|
||||
def set_null_result_ttl(cls, seconds: int) -> None:
|
||||
"""设置空结果缓存时间
|
||||
|
||||
参数:
|
||||
seconds: 缓存时间（秒）
|
||||
"""
|
||||
if seconds < 0:
|
||||
raise ValueError("缓存时间不能为负数")
|
||||
cls._NULL_RESULT_TTL = seconds
|
||||
logger.info(f"已设置DataAccess空结果缓存时间为 {seconds} 秒")
|
||||
|
||||
@classmethod
|
||||
def get_null_result_ttl(cls) -> int:
|
||||
"""获取空结果缓存时间
|
||||
|
||||
返回:
|
||||
int: 缓存时间（秒）
|
||||
"""
|
||||
return cls._NULL_RESULT_TTL
|
||||
|
||||
def __init__(
|
||||
self, model_cls: type[T], key_field: str = "id", cache_type: str | None = None
|
||||
):
|
||||
"""初始化数据访问对象
|
||||
|
||||
参数:
|
||||
model_cls: 模型类
|
||||
key_field: 主键字段
|
||||
"""
|
||||
self.model_cls = model_cls
|
||||
self.key_field = getattr(model_cls, "cache_key_field", key_field)
|
||||
self.cache_type = getattr(model_cls, "cache_type", cache_type)
|
||||
|
||||
if not self.cache_type:
|
||||
raise ValueError("缓存类型不能为空")
|
||||
self.cache = Cache(self.cache_type)
|
||||
|
||||
# 初始化缓存统计
|
||||
if self.cache_type not in self._cache_stats:
|
||||
self._cache_stats[self.cache_type] = {
|
||||
"hits": 0, # 缓存命中次数
|
||||
"misses": 0, # 缓存未命中次数
|
||||
"null_hits": 0, # 空结果缓存命中次数
|
||||
"sets": 0, # 缓存设置次数
|
||||
"null_sets": 0, # 空结果缓存设置次数
|
||||
"deletes": 0, # 缓存删除次数
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def get_cache_stats(cls):
|
||||
"""获取缓存统计信息"""
|
||||
result = []
|
||||
for cache_type, stats in cls._cache_stats.items():
|
||||
hits = stats["hits"]
|
||||
null_hits = stats.get("null_hits", 0)
|
||||
misses = stats["misses"]
|
||||
total = hits + null_hits + misses
|
||||
hit_rate = ((hits + null_hits) / total * 100) if total > 0 else 0
|
||||
result.append(
|
||||
{
|
||||
"cache_type": cache_type,
|
||||
"hits": hits,
|
||||
"null_hits": null_hits,
|
||||
"misses": misses,
|
||||
"sets": stats["sets"],
|
||||
"null_sets": stats.get("null_sets", 0),
|
||||
"deletes": stats["deletes"],
|
||||
"hit_rate": f"{hit_rate:.2f}%",
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def reset_cache_stats(cls):
|
||||
"""重置缓存统计信息"""
|
||||
for stats in cls._cache_stats.values():
|
||||
stats["hits"] = 0
|
||||
stats["null_hits"] = 0
|
||||
stats["misses"] = 0
|
||||
stats["sets"] = 0
|
||||
stats["null_sets"] = 0
|
||||
stats["deletes"] = 0
|
||||
|
||||
def _build_cache_key_from_kwargs(self, **kwargs) -> str | None:
|
||||
"""从关键字参数构建缓存键
|
||||
|
||||
参数:
|
||||
**kwargs: 关键字参数
|
||||
|
||||
返回:
|
||||
str | None: 缓存键，如果无法构建则返回None
|
||||
"""
|
||||
if isinstance(self.key_field, tuple):
|
||||
# 多字段主键
|
||||
key_parts = []
|
||||
key_parts.extend(str(kwargs.get(field, "")) for field in self.key_field)
|
||||
return COMPOSITE_KEY_SEPARATOR.join(key_parts) if key_parts else None
|
||||
elif self.key_field in kwargs:
|
||||
# 单字段主键
|
||||
return str(kwargs[self.key_field])
|
||||
return None
|
||||
|
||||
async def safe_get_or_none(self, *args, **kwargs) -> T | None:
|
||||
"""安全的获取单条数据
|
||||
|
||||
参数:
|
||||
*args: 查询参数
|
||||
**kwargs: 查询参数
|
||||
|
||||
返回:
|
||||
Optional[T]: 查询结果，如果不存在返回None
|
||||
"""
|
||||
# 如果没有缓存类型，直接从数据库获取
|
||||
if not self.cache_type or cache_config.cache_mode == CacheMode.NONE:
|
||||
logger.debug(f"{self.model_cls.__name__} 直接从数据库获取数据: {kwargs}")
|
||||
return await with_db_timeout(
|
||||
self.model_cls.safe_get_or_none(*args, **kwargs),
|
||||
operation=f"{self.model_cls.__name__}.safe_get_or_none",
|
||||
)
|
||||
|
||||
# 尝试从缓存获取
|
||||
cache_key = None
|
||||
try:
|
||||
# 尝试构建缓存键
|
||||
cache_key = self._build_cache_key_from_kwargs(**kwargs)
|
||||
|
||||
# 如果成功构建缓存键，尝试从缓存获取
|
||||
if cache_key is not None:
|
||||
data = await self.cache.get(cache_key)
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} self.cache.get(cache_key)"
|
||||
f" 从缓存获取到的数据 {type(data)}: {data}"
|
||||
)
|
||||
if data == self._NULL_RESULT:
|
||||
# 空结果缓存命中
|
||||
self._cache_stats[self.cache_type]["null_hits"] += 1
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} 从缓存获取到空结果: {cache_key}"
|
||||
)
|
||||
return None
|
||||
elif data:
|
||||
# 缓存命中
|
||||
self._cache_stats[self.cache_type]["hits"] += 1
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} 从缓存获取数据成功: {cache_key}"
|
||||
)
|
||||
return cast(T, data)
|
||||
else:
|
||||
# 缓存未命中
|
||||
self._cache_stats[self.cache_type]["misses"] += 1
|
||||
logger.debug(f"{self.model_cls.__name__} 缓存未命中: {cache_key}")
|
||||
except Exception as e:
|
||||
logger.error(f"{self.model_cls.__name__} 从缓存获取数据失败: {kwargs}", e=e)
|
||||
|
||||
# 如果缓存中没有，从数据库获取
|
||||
logger.debug(f"{self.model_cls.__name__} 从数据库获取数据: {kwargs}")
|
||||
data = await self.model_cls.safe_get_or_none(*args, **kwargs)
|
||||
|
||||
# 如果获取到数据，存入缓存
|
||||
if data:
|
||||
try:
|
||||
# 生成缓存键
|
||||
cache_key = self._build_cache_key_for_item(data)
|
||||
if cache_key is not None:
|
||||
# 存入缓存
|
||||
await self.cache.set(cache_key, data)
|
||||
self._cache_stats[self.cache_type]["sets"] += 1
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} 数据已存入缓存: {cache_key}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"{self.model_cls.__name__} 存入缓存失败，参数: {kwargs}", e=e
|
||||
)
|
||||
elif cache_key is not None:
|
||||
# 如果没有获取到数据，缓存空结果
|
||||
try:
|
||||
# 存入空结果缓存，使用较短的过期时间
|
||||
await self.cache.set(
|
||||
cache_key, self._NULL_RESULT, expire=self._NULL_RESULT_TTL
|
||||
)
|
||||
self._cache_stats[self.cache_type]["null_sets"] += 1
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} 空结果已存入缓存: {cache_key},"
|
||||
f" TTL={self._NULL_RESULT_TTL}秒"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"{self.model_cls.__name__} 存入空结果缓存失败，参数: {kwargs}", e=e
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
async def get_or_none(self, *args, **kwargs) -> T | None:
|
||||
"""获取单条数据
|
||||
|
||||
参数:
|
||||
*args: 查询参数
|
||||
**kwargs: 查询参数
|
||||
|
||||
返回:
|
||||
Optional[T]: 查询结果，如果不存在返回None
|
||||
"""
|
||||
# 如果没有缓存类型，直接从数据库获取
|
||||
if not self.cache_type or cache_config.cache_mode == CacheMode.NONE:
|
||||
logger.debug(f"{self.model_cls.__name__} 直接从数据库获取数据: {kwargs}")
|
||||
return await with_db_timeout(
|
||||
self.model_cls.get_or_none(*args, **kwargs),
|
||||
operation=f"{self.model_cls.__name__}.get_or_none",
|
||||
)
|
||||
|
||||
# 尝试从缓存获取
|
||||
cache_key = None
|
||||
try:
|
||||
# 尝试构建缓存键
|
||||
cache_key = self._build_cache_key_from_kwargs(**kwargs)
|
||||
|
||||
# 如果成功构建缓存键，尝试从缓存获取
|
||||
if cache_key is not None:
|
||||
data = await self.cache.get(cache_key)
|
||||
if data == self._NULL_RESULT:
|
||||
# 空结果缓存命中
|
||||
self._cache_stats[self.cache_type]["null_hits"] += 1
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} 从缓存获取到空结果: {cache_key}"
|
||||
)
|
||||
return None
|
||||
elif data:
|
||||
# 缓存命中
|
||||
self._cache_stats[self.cache_type]["hits"] += 1
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} 从缓存获取数据成功: {cache_key}"
|
||||
)
|
||||
return cast(T, data)
|
||||
else:
|
||||
# 缓存未命中
|
||||
self._cache_stats[self.cache_type]["misses"] += 1
|
||||
logger.debug(f"{self.model_cls.__name__} 缓存未命中: {cache_key}")
|
||||
except Exception as e:
|
||||
logger.error(f"{self.model_cls.__name__} 从缓存获取数据失败: {kwargs}", e=e)
|
||||
|
||||
# 如果缓存中没有，从数据库获取
|
||||
logger.debug(f"{self.model_cls.__name__} 从数据库获取数据: {kwargs}")
|
||||
data = await self.model_cls.get_or_none(*args, **kwargs)
|
||||
|
||||
# 如果获取到数据，存入缓存
|
||||
if data:
|
||||
try:
|
||||
cache_key = self._build_cache_key_for_item(data)
|
||||
# 生成缓存键
|
||||
if cache_key is not None:
|
||||
# 存入缓存
|
||||
await self.cache.set(cache_key, data)
|
||||
self._cache_stats[self.cache_type]["sets"] += 1
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} 数据已存入缓存: {cache_key}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"{self.model_cls.__name__} 存入缓存失败，参数: {kwargs}", e=e
|
||||
)
|
||||
elif cache_key is not None:
|
||||
# 如果没有获取到数据，缓存空结果
|
||||
try:
|
||||
# 存入空结果缓存，使用较短的过期时间
|
||||
await self.cache.set(
|
||||
cache_key, self._NULL_RESULT, expire=self._NULL_RESULT_TTL
|
||||
)
|
||||
self._cache_stats[self.cache_type]["null_sets"] += 1
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} 空结果已存入缓存: {cache_key},"
|
||||
f" TTL={self._NULL_RESULT_TTL}秒"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"{self.model_cls.__name__} 存入空结果缓存失败，参数: {kwargs}", e=e
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
async def clear_cache(self, **kwargs) -> bool:
|
||||
"""只清除缓存，不影响数据库数据
|
||||
|
||||
参数:
|
||||
**kwargs: 查询参数，必须包含主键字段
|
||||
|
||||
返回:
|
||||
bool: 是否成功清除缓存
|
||||
"""
|
||||
# 如果没有缓存类型，直接返回True
|
||||
if not self.cache_type or cache_config.cache_mode == CacheMode.NONE:
|
||||
return True
|
||||
|
||||
try:
|
||||
# 构建缓存键
|
||||
cache_key = self._build_cache_key_from_kwargs(**kwargs)
|
||||
if cache_key is None:
|
||||
if isinstance(self.key_field, tuple):
|
||||
# 如果是复合键，检查缺少哪些字段
|
||||
missing_fields = [
|
||||
field for field in self.key_field if field not in kwargs
|
||||
]
|
||||
logger.error(
|
||||
f"清除{self.model_cls.__name__}缓存失败: "
|
||||
f"缺少主键字段 {', '.join(missing_fields)}"
|
||||
)
|
||||
else:
|
||||
logger.error(
|
||||
f"清除{self.model_cls.__name__}缓存失败: "
|
||||
f"缺少主键字段 {self.key_field}"
|
||||
)
|
||||
return False
|
||||
|
||||
# 删除缓存
|
||||
await self.cache.delete(cache_key)
|
||||
self._cache_stats[self.cache_type]["deletes"] += 1
|
||||
logger.debug(f"已清除{self.model_cls.__name__}缓存: {cache_key}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"清除{self.model_cls.__name__}缓存失败", e=e)
|
||||
return False
|
||||
|
||||
def _build_composite_key(self, data: T) -> str | None:
|
||||
"""构建复合缓存键
|
||||
|
||||
参数:
|
||||
data: 数据对象
|
||||
|
||||
返回:
|
||||
str | None: 构建的缓存键，如果无法构建则返回None
|
||||
"""
|
||||
# 如果是元组，表示多个字段组成键
|
||||
if isinstance(self.key_field, tuple):
|
||||
# 构建键参数列表
|
||||
key_parts = []
|
||||
for field in self.key_field:
|
||||
value = getattr(data, field, "")
|
||||
key_parts.append(value if value is not None else "")
|
||||
|
||||
# 如果没有有效参数，返回None
|
||||
return COMPOSITE_KEY_SEPARATOR.join(key_parts) if key_parts else None
|
||||
elif hasattr(data, self.key_field):
|
||||
value = getattr(data, self.key_field, None)
|
||||
return str(value) if value is not None else None
|
||||
|
||||
return None
|
||||
|
||||
def _build_cache_key_for_item(self, item: T) -> str | None:
|
||||
"""为数据项构建缓存键
|
||||
|
||||
参数:
|
||||
item: 数据项
|
||||
|
||||
返回:
|
||||
str | None: 缓存键，如果无法生成则返回None
|
||||
"""
|
||||
# 如果没有缓存类型，返回None
|
||||
if not self.cache_type:
|
||||
return None
|
||||
|
||||
# 获取缓存类型的配置信息
|
||||
cache_model = CacheRoot.get_model(self.cache_type)
|
||||
|
||||
if not cache_model.key_format:
|
||||
# 常规处理，使用主键作为缓存键
|
||||
return self._build_composite_key(item)
|
||||
# 构建键参数字典
|
||||
key_parts = []
|
||||
# 从格式字符串中提取所需的字段名
|
||||
import re
|
||||
|
||||
field_names = re.findall(r"{([^}]+)}", cache_model.key_format)
|
||||
|
||||
# 收集所有字段值
|
||||
for field in field_names:
|
||||
value = getattr(item, field, "")
|
||||
key_parts.append(value if value is not None else "")
|
||||
|
||||
return COMPOSITE_KEY_SEPARATOR.join(key_parts)
|
||||
|
||||
async def _cache_items(self, data_list: list[T]) -> None:
|
||||
"""将数据列表存入缓存
|
||||
|
||||
参数:
|
||||
data_list: 数据列表
|
||||
"""
|
||||
if (
|
||||
not data_list
|
||||
or not self.cache_type
|
||||
or cache_config.cache_mode == CacheMode.NONE
|
||||
):
|
||||
return
|
||||
|
||||
try:
|
||||
# 遍历数据列表，将每条数据存入缓存
|
||||
cached_count = 0
|
||||
for item in data_list:
|
||||
cache_key = self._build_cache_key_for_item(item)
|
||||
if cache_key is not None:
|
||||
await self.cache.set(cache_key, item)
|
||||
cached_count += 1
|
||||
self._cache_stats[self.cache_type]["sets"] += 1
|
||||
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} 批量缓存: {cached_count}/{len(data_list)}项"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"{self.model_cls.__name__} 批量缓存失败", e=e)
|
||||
|
||||
async def filter(self, *args, **kwargs) -> list[T]:
|
||||
"""筛选数据
|
||||
|
||||
参数:
|
||||
*args: 查询参数
|
||||
**kwargs: 查询参数
|
||||
|
||||
返回:
|
||||
List[T]: 查询结果列表
|
||||
"""
|
||||
# 从数据库获取数据
|
||||
logger.debug(f"{self.model_cls.__name__} filter: 从数据库查询, 参数: {kwargs}")
|
||||
data_list = await self.model_cls.filter(*args, **kwargs)
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} filter: 查询结果数量: {len(data_list)}"
|
||||
)
|
||||
|
||||
# 将数据存入缓存
|
||||
await self._cache_items(data_list)
|
||||
|
||||
return data_list
|
||||
|
||||
async def all(self) -> list[T]:
|
||||
"""获取所有数据
|
||||
|
||||
返回:
|
||||
List[T]: 所有数据列表
|
||||
"""
|
||||
# 直接从数据库获取
|
||||
logger.debug(f"{self.model_cls.__name__} all: 从数据库查询所有数据")
|
||||
data_list = await self.model_cls.all()
|
||||
logger.debug(f"{self.model_cls.__name__} all: 查询结果数量: {len(data_list)}")
|
||||
|
||||
# 将数据存入缓存
|
||||
await self._cache_items(data_list)
|
||||
|
||||
return data_list
|
||||
|
||||
async def count(self, *args, **kwargs) -> int:
|
||||
"""获取数据数量
|
||||
|
||||
参数:
|
||||
*args: 查询参数
|
||||
**kwargs: 查询参数
|
||||
|
||||
返回:
|
||||
int: 数据数量
|
||||
"""
|
||||
# 直接从数据库获取数量
|
||||
return await self.model_cls.filter(*args, **kwargs).count()
|
||||
|
||||
async def exists(self, *args, **kwargs) -> bool:
|
||||
"""判断数据是否存在
|
||||
|
||||
参数:
|
||||
*args: 查询参数
|
||||
**kwargs: 查询参数
|
||||
|
||||
返回:
|
||||
bool: 是否存在
|
||||
"""
|
||||
# 直接从数据库判断是否存在
|
||||
return await self.model_cls.filter(*args, **kwargs).exists()
|
||||
|
||||
async def create(self, **kwargs) -> T:
|
||||
"""创建数据
|
||||
|
||||
参数:
|
||||
**kwargs: 创建参数
|
||||
|
||||
返回:
|
||||
T: 创建的数据
|
||||
"""
|
||||
# 创建数据
|
||||
logger.debug(f"{self.model_cls.__name__} create: 创建数据, 参数: {kwargs}")
|
||||
data = await self.model_cls.create(**kwargs)
|
||||
|
||||
# 如果有缓存类型，将数据存入缓存
|
||||
if self.cache_type and cache_config.cache_mode != CacheMode.NONE:
|
||||
try:
|
||||
# 生成缓存键
|
||||
cache_key = self._build_cache_key_for_item(data)
|
||||
if cache_key is not None:
|
||||
# 存入缓存
|
||||
await self.cache.set(cache_key, data)
|
||||
self._cache_stats[self.cache_type]["sets"] += 1
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} create: "
|
||||
f"新创建的数据已存入缓存: {cache_key}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"{self.model_cls.__name__} create: 存入缓存失败，参数: {kwargs}",
|
||||
e=e,
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
async def update_or_create(
|
||||
self, defaults: dict[str, Any] | None = None, **kwargs
|
||||
) -> tuple[T, bool]:
|
||||
"""更新或创建数据
|
||||
|
||||
参数:
|
||||
defaults: 默认值
|
||||
**kwargs: 查询参数
|
||||
|
||||
返回:
|
||||
tuple[T, bool]: (数据, 是否创建)
|
||||
"""
|
||||
# 更新或创建数据
|
||||
data, created = await self.model_cls.update_or_create(
|
||||
defaults=defaults, **kwargs
|
||||
)
|
||||
|
||||
# 如果有缓存类型，将数据存入缓存
|
||||
if self.cache_type and cache_config.cache_mode != CacheMode.NONE:
|
||||
try:
|
||||
# 生成缓存键
|
||||
cache_key = self._build_cache_key_for_item(data)
|
||||
if cache_key is not None:
|
||||
# 存入缓存
|
||||
await self.cache.set(cache_key, data)
|
||||
self._cache_stats[self.cache_type]["sets"] += 1
|
||||
logger.debug(f"更新或创建的数据已存入缓存: {cache_key}")
|
||||
except Exception as e:
|
||||
logger.error(f"存入缓存失败，参数: {kwargs}", e=e)
|
||||
|
||||
return data, created
|
||||
|
||||
async def delete(self, *args, **kwargs) -> int:
|
||||
"""删除数据
|
||||
|
||||
参数:
|
||||
*args: 查询参数
|
||||
**kwargs: 查询参数
|
||||
|
||||
返回:
|
||||
int: 删除的数据数量
|
||||
"""
|
||||
logger.debug(f"{self.model_cls.__name__} delete: 删除数据, 参数: {kwargs}")
|
||||
|
||||
# 如果有缓存类型且有key_field参数，先尝试删除缓存
|
||||
if self.cache_type and cache_config.cache_mode != CacheMode.NONE:
|
||||
try:
|
||||
# 尝试构建缓存键
|
||||
cache_key = self._build_cache_key_from_kwargs(**kwargs)
|
||||
|
||||
if cache_key is not None:
|
||||
# 如果成功构建缓存键，直接删除缓存
|
||||
await self.cache.delete(cache_key)
|
||||
self._cache_stats[self.cache_type]["deletes"] += 1
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} delete: 已删除缓存: {cache_key}"
|
||||
)
|
||||
else:
|
||||
# 否则需要先查询出要删除的数据，然后删除对应的缓存
|
||||
items = await self.model_cls.filter(*args, **kwargs)
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} delete:"
|
||||
f" 查询到 {len(items)} 条要删除的数据"
|
||||
)
|
||||
for item in items:
|
||||
item_cache_key = self._build_cache_key_for_item(item)
|
||||
if item_cache_key is not None:
|
||||
await self.cache.delete(item_cache_key)
|
||||
self._cache_stats[self.cache_type]["deletes"] += 1
|
||||
if items:
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} delete:"
|
||||
f" 已删除 {len(items)} 条数据的缓存"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"{self.model_cls.__name__} delete: 删除缓存失败", e=e)
|
||||
|
||||
# 删除数据
|
||||
result = await self.model_cls.filter(*args, **kwargs).delete()
|
||||
logger.debug(
|
||||
f"{self.model_cls.__name__} delete: 已从数据库删除 {result} 条数据"
|
||||
)
|
||||
return result
|
||||
|
||||
def _generate_cache_key(self, data: T) -> str:
|
||||
"""根据数据对象生成缓存键
|
||||
|
||||
参数:
|
||||
data: 数据对象
|
||||
|
||||
返回:
|
||||
str: 缓存键
|
||||
"""
|
||||
# 使用新方法构建复合键
|
||||
if composite_key := self._build_composite_key(data):
|
||||
return composite_key
|
||||
|
||||
# 如果无法生成复合键，生成一个唯一键
|
||||
return f"object_{id(data)}"
|
||||
+377
-13
@@ -1,37 +1,328 @@
|
||||
import nonebot
|
||||
import asyncio
|
||||
from collections.abc import Iterable
|
||||
import contextlib
|
||||
import time
|
||||
from typing import Any, ClassVar
|
||||
from typing_extensions import Self
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from nonebot import get_driver
|
||||
from nonebot.utils import is_coroutine_callable
|
||||
from tortoise import Tortoise
|
||||
from tortoise.backends.base.client import BaseDBAsyncClient
|
||||
from tortoise.connection import connections
|
||||
from tortoise.models import Model as Model_
|
||||
from tortoise.exceptions import IntegrityError, MultipleObjectsReturned
|
||||
from tortoise.models import Model as TortoiseModel
|
||||
from tortoise.transactions import in_transaction
|
||||
|
||||
from zhenxun.configs.config import BotConfig
|
||||
from zhenxun.services.cache import CacheRoot
|
||||
from zhenxun.services.log import logger
|
||||
from zhenxun.utils.enum import DbLockType
|
||||
from zhenxun.utils.exception import HookPriorityException
|
||||
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
|
||||
|
||||
from .log import logger
|
||||
driver = get_driver()
|
||||
|
||||
SCRIPT_METHOD = []
|
||||
MODELS: list[str] = []
|
||||
|
||||
# 数据库操作超时设置（秒）
|
||||
DB_TIMEOUT_SECONDS = 3.0
|
||||
|
||||
driver = nonebot.get_driver()
|
||||
# 性能监控阈值（秒）
|
||||
SLOW_QUERY_THRESHOLD = 0.5
|
||||
|
||||
LOG_COMMAND = "DbContext"
|
||||
|
||||
|
||||
class Model(Model_):
|
||||
async def with_db_timeout(
|
||||
coro, timeout: float = DB_TIMEOUT_SECONDS, operation: str | None = None
|
||||
):
|
||||
"""带超时控制的数据库操作"""
|
||||
start_time = time.time()
|
||||
try:
|
||||
result = await asyncio.wait_for(coro, timeout=timeout)
|
||||
elapsed = time.time() - start_time
|
||||
if elapsed > SLOW_QUERY_THRESHOLD and operation:
|
||||
logger.warning(f"慢查询: {operation} 耗时 {elapsed:.3f}s", LOG_COMMAND)
|
||||
return result
|
||||
except asyncio.TimeoutError:
|
||||
if operation:
|
||||
logger.error(f"数据库操作超时: {operation} (>{timeout}s)", LOG_COMMAND)
|
||||
raise
|
||||
|
||||
|
||||
class Model(TortoiseModel):
|
||||
"""
|
||||
自动添加模块
|
||||
|
||||
Args:
|
||||
Model_: Model
|
||||
增强的ORM基类，解决锁嵌套问题
|
||||
"""
|
||||
|
||||
sem_data: ClassVar[dict[str, dict[str, asyncio.Semaphore]]] = {}
|
||||
_current_locks: ClassVar[dict[int, DbLockType]] = {} # 跟踪当前协程持有的锁
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
if cls.__module__ not in MODELS:
|
||||
MODELS.append(cls.__module__)
|
||||
|
||||
if func := getattr(cls, "_run_script", None):
|
||||
SCRIPT_METHOD.append((cls.__module__, func))
|
||||
|
||||
@classmethod
|
||||
def get_cache_type(cls) -> str | None:
|
||||
"""获取缓存类型"""
|
||||
return getattr(cls, "cache_type", None)
|
||||
|
||||
@classmethod
|
||||
def get_cache_key_field(cls) -> str | tuple[str]:
|
||||
"""获取缓存键字段"""
|
||||
return getattr(cls, "cache_key_field", "id")
|
||||
|
||||
@classmethod
|
||||
def get_cache_key(cls, instance) -> str | None:
|
||||
"""获取缓存键
|
||||
|
||||
参数:
|
||||
instance: 模型实例
|
||||
|
||||
返回:
|
||||
str | None: 缓存键，如果无法获取则返回None
|
||||
"""
|
||||
from zhenxun.services.cache.config import COMPOSITE_KEY_SEPARATOR
|
||||
|
||||
key_field = cls.get_cache_key_field()
|
||||
|
||||
if isinstance(key_field, tuple):
|
||||
# 多字段主键
|
||||
key_parts = []
|
||||
for field in key_field:
|
||||
if hasattr(instance, field):
|
||||
value = getattr(instance, field, None)
|
||||
key_parts.append(value if value is not None else "")
|
||||
else:
|
||||
# 如果缺少任何必要的字段，返回None
|
||||
key_parts.append("")
|
||||
|
||||
# 如果没有有效参数，返回None
|
||||
return COMPOSITE_KEY_SEPARATOR.join(key_parts) if key_parts else None
|
||||
elif hasattr(instance, key_field):
|
||||
value = getattr(instance, key_field, None)
|
||||
return str(value) if value is not None else None
|
||||
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_semaphore(cls, lock_type: DbLockType):
|
||||
enable_lock = getattr(cls, "enable_lock", None)
|
||||
if not enable_lock or lock_type not in enable_lock:
|
||||
return None
|
||||
|
||||
if cls.__name__ not in cls.sem_data:
|
||||
cls.sem_data[cls.__name__] = {}
|
||||
if lock_type not in cls.sem_data[cls.__name__]:
|
||||
cls.sem_data[cls.__name__][lock_type] = asyncio.Semaphore(1)
|
||||
return cls.sem_data[cls.__name__][lock_type]
|
||||
|
||||
@classmethod
|
||||
def _require_lock(cls, lock_type: DbLockType) -> bool:
|
||||
"""检查是否需要真正加锁"""
|
||||
task_id = id(asyncio.current_task())
|
||||
return cls._current_locks.get(task_id) != lock_type
|
||||
|
||||
@classmethod
|
||||
@contextlib.asynccontextmanager
|
||||
async def _lock_context(cls, lock_type: DbLockType):
|
||||
"""带重入检查的锁上下文"""
|
||||
task_id = id(asyncio.current_task())
|
||||
need_lock = cls._require_lock(lock_type)
|
||||
|
||||
if need_lock and (sem := cls.get_semaphore(lock_type)):
|
||||
cls._current_locks[task_id] = lock_type
|
||||
async with sem:
|
||||
yield
|
||||
cls._current_locks.pop(task_id, None)
|
||||
else:
|
||||
yield
|
||||
|
||||
@classmethod
|
||||
async def create(
|
||||
cls, using_db: BaseDBAsyncClient | None = None, **kwargs: Any
|
||||
) -> Self:
|
||||
"""创建数据（使用CREATE锁）"""
|
||||
async with cls._lock_context(DbLockType.CREATE):
|
||||
# 直接调用父类的_create方法避免触发save的锁
|
||||
result = await super().create(using_db=using_db, **kwargs)
|
||||
if cache_type := cls.get_cache_type():
|
||||
await CacheRoot.invalidate_cache(cache_type, cls.get_cache_key(result))
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
async def get_or_create(
|
||||
cls,
|
||||
defaults: dict | None = None,
|
||||
using_db: BaseDBAsyncClient | None = None,
|
||||
**kwargs: Any,
|
||||
) -> tuple[Self, bool]:
|
||||
"""获取或创建数据（无锁版本，依赖数据库约束）"""
|
||||
result = await super().get_or_create(
|
||||
defaults=defaults, using_db=using_db, **kwargs
|
||||
)
|
||||
if cache_type := cls.get_cache_type():
|
||||
await CacheRoot.invalidate_cache(cache_type, cls.get_cache_key(result[0]))
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
async def update_or_create(
|
||||
cls,
|
||||
defaults: dict | None = None,
|
||||
using_db: BaseDBAsyncClient | None = None,
|
||||
**kwargs: Any,
|
||||
) -> tuple[Self, bool]:
|
||||
"""更新或创建数据（使用UPSERT锁）"""
|
||||
async with cls._lock_context(DbLockType.UPSERT):
|
||||
try:
|
||||
# 先尝试更新（带行锁）
|
||||
async with in_transaction():
|
||||
if obj := await cls.filter(**kwargs).select_for_update().first():
|
||||
await obj.update_from_dict(defaults or {})
|
||||
await obj.save()
|
||||
result = (obj, False)
|
||||
else:
|
||||
# 创建时不重复加锁
|
||||
result = await cls.create(**kwargs, **(defaults or {})), True
|
||||
|
||||
if cache_type := cls.get_cache_type():
|
||||
await CacheRoot.invalidate_cache(
|
||||
cache_type, cls.get_cache_key(result[0])
|
||||
)
|
||||
return result
|
||||
except IntegrityError:
|
||||
# 处理极端情况下的唯一约束冲突
|
||||
obj = await cls.get(**kwargs)
|
||||
return obj, False
|
||||
|
||||
async def save(
|
||||
self,
|
||||
using_db: BaseDBAsyncClient | None = None,
|
||||
update_fields: Iterable[str] | None = None,
|
||||
force_create: bool = False,
|
||||
force_update: bool = False,
|
||||
):
|
||||
"""保存数据（根据操作类型自动选择锁）"""
|
||||
lock_type = (
|
||||
DbLockType.CREATE
|
||||
if getattr(self, "id", None) is None
|
||||
else DbLockType.UPDATE
|
||||
)
|
||||
async with self._lock_context(lock_type):
|
||||
await super().save(
|
||||
using_db=using_db,
|
||||
update_fields=update_fields,
|
||||
force_create=force_create,
|
||||
force_update=force_update,
|
||||
)
|
||||
if cache_type := getattr(self, "cache_type", None):
|
||||
await CacheRoot.invalidate_cache(
|
||||
cache_type, self.__class__.get_cache_key(self)
|
||||
)
|
||||
|
||||
async def delete(self, using_db: BaseDBAsyncClient | None = None):
|
||||
cache_type = getattr(self, "cache_type", None)
|
||||
key = self.__class__.get_cache_key(self) if cache_type else None
|
||||
# 执行删除操作
|
||||
await super().delete(using_db=using_db)
|
||||
|
||||
# 清除缓存
|
||||
if cache_type:
|
||||
await CacheRoot.invalidate_cache(cache_type, key)
|
||||
|
||||
@classmethod
|
||||
async def safe_get_or_none(
|
||||
cls,
|
||||
*args,
|
||||
using_db: BaseDBAsyncClient | None = None,
|
||||
clean_duplicates: bool = True,
|
||||
**kwargs: Any,
|
||||
) -> Self | None:
|
||||
"""安全地获取一条记录或None，处理存在多个记录时返回最新的那个
|
||||
注意，默认会删除重复的记录，仅保留最新的
|
||||
|
||||
参数:
|
||||
*args: 查询参数
|
||||
using_db: 数据库连接
|
||||
clean_duplicates: 是否删除重复的记录，仅保留最新的
|
||||
**kwargs: 查询参数
|
||||
|
||||
返回:
|
||||
Self | None: 查询结果，如果不存在返回None
|
||||
"""
|
||||
try:
|
||||
# 先尝试使用 get_or_none 获取单个记录
|
||||
try:
|
||||
return await with_db_timeout(
|
||||
cls.get_or_none(*args, using_db=using_db, **kwargs),
|
||||
operation=f"{cls.__name__}.get_or_none",
|
||||
)
|
||||
except MultipleObjectsReturned:
|
||||
# 如果出现多个记录的情况，进行特殊处理
|
||||
logger.warning(
|
||||
f"{cls.__name__} safe_get_or_none 发现多个记录: {kwargs}",
|
||||
LOG_COMMAND,
|
||||
)
|
||||
|
||||
# 查询所有匹配记录
|
||||
records = await with_db_timeout(
|
||||
cls.filter(*args, **kwargs).all(),
|
||||
operation=f"{cls.__name__}.filter.all",
|
||||
)
|
||||
|
||||
if not records:
|
||||
return None
|
||||
|
||||
# 如果需要清理重复记录
|
||||
if clean_duplicates and hasattr(records[0], "id"):
|
||||
# 按 id 排序
|
||||
records = sorted(
|
||||
records, key=lambda x: getattr(x, "id", 0), reverse=True
|
||||
)
|
||||
for record in records[1:]:
|
||||
try:
|
||||
await with_db_timeout(
|
||||
record.delete(),
|
||||
operation=f"{cls.__name__}.delete_duplicate",
|
||||
)
|
||||
logger.info(
|
||||
f"{cls.__name__} 删除重复记录:"
|
||||
f" id={getattr(record, 'id', None)}",
|
||||
LOG_COMMAND,
|
||||
)
|
||||
except Exception as del_e:
|
||||
logger.error(f"删除重复记录失败: {del_e}")
|
||||
return records[0]
|
||||
# 如果不需要清理或没有 id 字段，则返回最新的记录
|
||||
if hasattr(cls, "id"):
|
||||
return await with_db_timeout(
|
||||
cls.filter(*args, **kwargs).order_by("-id").first(),
|
||||
operation=f"{cls.__name__}.filter.order_by.first",
|
||||
)
|
||||
# 如果没有 id 字段，则返回第一个记录
|
||||
return await with_db_timeout(
|
||||
cls.filter(*args, **kwargs).first(),
|
||||
operation=f"{cls.__name__}.filter.first",
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(
|
||||
f"数据库操作超时: {cls.__name__}.safe_get_or_none", LOG_COMMAND
|
||||
)
|
||||
return None
|
||||
except Exception as e:
|
||||
# 其他类型的错误则继续抛出
|
||||
logger.error(
|
||||
f"数据库操作异常: {cls.__name__}.safe_get_or_none, {e!s}", LOG_COMMAND
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
class DbUrlIsNode(HookPriorityException):
|
||||
"""
|
||||
@@ -49,6 +340,77 @@ class DbConnectError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
POSTGRESQL_CONFIG = {
|
||||
"max_size": 30, # 最大连接数
|
||||
"min_size": 5, # 最小保持的连接数（可选）
|
||||
}
|
||||
|
||||
|
||||
MYSQL_CONFIG = {
|
||||
"max_connections": 20, # 最大连接数
|
||||
"connect_timeout": 30, # 连接超时（可选）
|
||||
}
|
||||
|
||||
SQLITE_CONFIG = {
|
||||
"journal_mode": "WAL", # 提高并发写入性能
|
||||
"timeout": 30, # 锁等待超时（可选）
|
||||
}
|
||||
|
||||
|
||||
def get_config(db_url: str) -> dict:
|
||||
"""获取数据库配置"""
|
||||
parsed = urlparse(BotConfig.db_url)
|
||||
|
||||
# 基础配置
|
||||
config = {
|
||||
"connections": {
|
||||
"default": BotConfig.db_url # 默认直接使用连接字符串
|
||||
},
|
||||
"apps": {
|
||||
"models": {
|
||||
"models": MODELS,
|
||||
"default_connection": "default",
|
||||
}
|
||||
},
|
||||
"timezone": "Asia/Shanghai",
|
||||
}
|
||||
|
||||
# 根据数据库类型应用高级配置
|
||||
if parsed.scheme.startswith("postgres"):
|
||||
config["connections"]["default"] = {
|
||||
"engine": "tortoise.backends.asyncpg",
|
||||
"credentials": {
|
||||
"host": parsed.hostname,
|
||||
"port": parsed.port or 5432,
|
||||
"user": parsed.username,
|
||||
"password": parsed.password,
|
||||
"database": parsed.path[1:],
|
||||
},
|
||||
**POSTGRESQL_CONFIG,
|
||||
}
|
||||
elif parsed.scheme == "mysql":
|
||||
config["connections"]["default"] = {
|
||||
"engine": "tortoise.backends.mysql",
|
||||
"credentials": {
|
||||
"host": parsed.hostname,
|
||||
"port": parsed.port or 3306,
|
||||
"user": parsed.username,
|
||||
"password": parsed.password,
|
||||
"database": parsed.path[1:],
|
||||
},
|
||||
**MYSQL_CONFIG,
|
||||
}
|
||||
elif parsed.scheme == "sqlite":
|
||||
config["connections"]["default"] = {
|
||||
"engine": "tortoise.backends.sqlite",
|
||||
"credentials": {
|
||||
"file_path": parsed.path[1:] or ":memory:",
|
||||
},
|
||||
**SQLITE_CONFIG,
|
||||
}
|
||||
return config
|
||||
|
||||
|
||||
@PriorityLifecycle.on_startup(priority=1)
|
||||
async def init():
|
||||
if not BotConfig.db_url:
|
||||
@@ -64,9 +426,7 @@ async def init():
|
||||
raise DbUrlIsNode("\n" + error.strip())
|
||||
try:
|
||||
await Tortoise.init(
|
||||
db_url=BotConfig.db_url,
|
||||
modules={"models": MODELS},
|
||||
timezone="Asia/Shanghai",
|
||||
config=get_config(BotConfig.db_url),
|
||||
)
|
||||
if SCRIPT_METHOD:
|
||||
db = Tortoise.get_connection("default")
|
||||
@@ -85,13 +445,17 @@ async def init():
|
||||
for sql in sql_list:
|
||||
logger.debug(f"执行SQL: {sql}")
|
||||
try:
|
||||
await db.execute_query_dict(sql)
|
||||
await asyncio.wait_for(
|
||||
db.execute_query_dict(sql), timeout=DB_TIMEOUT_SECONDS
|
||||
)
|
||||
# await TestSQL.raw(sql)
|
||||
except Exception as e:
|
||||
logger.debug(f"执行SQL: {sql} 错误...", e=e)
|
||||
if sql_list:
|
||||
logger.debug("SCRIPT_METHOD方法执行完毕!")
|
||||
logger.debug("开始生成数据库表结构...")
|
||||
await Tortoise.generate_schemas()
|
||||
logger.debug("数据库表结构生成完毕!")
|
||||
logger.info("Database loaded successfully!")
|
||||
except Exception as e:
|
||||
raise DbConnectError(f"数据库连接错误... e:{e}") from e
|
||||
|
||||
Reference in New Issue
Block a user