400 lines
12 KiB
Python
400 lines
12 KiB
Python
from contextlib import asynccontextmanager
|
||
import os
|
||
from pathlib import Path
|
||
import re
|
||
from typing import Any
|
||
|
||
import aiosqlite
|
||
|
||
from zhenxun.services.log import logger
|
||
|
||
from ...utils.config import g_sDBFilePath, g_sDBPath
|
||
|
||
|
||
class CSqlManager:
|
||
def __init__(self):
|
||
dbPath = Path(g_sDBPath)
|
||
if dbPath and not dbPath.exists():
|
||
os.makedirs(dbPath, exist_ok=True)
|
||
|
||
@classmethod
|
||
async def cleanup(cls):
|
||
if hasattr(cls, "m_pDB") and cls.m_pDB:
|
||
await cls.m_pDB.close()
|
||
|
||
@classmethod
|
||
async def init(cls) -> bool:
|
||
try:
|
||
cls.m_pDB = await aiosqlite.connect(g_sDBFilePath)
|
||
cls.m_pDB.row_factory = aiosqlite.Row
|
||
return True
|
||
except Exception as e:
|
||
logger.debug("真寻农场初始化总数据库失败", e=e)
|
||
return False
|
||
|
||
@classmethod
|
||
@asynccontextmanager
|
||
async def _transaction(cls):
|
||
await cls.m_pDB.execute("BEGIN;")
|
||
try:
|
||
yield
|
||
except:
|
||
await cls.m_pDB.execute("ROLLBACK;")
|
||
raise
|
||
else:
|
||
await cls.m_pDB.execute("COMMIT;")
|
||
|
||
@classmethod
|
||
async def getTableInfo(cls, tableName: str) -> list:
|
||
if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", tableName):
|
||
raise ValueError(f"Illegal table name: {tableName}")
|
||
try:
|
||
cursor = await cls.m_pDB.execute(f'PRAGMA table_info("{tableName}")')
|
||
rows = await cursor.fetchall()
|
||
return [{"name": row[1], "type": row[2]} for row in rows]
|
||
except aiosqlite.Error:
|
||
return []
|
||
|
||
@classmethod
|
||
async def ensureTableSchema(cls, tableName: str, columns: dict) -> bool:
|
||
"""由AI生成
|
||
创建表或为已存在表添加缺失字段。
|
||
返回 True 表示有变更（创建或新增列），False 则无操作
|
||
|
||
Args:
|
||
tableName (_type_): 表名
|
||
columns (_type_): 字典
|
||
|
||
Returns:
|
||
_type_: _description_
|
||
"""
|
||
|
||
info = await cls.getTableInfo(tableName)
|
||
existing = {col["name"]: col["type"].upper() for col in info}
|
||
desired = {k: v.upper() for k, v in columns.items() if k != "PRIMARY KEY"}
|
||
primaryKey = columns.get("PRIMARY KEY", "")
|
||
|
||
if not existing:
|
||
colsDef = ", ".join(f'"{k}" {v}' for k, v in desired.items())
|
||
if primaryKey:
|
||
colsDef += f", PRIMARY KEY {primaryKey}"
|
||
await cls.m_pDB.execute(f'CREATE TABLE "{tableName}" ({colsDef});')
|
||
return True
|
||
|
||
toAdd = [k for k in desired if k not in existing]
|
||
toRemove = [k for k in existing if k not in desired]
|
||
typeMismatch = [
|
||
k for k in desired if k in existing and existing[k] != desired[k]
|
||
]
|
||
|
||
if toAdd and not toRemove and not typeMismatch:
|
||
for col in toAdd:
|
||
await cls.m_pDB.execute(
|
||
f'ALTER TABLE "{tableName}" ADD COLUMN "{col}" {columns[col]}'
|
||
)
|
||
return True
|
||
|
||
async with cls._transaction():
|
||
tmpTable = f"{tableName}_new"
|
||
colsDef = ", ".join(f'"{k}" {v}' for k, v in desired.items())
|
||
if primaryKey:
|
||
colsDef += f", PRIMARY KEY {primaryKey}"
|
||
await cls.m_pDB.execute(f'CREATE TABLE "{tmpTable}" ({colsDef});')
|
||
|
||
commonCols = [k for k in desired if k in existing]
|
||
if commonCols:
|
||
colsStr = ", ".join(f'"{c}"' for c in commonCols)
|
||
|
||
sql = (
|
||
f'INSERT INTO "{tmpTable}" ({colsStr}) '
|
||
f"SELECT {colsStr} "
|
||
f'FROM "{tableName}";'
|
||
)
|
||
|
||
await cls.m_pDB.execute(sql)
|
||
await cls.m_pDB.execute(f'DROP TABLE "{tableName}";')
|
||
await cls.m_pDB.execute(
|
||
f'ALTER TABLE "{tmpTable}" RENAME TO "{tableName}";'
|
||
)
|
||
return True
|
||
|
||
@classmethod
|
||
async def executeDB(cls, command: str) -> bool:
|
||
"""执行自定义SQL
|
||
|
||
Args:
|
||
command (str): SQL语句
|
||
|
||
Returns:
|
||
bool: 是否执行成功
|
||
"""
|
||
if not command:
|
||
return False
|
||
|
||
try:
|
||
async with cls._transaction():
|
||
await cls.m_pDB.execute(command)
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
@classmethod
|
||
async def insert(cls, tableName: str, data: dict) -> bool:
|
||
"""
|
||
插入数据
|
||
|
||
Args:
|
||
tableName: 表名
|
||
data: 要插入的数据字典，键为字段名，值为字段值
|
||
|
||
Returns:
|
||
bool: 是否执行成功
|
||
"""
|
||
if not data:
|
||
return False
|
||
|
||
try:
|
||
# 构建参数化查询
|
||
columns = ", ".join(f'"{k}"' for k in data.keys())
|
||
placeholders = ", ".join("?" for _ in data.keys())
|
||
values = list(data.values())
|
||
|
||
sql = f'INSERT INTO "{tableName}" ({columns}) VALUES ({placeholders})'
|
||
|
||
async with cls._transaction():
|
||
await cls.m_pDB.execute(sql, values)
|
||
return True
|
||
except Exception as e:
|
||
logger.debug("真寻农场插入数据失败！", e=e)
|
||
return False
|
||
|
||
@classmethod
|
||
async def batch_insert(cls, tableName: str, data_list: list) -> bool:
|
||
"""
|
||
批量插入数据
|
||
|
||
Args:
|
||
tableName: 表名
|
||
data_list: 要插入的数据字典列表
|
||
|
||
Returns:
|
||
bool: 是否执行成功
|
||
"""
|
||
if not data_list:
|
||
return False
|
||
|
||
try:
|
||
# 使用第一个字典的键作为所有记录的字段
|
||
columns = ", ".join(f'"{k}"' for k in data_list[0].keys())
|
||
placeholders = ", ".join("?" for _ in data_list[0].keys())
|
||
|
||
sql = f'INSERT INTO "{tableName}" ({columns}) VALUES ({placeholders})'
|
||
|
||
async with cls._transaction():
|
||
await cls.m_pDB.executemany(
|
||
sql, [list(data.values()) for data in data_list]
|
||
)
|
||
return True
|
||
except Exception as e:
|
||
logger.debug("真寻农场批量插入数据失败！", e=e)
|
||
return False
|
||
|
||
@classmethod
|
||
async def select(
|
||
cls,
|
||
tableName: str,
|
||
columns: list[Any] | None = None,
|
||
where: dict[str, Any] | None = None,
|
||
order_by: str | None = None,
|
||
limit: int | None = None,
|
||
) -> list[dict]:
|
||
"""
|
||
查询数据
|
||
|
||
Args:
|
||
tableName: 表名
|
||
columns: 要查询的字段列表，None表示所有字段
|
||
where: 查询条件字典
|
||
order_by: 排序字段
|
||
limit: 限制返回记录数
|
||
|
||
Returns:
|
||
list: 查询结果列表，每个元素是一个字典
|
||
"""
|
||
try:
|
||
# 构建SELECT部分
|
||
if columns:
|
||
select_clause = ", ".join(f'"{col}"' for col in columns)
|
||
else:
|
||
select_clause = "*"
|
||
|
||
sql = f'SELECT {select_clause} FROM "{tableName}"'
|
||
|
||
# 构建WHERE部分
|
||
params = []
|
||
if where:
|
||
where_conditions = []
|
||
for key, value in where.items():
|
||
if isinstance(value, (list, tuple)):
|
||
# 处理IN查询
|
||
placeholders = ", ".join("?" for _ in value)
|
||
where_conditions.append(f'"{key}" IN ({placeholders})')
|
||
params.extend(value)
|
||
else:
|
||
where_conditions.append(f'"{key}" = ?')
|
||
params.append(value)
|
||
|
||
if where_conditions:
|
||
sql += " WHERE " + " AND ".join(where_conditions)
|
||
|
||
# 构建ORDER BY部分
|
||
if order_by:
|
||
sql += f" ORDER BY {order_by}"
|
||
|
||
# 构建LIMIT部分
|
||
if limit:
|
||
sql += f" LIMIT {limit}"
|
||
|
||
cursor = await cls.m_pDB.execute(sql, params)
|
||
rows = await cursor.fetchall()
|
||
|
||
# 转换为字典列表
|
||
result = []
|
||
for row in rows:
|
||
result.append(dict(row))
|
||
|
||
return result
|
||
except Exception as e:
|
||
logger.debug("真寻农场查询数据失败！", e=e)
|
||
return []
|
||
|
||
@classmethod
|
||
async def update(cls, tableName: str, data: dict, where: dict) -> bool:
|
||
"""
|
||
更新数据
|
||
|
||
Args:
|
||
tableName: 表名
|
||
data: 要更新的数据字典
|
||
where: 更新条件字典
|
||
|
||
Returns:
|
||
bool: 是否执行成功
|
||
"""
|
||
if not data:
|
||
return False
|
||
|
||
if not where:
|
||
return False
|
||
|
||
try:
|
||
# 构建SET部分
|
||
set_conditions = []
|
||
params = []
|
||
for key, value in data.items():
|
||
set_conditions.append(f'"{key}" = ?')
|
||
params.append(value)
|
||
|
||
# 构建WHERE部分
|
||
where_conditions = []
|
||
for key, value in where.items():
|
||
where_conditions.append(f'"{key}" = ?')
|
||
params.append(value)
|
||
|
||
sql = f'UPDATE "{tableName}" SET {", ".join(set_conditions)} WHERE {" AND ".join(where_conditions)}'
|
||
|
||
async with cls._transaction():
|
||
cursor = await cls.m_pDB.execute(sql, params)
|
||
# 检查是否影响了行
|
||
return cursor.rowcount > 0
|
||
except Exception as e:
|
||
logger.debug("真寻农场更新数据失败！", e=e)
|
||
return False
|
||
|
||
@classmethod
|
||
async def delete(cls, tableName: str, where: dict) -> bool:
|
||
"""
|
||
删除数据
|
||
|
||
Args:
|
||
tableName: 表名
|
||
where: 删除条件字典
|
||
|
||
Returns:
|
||
bool: 是否执行成功
|
||
"""
|
||
if not where:
|
||
return False
|
||
|
||
try:
|
||
# 构建WHERE部分
|
||
where_conditions = []
|
||
params = []
|
||
for key, value in where.items():
|
||
where_conditions.append(f'"{key}" = ?')
|
||
params.append(value)
|
||
|
||
sql = f'DELETE FROM "{tableName}" WHERE {" AND ".join(where_conditions)}'
|
||
|
||
async with cls._transaction():
|
||
cursor = await cls.m_pDB.execute(sql, params)
|
||
# 检查是否影响了行
|
||
return cursor.rowcount > 0
|
||
except Exception as e:
|
||
logger.debug("真寻农场删除数据失败！", e=e)
|
||
return False
|
||
|
||
@classmethod
|
||
async def exists(cls, tableName: str, where: dict) -> bool:
|
||
"""
|
||
检查记录是否存在
|
||
|
||
Args:
|
||
tableName: 表名
|
||
where: 查询条件字典
|
||
|
||
Returns:
|
||
bool: 是否存在符合条件的记录
|
||
"""
|
||
try:
|
||
result = await cls.select(tableName, columns=["1"], where=where, limit=1)
|
||
return len(result) > 0
|
||
except Exception as e:
|
||
logger.debug("真寻农场检查数据失败！", e=e)
|
||
return False
|
||
|
||
@classmethod
|
||
async def count(cls, tableName: str, where: dict = {}) -> int:
|
||
"""
|
||
统计记录数量
|
||
|
||
Args:
|
||
tableName: 表名
|
||
where: 查询条件字典
|
||
|
||
Returns:
|
||
int: 记录数量
|
||
"""
|
||
try:
|
||
# 构建WHERE部分
|
||
sql = f'SELECT COUNT(*) as count FROM "{tableName}"'
|
||
params = []
|
||
|
||
if where:
|
||
where_conditions = []
|
||
for key, value in where.items():
|
||
where_conditions.append(f'"{key}" = ?')
|
||
params.append(value)
|
||
|
||
sql += " WHERE " + " AND ".join(where_conditions)
|
||
|
||
cursor = await cls.m_pDB.execute(sql, params)
|
||
row = await cursor.fetchone()
|
||
return row["count"] if row else 0
|
||
except Exception as e:
|
||
logger.debug("真寻农场统计数据失败！", e=e)
|
||
return 0
|
||
|
||
|
||
g_pSqlManager = CSqlManager()
|