mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-10-14 00:20:06 +08:00
bugfix:修复notice事件扩散问题 (#2132)
* bugfix:修复notice事件扩散问题 * 优化并发调度 * bugfix:修复签到样式 * bugfix：功能调用统计修复 * bugfix:修复私聊时功能调用统计显示已退群问题 * 提高插件适配兼容性 * 优化发送队列 * 修改权限检查设计 * 继续修改权限检查设计 * 完善权限检查设计 * 优化sqlite配置 * 优化数据库初始化 * 代码整理，无用代码清理 * bugfix:修复启动时数据库校验问题 * bugfix:修复预算裁剪过于激进问题
This commit is contained in:
Copaan
@@ -1,8 +1,10 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
from typing import Literal
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import aiofiles
|
||||
@@ -27,8 +29,12 @@ from .config import (
|
||||
prompt,
|
||||
)
|
||||
from .exceptions import DbConnectError, DbUrlIsNode
|
||||
from .schema_guard import repair_safe_schema_drift
|
||||
from .schema_ops import SchemaOpRisk, normalize_schema_ops
|
||||
from .utils import with_db_timeout
|
||||
|
||||
Dialect = Literal["sqlite", "postgres", "mysql", "unknown"]
|
||||
|
||||
MODELS = db_model.models
|
||||
SCRIPT_METHOD = db_model.script_method
|
||||
|
||||
@@ -47,7 +53,66 @@ __all__ = [
|
||||
|
||||
driver = nonebot.get_driver()
|
||||
|
||||
_SCRIPT_HASH_FILE = Path() / "data" / ".db_script_hash"
|
||||
_SCRIPT_HASH_DIR = Path() / "data" / ".db_script_hashes"
|
||||
_TRUE_VALUES = {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _connection_dialect() -> Dialect:
|
||||
try:
|
||||
connection = Tortoise.get_connection("default")
|
||||
capabilities = getattr(connection, "capabilities", None)
|
||||
raw = str(getattr(capabilities, "dialect", "") or "").lower()
|
||||
if raw.startswith("sqlite"):
|
||||
return "sqlite"
|
||||
if raw.startswith("postgres"):
|
||||
return "postgres"
|
||||
if raw.startswith("mysql"):
|
||||
return "mysql"
|
||||
except Exception:
|
||||
pass
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _allow_guarded_schema_ops() -> bool:
|
||||
"""Whether startup may run guarded SchemaOp migrations.
|
||||
|
||||
Safe SchemaOps are limited to non-destructive changes such as adding nullable
|
||||
columns and non-unique indexes. Guarded operations may rename, drop, or alter
|
||||
columns, so keep them opt-in to avoid damaging existing databases during
|
||||
normal startup.
|
||||
"""
|
||||
return os.getenv("DB_SCHEMA_RUN_GUARDED_OPS", "").strip().lower() in _TRUE_VALUES
|
||||
|
||||
|
||||
def _extract_alter_table_name(sql: str) -> str | None:
|
||||
match = re.match(r"ALTER\s+TABLE\s+[`\"]?(\w+)[`\"]?", sql, re.IGNORECASE)
|
||||
return match.group(1) if match else None
|
||||
|
||||
|
||||
def _extract_create_index_table_name(sql: str) -> str | None:
|
||||
match = re.search(r"\bON\s+[`\"]?(\w+)[`\"]?\s*\(", sql, re.IGNORECASE)
|
||||
return match.group(1) if match else None
|
||||
|
||||
|
||||
def _db_script_hash_file(script_fingerprint: str) -> Path:
|
||||
parsed = urlparse(BotConfig.db_url or "")
|
||||
dialect = parsed.scheme or "unknown"
|
||||
if dialect == "sqlite":
|
||||
db_identity = str(Path(parsed.path).resolve())
|
||||
else:
|
||||
db_identity = f"{parsed.hostname or ''}:{parsed.port or ''}{parsed.path}"
|
||||
db_hash = hashlib.md5(
|
||||
json.dumps(
|
||||
{
|
||||
"dialect": dialect,
|
||||
"db": db_identity,
|
||||
"script": script_fingerprint,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
sort_keys=True,
|
||||
).encode()
|
||||
).hexdigest()
|
||||
return _SCRIPT_HASH_DIR / f"{db_hash}.json"
|
||||
|
||||
|
||||
def get_config() -> dict:
|
||||
@@ -131,20 +196,39 @@ async def init():
|
||||
f"<u><y>{len(db_model.script_method)}</y></u> 个..."
|
||||
)
|
||||
sql_list = []
|
||||
allow_guarded_ops = _allow_guarded_schema_ops()
|
||||
for module, func in db_model.script_method:
|
||||
try:
|
||||
sql = await func() if is_coroutine_callable(func) else func()
|
||||
if sql:
|
||||
sql_list += sql
|
||||
items = await func() if is_coroutine_callable(func) else func()
|
||||
if not items:
|
||||
continue
|
||||
for item in items:
|
||||
if not isinstance(item, str):
|
||||
if item.risk == SchemaOpRisk.MANUAL:
|
||||
logger.debug(f"{module} 跳过手动迁移动作: {item}")
|
||||
continue
|
||||
if (
|
||||
item.risk == SchemaOpRisk.GUARDED
|
||||
and not allow_guarded_ops
|
||||
):
|
||||
logger.debug(f"{module} 跳过受保护迁移动作: {item}")
|
||||
continue
|
||||
if item.risk != SchemaOpRisk.SAFE and not allow_guarded_ops:
|
||||
logger.debug(f"{module} 跳过未知风险迁移动作: {item}")
|
||||
continue
|
||||
sql_list += normalize_schema_ops([item], _connection_dialect())
|
||||
except Exception as e:
|
||||
logger.debug(f"{module} 执行SCRIPT_METHOD方法出错...", e=e)
|
||||
if sql_list:
|
||||
fingerprint = hashlib.md5(
|
||||
json.dumps(sorted(sql_list), ensure_ascii=False).encode()
|
||||
).hexdigest()
|
||||
script_hash_file = _db_script_hash_file(fingerprint)
|
||||
need_run = not (
|
||||
_SCRIPT_HASH_FILE.exists()
|
||||
and _SCRIPT_HASH_FILE.read_text(encoding="utf-8").strip()
|
||||
script_hash_file.exists()
|
||||
and json.loads(script_hash_file.read_text(encoding="utf-8")).get(
|
||||
"script_fingerprint"
|
||||
)
|
||||
== fingerprint
|
||||
)
|
||||
if need_run:
|
||||
@@ -186,12 +270,16 @@ async def init():
|
||||
|
||||
for sql in sql_list:
|
||||
# 对于 ALTER TABLE 操作，先检查表是否存在
|
||||
if sql.strip().upper().startswith("ALTER TABLE"):
|
||||
match = re.match(
|
||||
r"ALTER\s+TABLE\s+(\w+)", sql, re.IGNORECASE
|
||||
)
|
||||
if match:
|
||||
table_name = match.group(1)
|
||||
sql_upper = sql.strip().upper()
|
||||
if sql_upper.startswith("ALTER TABLE"):
|
||||
table_name = _extract_alter_table_name(sql)
|
||||
if table_name:
|
||||
if not await table_exists(table_name):
|
||||
logger.debug(f"跳过SQL（表不存在）: {sql}")
|
||||
continue
|
||||
elif sql_upper.startswith("CREATE INDEX"):
|
||||
table_name = _extract_create_index_table_name(sql)
|
||||
if table_name:
|
||||
if not await table_exists(table_name):
|
||||
logger.debug(f"跳过SQL（表不存在）: {sql}")
|
||||
continue
|
||||
@@ -236,13 +324,31 @@ async def init():
|
||||
except Exception as e:
|
||||
logger.debug(f"执行SQL: {sql} 错误...", e=e)
|
||||
logger.debug("SCRIPT_METHOD方法执行完毕!")
|
||||
_SCRIPT_HASH_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
_SCRIPT_HASH_FILE.write_text(fingerprint, encoding="utf-8")
|
||||
script_hash_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
script_hash_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"dialect": urlparse(BotConfig.db_url or "").scheme,
|
||||
"db_url_hash": hashlib.md5(
|
||||
(BotConfig.db_url or "").encode()
|
||||
).hexdigest(),
|
||||
"script_fingerprint": fingerprint,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
else:
|
||||
logger.debug("迁移脚本无变化，跳过执行")
|
||||
# Tortoise may emit column comments/index SQL during generate_schemas().
|
||||
# On existing databases with newly added nullable fields, PostgreSQL can
|
||||
# fail before the post-generate SchemaGuard gets a chance to repair drift.
|
||||
await repair_safe_schema_drift()
|
||||
logger.debug("开始生成数据库表结构...")
|
||||
await Tortoise.generate_schemas()
|
||||
logger.debug("数据库表结构生成完毕!")
|
||||
await repair_safe_schema_drift()
|
||||
logger.info("Database loaded successfully!")
|
||||
except Exception as e:
|
||||
raise DbConnectError(f"数据库连接错误... e:{e}") from e
|
||||
|
||||
@@ -24,7 +24,8 @@ MYSQL_CONFIG = {
|
||||
|
||||
SQLITE_CONFIG = {
|
||||
"journal_mode": "WAL", # 提高并发写入性能
|
||||
"timeout": 30, # 锁等待超时（可选）
|
||||
"busy_timeout": 30000, # SQLite 锁等待超时，单位毫秒
|
||||
"foreign_keys": "ON",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,593 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
from tortoise import Tortoise
|
||||
from tortoise.exceptions import OperationalError
|
||||
|
||||
from zhenxun.services.log import logger
|
||||
|
||||
from .config import DB_TIMEOUT_SECONDS, LOG_COMMAND
|
||||
|
||||
Dialect = Literal["sqlite", "postgres", "mysql", "unknown"]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SchemaGuardResult:
|
||||
checked_tables: int = 0
|
||||
repaired_columns: int = 0
|
||||
repaired_indexes: int = 0
|
||||
skipped_columns: int = 0
|
||||
skipped_indexes: int = 0
|
||||
type_mismatches: int = 0
|
||||
warnings: int = 0
|
||||
drift: list[dict[str, Any]] | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ColumnInfo:
|
||||
name: str
|
||||
data_type: str
|
||||
nullable: bool | None = None
|
||||
default: str | None = None
|
||||
|
||||
|
||||
def _quote_identifier(identifier: str, dialect: Dialect) -> str:
|
||||
escaped = identifier.replace('"', '""')
|
||||
if dialect == "mysql":
|
||||
return f"`{identifier.replace('`', '``')}`"
|
||||
return f'"{escaped}"'
|
||||
|
||||
|
||||
def _connection_dialect(connection: Any) -> Dialect:
|
||||
capabilities = getattr(connection, "capabilities", None)
|
||||
raw = str(getattr(capabilities, "dialect", "") or "").lower()
|
||||
if raw.startswith("sqlite"):
|
||||
return "sqlite"
|
||||
if raw.startswith("postgres"):
|
||||
return "postgres"
|
||||
if raw.startswith("mysql"):
|
||||
return "mysql"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _is_safe_missing_field(field: Any) -> bool:
|
||||
if getattr(field, "pk", False):
|
||||
return False
|
||||
if getattr(field, "generated", False):
|
||||
return False
|
||||
if bool(getattr(field, "null", False)):
|
||||
return True
|
||||
default = getattr(field, "default", None)
|
||||
return default is not None
|
||||
|
||||
|
||||
def _field_sql_type(field: Any, dialect: Dialect) -> str | None:
|
||||
if dialect != "unknown" and hasattr(field, "get_for_dialect"):
|
||||
with contextlib.suppress(Exception):
|
||||
value = field.get_for_dialect(dialect, "SQL_TYPE")
|
||||
if value:
|
||||
return str(value)
|
||||
value = getattr(field, "SQL_TYPE", None)
|
||||
return str(value) if value else None
|
||||
|
||||
|
||||
def _is_db_field(field: Any, dialect: Dialect) -> bool:
|
||||
if getattr(field, "virtual", False):
|
||||
return False
|
||||
return _field_sql_type(field, dialect) is not None
|
||||
|
||||
|
||||
def _field_default_sql(field: Any, dialect: Dialect) -> str:
|
||||
default = getattr(field, "default", None)
|
||||
if default is None or callable(default):
|
||||
return ""
|
||||
if isinstance(default, bool):
|
||||
if dialect == "postgres":
|
||||
return " DEFAULT TRUE" if default else " DEFAULT FALSE"
|
||||
return " DEFAULT 1" if default else " DEFAULT 0"
|
||||
if isinstance(default, int | float):
|
||||
return f" DEFAULT {default}"
|
||||
if isinstance(default, str):
|
||||
escaped = default.replace("'", "''")
|
||||
return f" DEFAULT '{escaped}'"
|
||||
return ""
|
||||
|
||||
|
||||
def _missing_column_sql(
|
||||
table: str,
|
||||
column: str,
|
||||
field: Any,
|
||||
dialect: Dialect,
|
||||
) -> str | None:
|
||||
sql_type = _field_sql_type(field, dialect)
|
||||
if not sql_type:
|
||||
return None
|
||||
table_sql = _quote_identifier(table, dialect)
|
||||
column_sql = _quote_identifier(column, dialect)
|
||||
null_sql = "" if bool(getattr(field, "null", False)) else " NOT NULL"
|
||||
default_sql = _field_default_sql(field, dialect)
|
||||
if not bool(getattr(field, "null", False)) and not default_sql:
|
||||
return None
|
||||
if dialect == "postgres":
|
||||
return (
|
||||
f"ALTER TABLE {table_sql} ADD COLUMN IF NOT EXISTS "
|
||||
f"{column_sql} {sql_type}{default_sql}{null_sql}"
|
||||
)
|
||||
return (
|
||||
f"ALTER TABLE {table_sql} ADD COLUMN "
|
||||
f"{column_sql} {sql_type}{default_sql}{null_sql}"
|
||||
)
|
||||
|
||||
|
||||
async def _table_columns(
|
||||
connection: Any, table: str, dialect: Dialect
|
||||
) -> dict[str, ColumnInfo] | None:
|
||||
if dialect == "sqlite":
|
||||
rows = await connection.execute_query_dict(
|
||||
f"PRAGMA table_xinfo({_quote_identifier(table, dialect)})"
|
||||
)
|
||||
if not rows:
|
||||
return None
|
||||
return {
|
||||
str(row.get("name")): ColumnInfo(
|
||||
name=str(row.get("name")),
|
||||
data_type=str(row.get("type") or ""),
|
||||
nullable=not bool(row.get("notnull")),
|
||||
default=str(row.get("dflt_value"))
|
||||
if row.get("dflt_value") is not None
|
||||
else None,
|
||||
)
|
||||
for row in rows
|
||||
if row.get("name")
|
||||
}
|
||||
if dialect == "postgres":
|
||||
rows = await connection.execute_query_dict(
|
||||
"SELECT column_name, data_type, is_nullable, column_default "
|
||||
"FROM information_schema.columns "
|
||||
"WHERE table_schema = current_schema() AND table_name = $1",
|
||||
[table],
|
||||
)
|
||||
if not rows:
|
||||
return None
|
||||
return {
|
||||
str(row.get("column_name")): ColumnInfo(
|
||||
name=str(row.get("column_name")),
|
||||
data_type=str(row.get("data_type") or ""),
|
||||
nullable=str(row.get("is_nullable") or "").upper() == "YES",
|
||||
default=str(row.get("column_default"))
|
||||
if row.get("column_default") is not None
|
||||
else None,
|
||||
)
|
||||
for row in rows
|
||||
if row.get("column_name")
|
||||
}
|
||||
if dialect == "mysql":
|
||||
rows = await connection.execute_query_dict(
|
||||
"SELECT COLUMN_NAME AS column_name, COLUMN_TYPE AS column_type, "
|
||||
"IS_NULLABLE AS is_nullable, COLUMN_DEFAULT AS column_default "
|
||||
"FROM INFORMATION_SCHEMA.COLUMNS "
|
||||
"WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s",
|
||||
[table],
|
||||
)
|
||||
if not rows:
|
||||
return None
|
||||
return {
|
||||
str(row.get("column_name")): ColumnInfo(
|
||||
name=str(row.get("column_name")),
|
||||
data_type=str(row.get("column_type") or ""),
|
||||
nullable=str(row.get("is_nullable") or "").upper() == "YES",
|
||||
default=str(row.get("column_default"))
|
||||
if row.get("column_default") is not None
|
||||
else None,
|
||||
)
|
||||
for row in rows
|
||||
if row.get("column_name")
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_type(value: str, dialect: Dialect) -> str:
|
||||
text = value.lower().strip()
|
||||
if not text:
|
||||
return ""
|
||||
if dialect == "sqlite":
|
||||
if "int" in text:
|
||||
return "integer"
|
||||
if any(part in text for part in ("char", "clob", "text", "varchar")):
|
||||
return "text"
|
||||
if any(part in text for part in ("real", "floa", "doub")):
|
||||
return "real"
|
||||
if "blob" in text:
|
||||
return "blob"
|
||||
if "bool" in text:
|
||||
return "integer"
|
||||
if "timestamp" in text or "datetime" in text:
|
||||
return "text"
|
||||
return text
|
||||
if dialect == "postgres":
|
||||
text = text.replace("character varying", "varchar")
|
||||
if text.startswith("timestamp"):
|
||||
return "timestamp"
|
||||
if text in {"boolean", "bool"}:
|
||||
return "boolean"
|
||||
return text
|
||||
if dialect == "mysql":
|
||||
if text.startswith("tinyint(1)") or text == "bool" or text == "boolean":
|
||||
return "boolean"
|
||||
if text.startswith("datetime") or text.startswith("timestamp"):
|
||||
return "datetime"
|
||||
return text
|
||||
return text
|
||||
|
||||
|
||||
def _field_type_compatible(field: Any, column: ColumnInfo, dialect: Dialect) -> bool:
|
||||
expected = _field_sql_type(field, dialect)
|
||||
if not expected:
|
||||
return True
|
||||
expected_norm = _normalize_type(expected, dialect)
|
||||
actual_norm = _normalize_type(column.data_type, dialect)
|
||||
if not expected_norm or not actual_norm:
|
||||
return True
|
||||
if expected_norm == actual_norm:
|
||||
return True
|
||||
if dialect == "sqlite":
|
||||
# SQLite affinity is intentionally loose; varchar/text and bool/integer are
|
||||
# compatible enough for startup validation.
|
||||
compatible = {
|
||||
("varchar", "text"),
|
||||
("text", "varchar"),
|
||||
("boolean", "integer"),
|
||||
("integer", "boolean"),
|
||||
}
|
||||
return (expected_norm, actual_norm) in compatible
|
||||
return False
|
||||
|
||||
|
||||
def _index_name(table: str, columns: tuple[str, ...]) -> str:
|
||||
raw = f"idx_{table}_{'_'.join(columns)}"
|
||||
return raw[:62]
|
||||
|
||||
|
||||
async def _table_indexes(
|
||||
connection: Any, table: str, dialect: Dialect
|
||||
) -> set[tuple[str, ...]]:
|
||||
indexes: set[tuple[str, ...]] = set()
|
||||
if dialect == "sqlite":
|
||||
rows = await connection.execute_query_dict(
|
||||
f"PRAGMA index_list({_quote_identifier(table, dialect)})"
|
||||
)
|
||||
for row in rows:
|
||||
if bool(row.get("unique")):
|
||||
continue
|
||||
index_name = row.get("name")
|
||||
if not index_name:
|
||||
continue
|
||||
info = await connection.execute_query_dict(
|
||||
f"PRAGMA index_info({_quote_identifier(str(index_name), dialect)})"
|
||||
)
|
||||
columns = tuple(
|
||||
str(item.get("name"))
|
||||
for item in sorted(info, key=lambda item: int(item.get("seqno") or 0))
|
||||
if item.get("name")
|
||||
)
|
||||
if columns:
|
||||
indexes.add(columns)
|
||||
return indexes
|
||||
if dialect == "postgres":
|
||||
rows = await connection.execute_query_dict(
|
||||
"SELECT indexname, indexdef FROM pg_indexes "
|
||||
"WHERE schemaname = current_schema() AND tablename = $1",
|
||||
[table],
|
||||
)
|
||||
for row in rows:
|
||||
indexdef = str(row.get("indexdef") or "")
|
||||
if " UNIQUE INDEX " in indexdef.upper():
|
||||
continue
|
||||
start = indexdef.rfind("(")
|
||||
end = indexdef.rfind(")")
|
||||
if start < 0 or end <= start:
|
||||
continue
|
||||
columns = tuple(
|
||||
part.strip().strip('"')
|
||||
for part in indexdef[start + 1 : end].split(",")
|
||||
if part.strip()
|
||||
)
|
||||
if columns:
|
||||
indexes.add(columns)
|
||||
return indexes
|
||||
if dialect == "mysql":
|
||||
rows = await connection.execute_query_dict(
|
||||
"SELECT INDEX_NAME, COLUMN_NAME, SEQ_IN_INDEX, NON_UNIQUE "
|
||||
"FROM INFORMATION_SCHEMA.STATISTICS "
|
||||
"WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s",
|
||||
[table],
|
||||
)
|
||||
grouped: dict[str, list[tuple[int, str]]] = {}
|
||||
unique_names: set[str] = set()
|
||||
for row in rows:
|
||||
name = str(row.get("INDEX_NAME") or "")
|
||||
column = str(row.get("COLUMN_NAME") or "")
|
||||
if not name or not column:
|
||||
continue
|
||||
if int(row.get("NON_UNIQUE") or 0) == 0:
|
||||
unique_names.add(name)
|
||||
continue
|
||||
grouped.setdefault(name, []).append(
|
||||
(int(row.get("SEQ_IN_INDEX") or 0), column)
|
||||
)
|
||||
for name, items in grouped.items():
|
||||
if name in unique_names:
|
||||
continue
|
||||
columns = tuple(column for _, column in sorted(items))
|
||||
if columns:
|
||||
indexes.add(columns)
|
||||
return indexes
|
||||
return indexes
|
||||
|
||||
|
||||
def _index_sql(table: str, columns: tuple[str, ...], dialect: Dialect) -> str | None:
|
||||
if not columns:
|
||||
return None
|
||||
table_sql = _quote_identifier(table, dialect)
|
||||
index_sql = _quote_identifier(_index_name(table, columns), dialect)
|
||||
columns_sql = ", ".join(_quote_identifier(column, dialect) for column in columns)
|
||||
if dialect in {"sqlite", "postgres"}:
|
||||
return (
|
||||
f"CREATE INDEX IF NOT EXISTS {index_sql} " f"ON {table_sql}({columns_sql})"
|
||||
)
|
||||
if dialect == "mysql":
|
||||
return f"CREATE INDEX {index_sql} ON {table_sql}({columns_sql})"
|
||||
return None
|
||||
|
||||
|
||||
def _meta_index_columns(index: Any) -> tuple[str, ...]:
|
||||
if isinstance(index, str):
|
||||
return (index,)
|
||||
if isinstance(index, list | tuple):
|
||||
return tuple(str(column) for column in index if column)
|
||||
fields = getattr(index, "fields", None)
|
||||
if fields:
|
||||
return tuple(str(column) for column in fields if column)
|
||||
return ()
|
||||
|
||||
|
||||
def _empty_drift(table: str) -> dict[str, Any]:
|
||||
return {
|
||||
"table": table,
|
||||
"missing_columns": [],
|
||||
"extra_columns": [],
|
||||
"type_mismatches": [],
|
||||
"missing_indexes": [],
|
||||
"unsafe_changes": [],
|
||||
}
|
||||
|
||||
|
||||
async def _write_schema_report(result: SchemaGuardResult) -> None:
|
||||
report = {
|
||||
"created_at": datetime.now().isoformat(timespec="seconds"),
|
||||
"summary": {
|
||||
"checked_tables": result.checked_tables,
|
||||
"repaired_columns": result.repaired_columns,
|
||||
"repaired_indexes": result.repaired_indexes,
|
||||
"skipped_columns": result.skipped_columns,
|
||||
"skipped_indexes": result.skipped_indexes,
|
||||
"type_mismatches": result.type_mismatches,
|
||||
"warnings": result.warnings,
|
||||
},
|
||||
"tables": result.drift or [],
|
||||
}
|
||||
path = Path() / "data" / "db" / "schema_report.json"
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
await asyncio.to_thread(
|
||||
path.write_text,
|
||||
json.dumps(report, ensure_ascii=False, indent=2),
|
||||
"utf-8",
|
||||
)
|
||||
|
||||
|
||||
async def repair_safe_schema_drift() -> SchemaGuardResult:
|
||||
"""Repair low-risk missing columns after Tortoise schema creation.
|
||||
|
||||
This guard intentionally avoids type changes, constraints, unique indexes, foreign
|
||||
keys, and SQLite table rebuilds. It is a startup-only safety net for model/schema
|
||||
drift caused by copied databases or skipped legacy script hashes.
|
||||
"""
|
||||
result = SchemaGuardResult()
|
||||
result.drift = []
|
||||
connection = Tortoise.get_connection("default")
|
||||
dialect = _connection_dialect(connection)
|
||||
if dialect == "unknown":
|
||||
logger.debug("SchemaGuard 跳过未知数据库方言", LOG_COMMAND)
|
||||
return result
|
||||
|
||||
app = Tortoise.apps.get("models", {})
|
||||
for model in app.values():
|
||||
meta = getattr(model, "_meta", None)
|
||||
table = getattr(meta, "db_table", None) or getattr(meta, "table", None)
|
||||
if not table:
|
||||
continue
|
||||
try:
|
||||
columns = await _table_columns(connection, table, dialect)
|
||||
except Exception as exc:
|
||||
result.warnings += 1
|
||||
logger.debug(f"SchemaGuard 检查表 {table} 失败", LOG_COMMAND, e=exc)
|
||||
continue
|
||||
if columns is None:
|
||||
continue
|
||||
result.checked_tables += 1
|
||||
drift = _empty_drift(table)
|
||||
fields_map = getattr(meta, "fields_map", {}) or {}
|
||||
expected_sources: set[str] = set()
|
||||
for field_name, field in fields_map.items():
|
||||
if not _is_db_field(field, dialect):
|
||||
continue
|
||||
source = str(getattr(field, "source_field", None) or field_name)
|
||||
expected_sources.add(source)
|
||||
if source in columns:
|
||||
if not _field_type_compatible(field, columns[source], dialect):
|
||||
result.type_mismatches += 1
|
||||
drift["type_mismatches"].append(
|
||||
{
|
||||
"column": source,
|
||||
"expected": _field_sql_type(field, dialect),
|
||||
"actual": columns[source].data_type,
|
||||
}
|
||||
)
|
||||
continue
|
||||
drift["missing_columns"].append(source)
|
||||
if not _is_safe_missing_field(field):
|
||||
result.skipped_columns += 1
|
||||
drift["unsafe_changes"].append(
|
||||
{
|
||||
"kind": "missing_required_column",
|
||||
"column": source,
|
||||
}
|
||||
)
|
||||
logger.debug(
|
||||
f"SchemaGuard 跳过非低风险缺字段: {table}.{source}",
|
||||
LOG_COMMAND,
|
||||
)
|
||||
continue
|
||||
sql = _missing_column_sql(table, source, field, dialect)
|
||||
if not sql:
|
||||
result.skipped_columns += 1
|
||||
logger.debug(
|
||||
f"SchemaGuard 无法生成补字段 SQL: {table}.{source}",
|
||||
LOG_COMMAND,
|
||||
)
|
||||
continue
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
connection.execute_query_dict(sql),
|
||||
timeout=DB_TIMEOUT_SECONDS,
|
||||
)
|
||||
columns[source] = ColumnInfo(name=source, data_type="")
|
||||
result.repaired_columns += 1
|
||||
logger.info(f"SchemaGuard 已补齐字段: {table}.{source}", LOG_COMMAND)
|
||||
except OperationalError as exc:
|
||||
err = str(exc).lower()
|
||||
if any(
|
||||
text in err
|
||||
for text in ("duplicate column", "already exists", "已存在")
|
||||
):
|
||||
columns[source] = ColumnInfo(name=source, data_type="")
|
||||
continue
|
||||
result.warnings += 1
|
||||
logger.warning(
|
||||
f"SchemaGuard 补齐字段失败: {table}.{source}",
|
||||
LOG_COMMAND,
|
||||
e=exc,
|
||||
)
|
||||
except Exception as exc:
|
||||
result.warnings += 1
|
||||
logger.warning(
|
||||
f"SchemaGuard 补齐字段失败: {table}.{source}",
|
||||
LOG_COMMAND,
|
||||
e=exc,
|
||||
)
|
||||
drift["extra_columns"] = sorted(
|
||||
column for column in columns.keys() if column not in expected_sources
|
||||
)
|
||||
try:
|
||||
existing_indexes = await _table_indexes(connection, table, dialect)
|
||||
except Exception as exc:
|
||||
result.warnings += 1
|
||||
logger.debug(f"SchemaGuard 检查索引 {table} 失败", LOG_COMMAND, e=exc)
|
||||
existing_indexes = set()
|
||||
indexes = getattr(meta, "indexes", ()) or ()
|
||||
for index in indexes:
|
||||
index_columns = _meta_index_columns(index)
|
||||
if not index_columns or index_columns in existing_indexes:
|
||||
continue
|
||||
drift["missing_indexes"].append(list(index_columns))
|
||||
if not all(column in columns for column in index_columns):
|
||||
result.skipped_indexes += 1
|
||||
logger.debug(
|
||||
f"SchemaGuard 跳过缺字段索引: {table}.{index_columns}",
|
||||
LOG_COMMAND,
|
||||
)
|
||||
continue
|
||||
sql = _index_sql(table, index_columns, dialect)
|
||||
if not sql:
|
||||
result.skipped_indexes += 1
|
||||
continue
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
connection.execute_query_dict(sql),
|
||||
timeout=DB_TIMEOUT_SECONDS,
|
||||
)
|
||||
existing_indexes.add(index_columns)
|
||||
result.repaired_indexes += 1
|
||||
logger.debug(
|
||||
f"SchemaGuard 已补齐索引: {table}.{index_columns}",
|
||||
LOG_COMMAND,
|
||||
)
|
||||
except OperationalError as exc:
|
||||
err = str(exc).lower()
|
||||
if any(
|
||||
text in err for text in ("already exists", "duplicate", "已存在")
|
||||
):
|
||||
existing_indexes.add(index_columns)
|
||||
continue
|
||||
result.warnings += 1
|
||||
logger.warning(
|
||||
f"SchemaGuard 补齐索引失败: {table}.{index_columns}",
|
||||
LOG_COMMAND,
|
||||
e=exc,
|
||||
)
|
||||
if any(
|
||||
drift[key]
|
||||
for key in (
|
||||
"missing_columns",
|
||||
"extra_columns",
|
||||
"type_mismatches",
|
||||
"missing_indexes",
|
||||
"unsafe_changes",
|
||||
)
|
||||
):
|
||||
result.drift.append(drift)
|
||||
|
||||
logger.info(
|
||||
"SchemaGuard 完成: "
|
||||
f"checked={result.checked_tables}, "
|
||||
f"repaired_columns={result.repaired_columns}, "
|
||||
f"repaired_indexes={result.repaired_indexes}, "
|
||||
f"skipped_columns={result.skipped_columns}, "
|
||||
f"skipped_indexes={result.skipped_indexes}, "
|
||||
f"type_mismatches={result.type_mismatches}, "
|
||||
f"warnings={result.warnings}",
|
||||
LOG_COMMAND,
|
||||
)
|
||||
with contextlib.suppress(Exception):
|
||||
await _write_schema_report(result)
|
||||
return result
|
||||
|
||||
|
||||
async def repair_table_schema(table_name: str) -> SchemaGuardResult:
|
||||
"""Repair one table by reusing the startup guard path.
|
||||
|
||||
The current guard is cheap enough and already table-scoped internally by model
|
||||
metadata, so this helper keeps write-path recovery simple and conservative.
|
||||
"""
|
||||
result = SchemaGuardResult()
|
||||
result.drift = []
|
||||
connection = Tortoise.get_connection("default")
|
||||
dialect = _connection_dialect(connection)
|
||||
if dialect == "unknown":
|
||||
return result
|
||||
app = Tortoise.apps.get("models", {})
|
||||
for model in app.values():
|
||||
meta = getattr(model, "_meta", None)
|
||||
table = getattr(meta, "db_table", None) or getattr(meta, "table", None)
|
||||
if table == table_name:
|
||||
# Keep the implementation conservative: repairing all tables is still
|
||||
# startup-style work and avoids a second partial code path.
|
||||
return await repair_safe_schema_drift()
|
||||
return result
|
||||
@@ -0,0 +1,180 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Literal, Protocol
|
||||
|
||||
Dialect = Literal["sqlite", "postgres", "mysql", "unknown"]
|
||||
|
||||
|
||||
class SchemaOpRisk(str, Enum):
|
||||
SAFE = "safe"
|
||||
GUARDED = "guarded"
|
||||
MANUAL = "manual"
|
||||
|
||||
|
||||
class SchemaOp(Protocol):
|
||||
risk: SchemaOpRisk
|
||||
|
||||
def to_sql(self, dialect: Dialect) -> list[str]: ...
|
||||
|
||||
|
||||
def quote_identifier(identifier: str, dialect: Dialect) -> str:
|
||||
if dialect == "mysql":
|
||||
return f"`{identifier.replace('`', '``')}`"
|
||||
return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"'
|
||||
|
||||
|
||||
def _column_type(column_type: str | dict[str, str], dialect: Dialect) -> str:
|
||||
if isinstance(column_type, dict):
|
||||
return column_type.get(dialect) or column_type.get("default") or "TEXT"
|
||||
return column_type
|
||||
|
||||
|
||||
def _default_sql(default: str | float | bool | None, dialect: Dialect) -> str:
|
||||
if default is None:
|
||||
return ""
|
||||
if isinstance(default, bool):
|
||||
if dialect == "postgres":
|
||||
return " DEFAULT TRUE" if default else " DEFAULT FALSE"
|
||||
return " DEFAULT 1" if default else " DEFAULT 0"
|
||||
if isinstance(default, int | float):
|
||||
return f" DEFAULT {default}"
|
||||
escaped = default.replace("'", "''")
|
||||
return f" DEFAULT '{escaped}'"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AddColumn:
|
||||
table: str
|
||||
column: str
|
||||
column_type: str | dict[str, str]
|
||||
nullable: bool = True
|
||||
default: str | float | bool | None = None
|
||||
risk: SchemaOpRisk = SchemaOpRisk.SAFE
|
||||
|
||||
def to_sql(self, dialect: Dialect) -> list[str]:
|
||||
table = quote_identifier(self.table, dialect)
|
||||
column = quote_identifier(self.column, dialect)
|
||||
column_type = _column_type(self.column_type, dialect)
|
||||
null_sql = "" if self.nullable else " NOT NULL"
|
||||
default_sql = _default_sql(self.default, dialect)
|
||||
if not self.nullable and not default_sql:
|
||||
return []
|
||||
if dialect == "postgres":
|
||||
return [
|
||||
"ALTER TABLE "
|
||||
f"{table} ADD COLUMN IF NOT EXISTS {column} "
|
||||
f"{column_type}{default_sql}{null_sql}"
|
||||
]
|
||||
return [
|
||||
"ALTER TABLE "
|
||||
f"{table} ADD COLUMN {column} {column_type}{default_sql}{null_sql}"
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CreateIndex:
|
||||
table: str
|
||||
columns: tuple[str, ...]
|
||||
name: str | None = None
|
||||
if_not_exists: bool = True
|
||||
unique: bool = False
|
||||
where: str | None = None
|
||||
risk: SchemaOpRisk = SchemaOpRisk.SAFE
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
table: str,
|
||||
columns: tuple[str, ...] | list[str],
|
||||
name: str | None = None,
|
||||
if_not_exists: bool = True,
|
||||
unique: bool = False,
|
||||
where: str | None = None,
|
||||
risk: SchemaOpRisk | None = None,
|
||||
) -> None:
|
||||
object.__setattr__(self, "table", table)
|
||||
object.__setattr__(self, "columns", tuple(columns))
|
||||
object.__setattr__(self, "name", name)
|
||||
object.__setattr__(self, "if_not_exists", if_not_exists)
|
||||
object.__setattr__(self, "unique", unique)
|
||||
object.__setattr__(self, "where", where)
|
||||
resolved_risk = risk or (SchemaOpRisk.MANUAL if unique else SchemaOpRisk.SAFE)
|
||||
object.__setattr__(self, "risk", resolved_risk)
|
||||
|
||||
def to_sql(self, dialect: Dialect) -> list[str]:
|
||||
if self.unique or not self.columns:
|
||||
return []
|
||||
name = self.name or f"idx_{self.table}_{'_'.join(self.columns)}"[:62]
|
||||
if dialect == "mysql" and self.where:
|
||||
return []
|
||||
exists_sql = (
|
||||
"IF NOT EXISTS " if self.if_not_exists and dialect != "mysql" else ""
|
||||
)
|
||||
table = quote_identifier(self.table, dialect)
|
||||
index = quote_identifier(name, dialect)
|
||||
columns = ", ".join(
|
||||
quote_identifier(column, dialect) for column in self.columns
|
||||
)
|
||||
where = f" WHERE {self.where}" if self.where and dialect != "mysql" else ""
|
||||
return [f"CREATE INDEX {exists_sql}{index} ON {table}({columns}){where}"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RenameColumn:
|
||||
table: str
|
||||
old: str
|
||||
new: str
|
||||
risk: SchemaOpRisk = SchemaOpRisk.GUARDED
|
||||
|
||||
def to_sql(self, dialect: Dialect) -> list[str]:
|
||||
table = quote_identifier(self.table, dialect)
|
||||
old = quote_identifier(self.old, dialect)
|
||||
new = quote_identifier(self.new, dialect)
|
||||
return [f"ALTER TABLE {table} RENAME COLUMN {old} TO {new}"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DropColumn:
|
||||
table: str
|
||||
column: str
|
||||
risk: SchemaOpRisk = SchemaOpRisk.GUARDED
|
||||
|
||||
def to_sql(self, dialect: Dialect) -> list[str]:
|
||||
table = quote_identifier(self.table, dialect)
|
||||
column = quote_identifier(self.column, dialect)
|
||||
return [f"ALTER TABLE {table} DROP COLUMN {column}"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AlterColumnType:
|
||||
table: str
|
||||
column: str
|
||||
column_type: str | dict[str, str]
|
||||
nullable: bool | None = None
|
||||
risk: SchemaOpRisk = SchemaOpRisk.GUARDED
|
||||
|
||||
def to_sql(self, dialect: Dialect) -> list[str]:
|
||||
column_type = _column_type(self.column_type, dialect)
|
||||
table = quote_identifier(self.table, dialect)
|
||||
column = quote_identifier(self.column, dialect)
|
||||
if dialect == "postgres":
|
||||
return [f"ALTER TABLE {table} ALTER COLUMN {column} TYPE {column_type}"]
|
||||
if dialect == "mysql":
|
||||
null_sql = " NULL" if self.nullable else " NOT NULL"
|
||||
if self.nullable is None:
|
||||
null_sql = ""
|
||||
return [
|
||||
f"ALTER TABLE {table} MODIFY COLUMN {column} {column_type}{null_sql}"
|
||||
]
|
||||
return []
|
||||
|
||||
|
||||
def normalize_schema_ops(items: list[str | SchemaOp], dialect: Dialect) -> list[str]:
|
||||
sql_list: list[str] = []
|
||||
for item in items:
|
||||
if isinstance(item, str):
|
||||
sql_list.append(item)
|
||||
else:
|
||||
sql_list.extend(item.to_sql(dialect))
|
||||
return sql_list
|
||||
Reference in New Issue
Block a user