"""数据库管理工具模块
提供数据库引擎创建、会话管理和表初始化等功能,
基于SQLAlchemy的SQLite数据库实现。
支持从WebDAV服务器同步数据库文件。
"""
from __future__ import annotations
import os
import re
import shutil
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Any
from loguru import logger
from sqlalchemy import Boolean, create_engine, inspect, text
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, sessionmaker
from certflow.models.base import Base
# 尝试导入requests(用于WebDAV)
try:
import requests
from requests.auth import HTTPBasicAuth
REQUESTS_AVAILABLE = True
except ImportError:
REQUESTS_AVAILABLE = False
logger.warning("requests未安装,WebDAV功能不可用。请安装: pip install requests")
[文档]
class NutstoreWebDAVConfig:
"""坚果云WebDAV配置类"""
def __init__(
self,
url: str,
username: str,
password: str,
remote_path: str = "CertFlow/backups",
keep_backups: int = 10,
):
"""
Args:
url: WebDAV服务器地址(如: https://dav.jianguoyun.com/dav/)
username: 用户名(通常是邮箱)
password: 应用密码
remote_path: 远程备份目录路径
keep_backups: 保留的备份数量
"""
self.url = url.rstrip("/")
self.username = username
self.password = password
self.remote_path = remote_path
self.keep_backups = keep_backups
[文档]
def get_auth(self) -> HTTPBasicAuth:
"""获取HTTP Basic认证
Returns:
HTTPBasicAuth: 用于 WebDAV 请求的 Basic Auth 认证对象。
"""
return HTTPBasicAuth(self.username, self.password)
[文档]
def get_remote_url(self, filename: str = "") -> str:
"""获取完整的远程URL
Args:
filename: 文件名,为空时返回目录 URL。
Returns:
str: 完整的远程 URL 路径。
"""
base_url = f"{self.url}/{self.remote_path}"
if filename:
return f"{base_url}/{filename}"
return base_url
[文档]
class DatabaseManager:
"""数据库管理器
管理数据库引擎、会话工厂和表创建,提供统一的数据库操作入口。
支持数据库初始化、表迁移(自动补齐缺失列)和会话管理。
支持从WebDAV服务器同步数据库文件。
Attributes:
db_path: 数据库文件路径
engine: SQLAlchemy引擎实例,负责数据库连接
SessionLocal: 会话工厂实例,用于创建数据库会话
webdav_config: WebDAV配置(如果启用)
Examples:
>>> from certflow.utils.database import DatabaseManager
>>>
>>> # 创建数据库管理器
>>> db = DatabaseManager("path/to/database.db")
>>>
>>> # 初始化数据库
>>> db.init_db(create_tables=True)
>>>
>>> # 从WebDAV拉取数据库
>>> webdav_config = NutstoreWebDAVConfig(
... url="https://dav.jianguoyun.com/dav/",
... username="user@example.com",
... password="app_password"
... )
>>> db.init_db(sync_from_webdav=True)
>>>
>>> # 获取会话并执行操作
>>> session = db.get_session()
>>> try:
... result = session.query(SalePlan).all()
... finally:
... session.close()
"""
def __init__(
self,
db_path: str | None = None,
webdav_config: NutstoreWebDAVConfig | None = None,
auto_backup_before_sync: bool = True,
) -> None:
"""初始化数据库管理器
Args:
db_path: 数据库文件路径。如果为None,则使用默认路径:
项目根目录/database/certflow.db
webdav_config: WebDAV配置,如果提供则启用WebDAV同步功能
auto_backup_before_sync: 同步前是否自动备份本地数据库
"""
if db_path is None:
# 调试库切换点:设置 CERTFLOW_DB_PATH 即可让 GUI / 全部 CLI 改用调试库,
# 不污染主库 database/certflow.db。显式 db_path 参数仍优先于环境变量。
db_path = os.getenv("CERTFLOW_DB_PATH")
if db_path is None:
from certflow.utils.path_utils import get_project_root
project_root = get_project_root()
db_path = str(project_root / "database" / "certflow.db")
self.db_path: Path = Path(db_path)
self.backup_dir: Path = self.db_path.parent / "backups"
self.engine: Engine | None = None
self.SessionLocal: sessionmaker | None = None
self.webdav_config: NutstoreWebDAVConfig | None = webdav_config
self.auto_backup_before_sync = auto_backup_before_sync
[文档]
def init_db(
self, create_tables: bool = True, sync_from_webdav: bool = False, force_sync: bool = False
) -> None:
"""初始化数据库
创建数据库目录、SQLAlchemy引擎和会话工厂,
可选创建所有数据表。已有表会自动补齐缺失列。
可选从WebDAV同步数据库文件。
Args:
create_tables: 是否创建数据表,默认为True
sync_from_webdav: 是否从WebDAV同步数据库,默认为False
force_sync: 是否强制同步(即使本地数据库存在也尝试从远程拉取)
Raises:
Exception: 当数据库初始化失败时抛出异常
Examples:
>>> db = DatabaseManager("test.db")
>>> db.init_db(create_tables=True)
>>> print("数据库初始化完成")
"""
try:
# 确保数据库目录存在
self.db_path.parent.mkdir(parents=True, exist_ok=True)
self.backup_dir.mkdir(parents=True, exist_ok=True)
# 如果启用WebDAV同步,在创建引擎前先拉取数据库
if sync_from_webdav and self.webdav_config:
self._sync_from_webdav_on_init(force_sync)
# 如果数据库不存在,创建空数据库
if not self.db_path.exists():
self._create_empty_database()
logger.info(f"创建空白数据库: {self.db_path}")
# 创建引擎
self.engine = create_engine(
f"sqlite:///{self.db_path}", connect_args={"check_same_thread": False}, echo=False
)
# 消除「database is locked」写保护锁:开启 WAL 模式 + 写忙超时。
# 普通 SQLite 删除/回滚日志模式下,UI 主线程与队列工作线程同时写
# 同一个 .db 文件时会直接抛锁;WAL 允许读写并发、busy_timeout 让
# 写方在锁被占用时阻塞等待而非立即失败。详见 certflow.utils.retry。
self._apply_sqlite_pragmas()
# 创建会话工厂
self.SessionLocal = sessionmaker(bind=self.engine)
if create_tables:
# 先创建新表(不报错如果表已存在)
Base.metadata.create_all(self.engine)
logger.info("数据库表检查/创建完成")
# 自动补齐已有表中缺失的列
self._migrate_missing_columns()
# 修复历史布尔列默认值(字符串 'False' → 整数 0)
self._repair_bool_defaults()
# 执行列重命名迁移(如 certificates.status → print_status)
self._rename_columns()
# 自动收缩列类型(VARCHAR 长度缩小的场景)
self._migrate_column_types()
# 删除旧索引(如 idx_cert_status → idx_cert_print_status)
self._drop_old_indexes()
# 删除数据库重构阶段1b 识别的冗余索引(零数据风险)
self._drop_redundant_indexes()
# 创建数据库索引以优化查询性能
self._create_indexes()
# 初始化种子数据(VBAMapping 等)
self._seed_data()
# S1 安全缺口修复(M4 残留):所有写库操作完成后再限定权限为 0o600,
# 仅属主可读写,避免同机其他用户读取含业务数据的 .db。
# 必须在 WAL/建表/迁移之后执行,否则 SQLite 重写文件会重置 mode。
self._harden_db_file_permissions()
logger.info(f"数据库初始化成功: {self.db_path}")
except Exception as e:
logger.error(f"数据库初始化失败: {e}")
raise
[文档]
def ensure_initialized_if_needed(self) -> None:
"""仅在引擎尚未创建时初始化数据库(幂等、低开销)
供装配层(``BaseController`` / ``AppContext``)在「不确定当前是否已完成
引导」的构造路径上调用:已初始化则直接跳过,未初始化则补建表/迁移。
区别于 ``init_db`` 的「总是执行完整初始化」,本方法避免在每次
Controller 构造时重复跑建表与迁移流程。
"""
if self.engine is None:
self.init_db(create_tables=True)
[文档]
def warn_if_base_data_empty(self) -> bool:
"""检查基础字典底座是否为空,为空则给出 seed 提示(只读、不自动播种)。
对应 ``DATABASE_BASE_DATA.md §4.2`` 的「可选加固」:在 ``init_db`` 之后挂载
一个轻量检查——若 ``material_grades`` / ``model_param_mappings`` /
``caliber_mappings`` 任一为空,记录 WARNING 并提示运行
``scripts/data/seed_base_data.py``。
刻意**不**自动播种(§4.2 明确「不要每次启动自动 seed」,避免依赖 xlsx /
拖慢启动);自动播种由 ``main._maybe_seed_base_data`` 在 GUI 启动路径单独负责。
CLI / 双库等不走自动播种的入口调用本方法即可获得「空库提示」安全网。
Returns:
bool: 字典表是否为空(True=空,需 seed)。
"""
if self.SessionLocal is None:
return False
from certflow.models.caliber_mapping import CaliberMapping
from certflow.models.material_grade import MaterialGrade
from certflow.models.model_param_mapping import ModelParamMapping
tables = (
("material_grades", MaterialGrade),
("model_param_mappings", ModelParamMapping),
("caliber_mappings", CaliberMapping),
)
session = self.get_session()
try:
empty_tables = [name for name, model in tables if session.query(model).count() == 0]
finally:
session.close()
if not empty_tables:
return False
logger.warning(
f"基础字典底座为空({', '.join(empty_tables)})。请运行 "
"`python scripts/data/seed_base_data.py` 一键幂等播种材质/型号压力/口径字典;"
"该检查不自动播种,以免影响启动速度。"
)
return True
def _sync_from_webdav_on_init(self, force_sync: bool = False) -> None:
"""在初始化时从WebDAV同步数据库
策略:
1. 优先使用远程最新备份(如果存在且可用)
2. 远程连接失败或没有备份时,使用本地数据库(或创建空白数据库)
Args:
force_sync: 是否强制同步(忽略本地数据库存在性检查)
"""
# 如果本地数据库存在且不强制同步,询问或跳过
if self.db_path.exists() and not force_sync:
logger.info(f"本地数据库已存在: {self.db_path}")
# 可以选择检查远程是否有更新,这里先使用本地
return
# 尝试从远程获取最新备份
logger.info("尝试从坚果云同步数据库...")
try:
latest_backup = self._get_latest_remote_backup()
if latest_backup:
# 找到远程备份,下载并替换本地
logger.info(f"发现远程备份: {latest_backup}")
success = self._download_remote_backup(latest_backup)
if success:
logger.info("✅ 已从坚果云同步数据库")
return
logger.warning("下载远程备份失败")
else:
logger.info("远程没有备份文件")
except Exception as e:
logger.warning(f"连接坚果云失败: {e}")
# 远程不可用或没有备份,使用本地或创建空白数据库
if not self.db_path.exists():
logger.info("远程不可用且本地无数据库,将创建空白数据库")
else:
logger.info("远程不可用,使用本地数据库")
def _apply_sqlite_pragmas(self) -> None:
"""为 SQLite 连接设置 WAL + busy_timeout PRAGMA。
- ``journal_mode=WAL`` 是**库级**属性,用独立临时连接执行一次即可
对后续所有连接生效,让读写并发(写不阻塞读);
- ``busy_timeout`` 是**每个连接**的属性,新连接默认不继承,因此
通过 SQLAlchemy 的 ``connect`` 事件钩子在*每次新建连接*时自动执行,
保证本引擎产生的所有业务会话都带 15s 写忙超时。
- ``synchronous=NORMAL``:WAL 下兼顾安全与性能。
这样 UI 主线程(保存订正、写打印日志)与队列工作线程(写
print_jobs)并发写同一 .db 时不再直接崩 ``database is locked``。
"""
try:
from sqlalchemy import create_engine as _create_engine
from sqlalchemy import event
# 1) 库级 WAL:对库文件执行一次即可持久生效
tmp_engine = _create_engine(
f"sqlite:///{self.db_path}",
connect_args={"check_same_thread": False},
)
with tmp_engine.begin() as conn:
conn.exec_driver_sql("PRAGMA journal_mode=WAL")
conn.exec_driver_sql("PRAGMA synchronous=NORMAL")
tmp_engine.dispose()
# 2) 每连接 busy_timeout:注册在业务引擎上,新建连接自动带上
def _set_busy_timeout(dbapi_conn, _record): # noqa: ANN001
cur = dbapi_conn.cursor()
cur.execute("PRAGMA busy_timeout=15000")
cur.close()
event.listen(self.engine, "connect", _set_busy_timeout)
logger.info(f"SQLite PRAGMA 已应用 (WAL + busy_timeout=15000): {self.db_path}")
except Exception as e: # noqa: BLE001 - PRAGMA 失败不应阻断启动
logger.warning(f"应用 SQLite PRAGMA 失败(将退化为默认日志模式): {e}")
[文档]
@staticmethod
def commit_with_retry(session: Any, max_attempts: int = 12) -> None:
"""带退避重试的 session.commit,专治瞬时写锁。
直接在业务 session 上调用,遇 ``database is locked`` 时指数退避
重试(base 0.05s / 上限 0.8s / 累计约 8s)。配合 WAL+busy_timeout
使用,几乎可消除打印流程中的写保护锁现象。
Args:
session: SQLAlchemy Session(其引擎需已设 busy_timeout)。
max_attempts: 最大尝试次数。
Raises:
非锁类异常或重试耗尽后仍失败时原样抛出。
"""
from certflow.utils.retry import retry_on_locked
retry_on_locked(session.commit)
[文档]
@staticmethod
def flush_with_retry(session: Any, max_attempts: int = 12) -> None:
"""带退避重试的 session.flush(写库前先落盘也可能需要抢锁)。"""
from certflow.utils.retry import retry_on_locked
retry_on_locked(session.flush)
def _create_empty_database(self) -> None:
"""创建空白的SQLite数据库文件"""
# 创建一个空的SQLite数据库
empty_engine = create_engine(f"sqlite:///{self.db_path}")
empty_engine.dispose()
logger.info(f"已创建空白数据库: {self.db_path}")
def _harden_db_file_permissions(self) -> None:
"""创建/同步后收紧数据库文件权限为 0o600(仅属主可读写)。
对应 BLUEPRINT §2.3 独立轨 S1(M4 残留安全缺口)。需在 WAL/建表/
迁移等写库操作全部完成之后调用,否则 SQLite 重写主库文件会重置 mode。
降级策略:权限位不受支持(如部分挂载文件系统)时仅告警,不阻断初始化。
"""
try:
os.chmod(self.db_path, 0o600)
except OSError as e: # noqa: BLE001
logger.warning(f"无法设置数据库文件权限 0o600({self.db_path}): {e}")
def _get_latest_remote_backup(self) -> str | None:
"""获取远程最新的备份文件名
Returns:
最新的备份文件名,如果没有则返回None
"""
if not self.webdav_config:
return None
if not REQUESTS_AVAILABLE:
raise ImportError("请先安装requests: pip install requests")
url = self.webdav_config.get_remote_url()
auth = self.webdav_config.get_auth()
# 使用PROPFIND方法列出目录内容
resp = requests.request("PROPFIND", url, auth=auth, timeout=30)
if resp.status_code != 207:
logger.warning(f"无法访问远程目录: HTTP {resp.status_code}")
return None
# 查找所有备份文件(格式: certflow_YYYYMMDD_HHMMSS.db)
files = re.findall(r"certflow_\d{8}_\d{6}\.db", resp.text)
if not files:
return None
# 返回最新的文件(按文件名排序,因为时间戳在前)
return sorted(files)[-1]
def _download_remote_backup(self, filename: str) -> bool:
"""下载远程备份文件到本地
Args:
filename: 远程备份文件名
Returns:
是否下载成功
"""
if not self.webdav_config:
return False
# 备份当前本地数据库(如果存在)
if self.auto_backup_before_sync and self.db_path.exists():
self._backup_local_database("before_sync")
# 下载远程文件
remote_url = self.webdav_config.get_remote_url(filename)
auth = self.webdav_config.get_auth()
logger.info(f"下载远程备份: {filename}")
resp = requests.get(remote_url, auth=auth, timeout=60)
if resp.status_code == 200:
# 临时文件,然后替换
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp_file:
tmp_path = Path(tmp_file.name)
tmp_path.write_bytes(resp.content)
# 替换本地数据库
if self.engine:
self.engine.dispose()
shutil.move(str(tmp_path), str(self.db_path))
logger.info(f"数据库已更新: {self.db_path}")
return True
logger.error(f"下载失败: HTTP {resp.status_code}")
return False
def _backup_local_database(self, suffix: str = "") -> str:
"""备份本地数据库
Args:
suffix: 备份文件后缀标识
Returns:
备份文件路径
"""
if not self.db_path.exists():
return ""
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
backup_name = f"certflow_{timestamp}_{suffix}.db" if suffix else f"certflow_{timestamp}.db"
backup_path = self.backup_dir / backup_name
shutil.copy2(self.db_path, backup_path)
logger.info(f"本地数据库已备份: {backup_path}")
# 清理旧的本地备份(保留最近20个)
self._cleanup_local_backups(keep=20)
return str(backup_path)
def _cleanup_local_backups(self, keep: int = 20) -> None:
"""清理本地旧备份
Args:
keep: 保留的备份数量
"""
backups = sorted(self.backup_dir.glob("certflow_*.db"))
if len(backups) <= keep:
return
for old_backup in backups[:-keep]:
old_backup.unlink()
logger.debug(f"删除旧备份: {old_backup.name}")
[文档]
def pull_from_webdav(self, backup_before_pull: bool = True) -> bool:
"""从坚果云拉取最新数据库备份
Args:
backup_before_pull: 拉取前是否备份本地数据库
Returns:
是否成功拉取并替换
"""
if not self.webdav_config:
logger.warning("未配置WebDAV,无法拉取")
return False
logger.info("从坚果云拉取最新数据库...")
try:
# 获取最新备份
latest = self._get_latest_remote_backup()
if not latest:
logger.warning("远程没有备份文件")
return False
# 备份本地数据库
if backup_before_pull and self.db_path.exists():
self._backup_local_database("before_pull")
# 下载并替换
success = self._download_remote_backup(latest)
if success:
# 重新初始化引擎
if self.engine:
self.engine.dispose()
self.engine = create_engine(
f"sqlite:///{self.db_path}",
connect_args={"check_same_thread": False},
echo=False,
)
# pull 同步后用同一份 PRAGMA 配置,保证从远程拉回的库也走 WAL
self._apply_sqlite_pragmas()
self.SessionLocal = sessionmaker(bind=self.engine)
# 同步拉回的文件同样收紧权限(S1)
self._harden_db_file_permissions()
logger.info("✅ 数据库已从坚果云同步")
return True
return False
except Exception as e:
logger.error(f"从坚果云拉取失败: {e}")
return False
[文档]
def push_to_webdav(self, backup_remote: bool = True) -> bool:
"""上传本地数据库到坚果云
Args:
backup_remote: 是否备份远程数据库(保留旧版本)
Returns:
是否成功上传
"""
if not self.webdav_config:
logger.warning("未配置WebDAV,无法上传")
return False
if not self.db_path.exists():
logger.error(f"本地数据库不存在: {self.db_path}")
return False
if not REQUESTS_AVAILABLE:
raise ImportError("请先安装requests: pip install requests")
# 关闭数据库连接确保文件完整性
if self.engine:
self.engine.dispose()
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"certflow_{timestamp}.db"
remote_url = self.webdav_config.get_remote_url(filename)
auth = self.webdav_config.get_auth()
logger.info(f"上传数据库到坚果云: {filename}")
try:
with open(self.db_path, "rb") as f:
resp = requests.put(remote_url, data=f, auth=auth, timeout=60)
if resp.status_code in (200, 201, 204):
logger.info("✅ 数据库上传成功")
# 清理远程旧备份
self._cleanup_remote_backups()
return True
logger.error(f"上传失败: HTTP {resp.status_code}")
return False
except Exception as e:
logger.error(f"上传到坚果云失败: {e}")
return False
def _cleanup_remote_backups(self) -> None:
"""清理远程旧备份,只保留最新的keep_backups个"""
if not self.webdav_config:
return
url = self.webdav_config.get_remote_url()
auth = self.webdav_config.get_auth()
resp = requests.request("PROPFIND", url, auth=auth, timeout=30)
if resp.status_code != 207:
return
files = sorted(re.findall(r"certflow_\d{8}_\d{6}\.db", resp.text))
keep = self.webdav_config.keep_backups
if len(files) <= keep:
return
for old_file in files[:-keep]:
del_url = self.webdav_config.get_remote_url(old_file)
r = requests.delete(del_url, auth=auth, timeout=10)
if r.status_code in (200, 204):
logger.info(f"删除远程旧备份: {old_file}")
[文档]
def sync_with_webdav(self, direction: str = "pull") -> bool:
"""与WebDAV服务器同步数据库
Args:
direction: 同步方向,"pull"(从服务器拉取)或 "push"(推送到服务器)
Returns:
是否同步成功
"""
if direction == "pull":
return self.pull_from_webdav()
if direction == "push":
return self.push_to_webdav()
raise ValueError(f"不支持的同步方向: {direction}")
[文档]
def check_remote_status(self) -> dict[str, Any]:
"""检查远程备份状态
Returns:
包含远程备份信息的字典
"""
if not self.webdav_config:
return {"available": False, "error": "未配置WebDAV"}
try:
url = self.webdav_config.get_remote_url()
auth = self.webdav_config.get_auth()
resp = requests.request("PROPFIND", url, auth=auth, timeout=30)
if resp.status_code != 207:
return {"available": False, "error": f"HTTP {resp.status_code}"}
files = sorted(re.findall(r"certflow_\d{8}_\d{6}\.db", resp.text))
return {
"available": True,
"backup_count": len(files),
"latest_backup": files[-1] if files else None,
"backups": files,
}
except Exception as e:
return {"available": False, "error": str(e)}
def _migrate_missing_columns(self) -> None:
"""自动补齐已有表中模型定义但数据库缺失的列
遍历所有映射表,对比模型定义与实际数据库列,
对缺失列执行 ALTER TABLE ADD COLUMN。
用于实现数据库架构的平滑升级。
Note:
该方法仅添加列,不删除或修改现有列。
不支持设置外键约束和唯一约束。
"""
insp = inspect(self.engine)
for table_name, table_obj in Base.metadata.tables.items():
if table_name not in insp.get_table_names():
continue # 表不存在则跳过
existing_cols = {col["name"] for col in insp.get_columns(table_name)}
model_cols = {c.name for c in table_obj.columns}
missing_cols = model_cols - existing_cols
if not missing_cols:
continue
with self.engine.begin() as conn:
for col_name in missing_cols:
col = table_obj.columns[col_name]
col_type = col.type.compile(self.engine.dialect)
nullable = "" if col.nullable else " NOT NULL"
default_val = ""
if col.default and col.default.arg is not None:
arg = col.default.arg
if isinstance(arg, bool):
default_val = f" DEFAULT {1 if arg else 0}"
elif isinstance(arg, (int, float)):
default_val = f" DEFAULT {arg}"
else:
default_val = f" DEFAULT '{arg}'"
sql = f'ALTER TABLE "{table_name}" ADD COLUMN "{col_name}" {col_type}{nullable}{default_val}'
logger.info(f"迁移列: {table_name}.{col_name} (type={col_type})")
conn.execute(text(sql))
if missing_cols:
logger.info(f"表 {table_name} 补齐了 {len(missing_cols)} 个缺失列: {missing_cols}")
def _repair_bool_defaults(self) -> None:
"""修复历史迁移中布尔列默认值被存为字符串 'False'/'True' 的问题
原因:旧版 _migrate_missing_columns 对所有类型使用
DEFAULT '{col.default.arg}',导致 SQLite 将 Python False
存为字符串 'False',而 'False' 在 Python 中是 truthy 的。
修复:
1. 将布尔列中的字符串 'False'→0, 'True'→1
2. 重算 sale_plans 的 progress_percent(因脏布尔值可能已误算为 100)
"""
insp = inspect(self.engine)
bool_cols_info: list[tuple[str, str]] = []
for table_name, table_obj in Base.metadata.tables.items():
if table_name not in insp.get_table_names():
continue
for col in table_obj.columns:
if isinstance(col.type, Boolean):
bool_cols_info.append((table_name, col.name))
if not bool_cols_info:
return
repaired = 0
with self.engine.begin() as conn:
for table_name, col_name in bool_cols_info:
for bad_val, correct_val in (("False", 0), ("True", 1)):
result = conn.execute(
text(
f'UPDATE "{table_name}" SET "{col_name}" = :correct '
f'WHERE "{col_name}" = :bad'
),
{"correct": correct_val, "bad": bad_val},
)
repaired += result.rowcount
# 修复后重算 sale_plans 的进度百分比
# (certificate_done + nameplate_done + test_report_done + warranty_done + cert_scan_done) / 5 * 100
result = conn.execute(
text(
"UPDATE sale_plans SET progress_percent = CAST("
" (certificate_done + nameplate_done + test_report_done "
" + warranty_done + cert_scan_done) * 100 / 5 AS INTEGER"
") WHERE certificate_done IS NOT NULL"
)
)
progress_fixed = result.rowcount
# 诊断:检查是否存在全部 _done=1 但对应 _time 为空的行
# (说明这些行是历史迁移 bug 导致的虚假完成,而非真实操作)
suspect = conn.execute(
text(
"SELECT COUNT(*) FROM sale_plans WHERE "
"certificate_done = 1 AND certificate_time IS NULL AND "
"nameplate_done = 1 AND nameplate_time IS NULL AND "
"test_report_done = 1 AND test_report_time IS NULL"
)
).scalar()
if repaired:
logger.info(
f"布尔列默认值修复完成: 修复了 {repaired} 行"
f"{f', 重算了 {progress_fixed} 行进度' if progress_fixed else ''}"
)
if suspect:
logger.warning(
f"检测到 {suspect} 行进度字段全部为 True 但无对应完成时间,"
f"可能是历史迁移 bug 导致的虚假完成,请手动核实"
)
def _detect_shrink_cols(
self, model_cols: dict[str, Any], existing_cols: dict[str, Any]
) -> dict[str, str]:
"""检测需要收缩 VARCHAR 长度的列
Args:
model_cols: 模型定义的列 {name: Column}
existing_cols: 数据库中已有的列 {name: inspector_dict}
Returns:
需要收缩的列 {col_name: new_varchar_type}
"""
shrink_cols: dict[str, str] = {}
for col_name, model_col in model_cols.items():
existing = existing_cols.get(col_name)
if not existing:
continue
model_type = str(model_col.type).upper()
existing_type = str(existing["type"]).upper()
if not (
model_type.startswith("VARCHAR(")
and existing_type.startswith("VARCHAR(")
and model_type != existing_type
):
continue
try:
new_len = int(model_type.split("(")[1].rstrip(")"))
old_len = int(existing_type.split("(")[1].rstrip(")"))
if new_len < old_len:
shrink_cols[col_name] = model_type
except (ValueError, IndexError):
continue
return shrink_cols
def _build_column_defs(self, table_obj: Any) -> list[str]:
"""构建新表的列定义列表"""
model_cols = {c.name: c for c in table_obj.columns}
col_defs = []
for col_name in (c.name for c in table_obj.columns):
col = model_cols.get(col_name)
if col is None:
continue
col_type = col.type.compile(self.engine.dialect)
nullable = "" if col.nullable else " NOT NULL"
default_val = ""
if col.default and col.default.arg is not None:
arg = col.default.arg
if isinstance(arg, bool):
default_val = f" DEFAULT {1 if arg else 0}"
elif isinstance(arg, (int, float)):
default_val = f" DEFAULT {arg}"
else:
default_val = f" DEFAULT '{arg}'"
col_defs.append(f'"{col_name}" {col_type}{nullable}{default_val}')
return col_defs
def _drop_old_indexes(self) -> None:
"""删除旧的索引(字段重命名后)"""
insp = inspect(self.engine)
indexes = insp.get_indexes("certificates")
for idx in indexes:
if idx["name"] == "idx_cert_status":
with self.engine.begin() as conn:
conn.execute(text('DROP INDEX IF EXISTS "idx_cert_status"'))
logger.info("删除旧索引: idx_cert_status")
break
def _drop_redundant_indexes(self) -> None:
"""删除数据库重构阶段1b 识别的冗余索引
这些索引或被复合索引前缀覆盖、或与 unique 约束自动生成的索引重复,
删除后不影响查询(复合/唯一索引仍可被复用),且零数据风险。
使用 DROP INDEX IF EXISTS 保证幂等,空白库与已有库均安全。
"""
redundant_indexes = [
# sale_plans:被复合索引前缀覆盖 / 重复 unique_key 约束
'DROP INDEX IF EXISTS "idx_sale_plan_product_model"',
'DROP INDEX IF EXISTS "idx_sale_plan_sales_order_no"',
'DROP INDEX IF EXISTS "idx_sale_plan_customer"',
'DROP INDEX IF EXISTS "idx_sale_plan_unique_key"',
# certificates:重复 certificate_no unique 约束
'DROP INDEX IF EXISTS "idx_cert_certificate_no"',
# auto_number:被复合索引前缀覆盖(CounterRule / Counter 各一)
'DROP INDEX IF EXISTS "idx_auto_number_rule_code"',
'DROP INDEX IF EXISTS "idx_counter_rule_code"',
# material_grade:被 unique(grade, standard_code, element) 前缀覆盖
'DROP INDEX IF EXISTS "idx_material_grade_standard"',
# caliber_mapping:重复 raw_text unique 约束
'DROP INDEX IF EXISTS "idx_caliber_raw_text"',
]
with self.engine.begin() as conn:
for sql in redundant_indexes:
conn.execute(text(sql))
logger.info(f"阶段1b:已清理 {len(redundant_indexes)} 个冗余索引")
def _rename_columns(self) -> None:
"""执行列重命名迁移
处理需要重命名的列,如 certificates.status → print_status
"""
insp = inspect(self.engine)
# 定义重命名映射: {表名: [(旧列名, 新列名, 类型), ...]}
rename_map = {
"certificates": [
("status", "print_status", "VARCHAR(20)"), # 原 status → print_status
]
}
for table_name, columns in rename_map.items():
if table_name not in insp.get_table_names():
continue
existing_cols = {col["name"] for col in insp.get_columns(table_name)}
for old_name, new_name, col_type in columns:
if old_name not in existing_cols:
logger.debug(f"列 {table_name}.{old_name} 不存在,跳过重命名")
continue
if new_name in existing_cols:
logger.debug(f"列 {table_name}.{new_name} 已存在,跳过重命名")
continue
# SQLite 不支持直接重命名列,需要重建表
self._rename_column_sqlite(table_name, old_name, new_name, col_type)
logger.info(f"迁移: {table_name}.{old_name} → {new_name}")
def _rename_column_sqlite(
self, table_name: str, old_name: str, new_name: str, col_type: str
) -> None:
"""SQLite 列重命名(通过重建表实现)"""
insp = inspect(self.engine)
# 获取表的所有列
columns = insp.get_columns(table_name)
# 获取主键列
pk = insp.get_pk_constraint(table_name)
pk_columns = set(pk.get("constrained_columns", []))
# 构建列定义(明确标记主键)
col_defs = []
for col in columns:
col_name = col["name"]
if col_name == old_name:
col_name = new_name
col_type_use = col_type
else:
col_type_use = col["type"]
# 添加主键标记(如果是主键列)
pk_suffix = " PRIMARY KEY" if col_name in pk_columns else ""
col_defs.append(f'"{col_name}" {col_type_use}{pk_suffix}')
# 获取外键约束
fks = insp.get_foreign_keys(table_name)
with self.engine.begin() as conn:
temp_table = f"{table_name}_tmp"
# 创建临时表
conn.execute(text(f'DROP TABLE IF EXISTS "{temp_table}"'))
conn.execute(text(f'CREATE TABLE "{temp_table}" ({", ".join(col_defs)})'))
# 复制数据
old_cols = [f'"{c["name"]}"' for c in columns]
new_cols = [f'"{new_name if c["name"] == old_name else c["name"]}"' for c in columns]
conn.execute(
text(
f'INSERT INTO "{temp_table}" ({", ".join(new_cols)}) '
f'SELECT {", ".join(old_cols)} FROM "{table_name}"'
)
)
# 删除旧表
conn.execute(text(f'DROP TABLE "{table_name}"'))
# 重命名新表
conn.execute(text(f'ALTER TABLE "{temp_table}" RENAME TO "{table_name}"'))
# 重建外键
for fk in fks:
if fk.get("referred_table"):
conn.execute(
text(
f'ALTER TABLE "{table_name}" ADD FOREIGN KEY({fk["constrained_columns"][0]}) '
f"REFERENCES {fk['referred_table']}({fk['referred_columns'][0]})"
)
)
# 重建索引(调用现有方法)
self._create_indexes()
def _migrate_column_types(self) -> None:
"""自动修正已有表中列类型与模型不一致的列
SQLite 不支持 ALTER COLUMN TYPE,因此采用重建表方式:
1. 创建临时新表(使用模型定义的类型)
2. 将旧表数据迁移到新表
3. 删除旧表并重命名新表
4. 重建索引
仅对 VARCHAR -> VARCHAR 且长度缩小的场景执行迁移。
"""
insp = inspect(self.engine)
for table_name, table_obj in Base.metadata.tables.items():
if table_name not in insp.get_table_names():
continue
existing_cols = {col["name"]: col for col in insp.get_columns(table_name)}
model_cols = {c.name: c for c in table_obj.columns}
shrink_cols = self._detect_shrink_cols(model_cols, existing_cols)
if not shrink_cols:
continue
logger.info(
f"表 {table_name} 检测到 {len(shrink_cols)} 个列类型需要收缩: {shrink_cols}"
)
all_col_names = [c.name for c in table_obj.columns]
col_defs = self._build_column_defs(table_obj)
with self.engine.begin() as conn:
temp_table = f"{table_name}_tmp"
conn.execute(text(f'DROP TABLE IF EXISTS "{temp_table}"'))
conn.execute(text(f'CREATE TABLE "{temp_table}" ({", ".join(col_defs)})'))
cols_str = ", ".join(f'"{c}"' for c in all_col_names)
conn.execute(
text(
f'INSERT INTO "{temp_table}" ({cols_str}) '
f'SELECT {cols_str} FROM "{table_name}"'
)
)
conn.execute(text(f'DROP TABLE "{table_name}"'))
conn.execute(text(f'ALTER TABLE "{temp_table}" RENAME TO "{table_name}"'))
logger.info(f"表 {table_name} 列类型收缩完成,重建索引...")
Base.metadata.create_all(self.engine)
def _create_indexes(self) -> None:
"""创建数据库索引(已迁移至模型 __table_args__ 中定义)
索引定义位于各模型类的 __table_args__ 中,由 Base.metadata.create_all()
自动创建。此方法保留空实现以保持向后兼容。
Note:
当前所有索引定义已在各模型的 __table_args__ 中配置,
因此此方法不需要额外实现。
"""
# #30 P0 A 方案:为新增的发货状态溯源列补建索引。
# 模型 __table_args__ 中的索引仅对全新库生效;已有库经 _migrate_missing_columns
# 补列后不会自动建索引,故在此用 IF NOT EXISTS 显式补齐,确保新旧库一致。
# #88 数据库重构(阶段1a/2):为新增索引在已有库中补齐。
# 模型 __table_args__ 中的索引仅对全新库生效;已有库经迁移后不会自动建索引,
# 故在此用 IF NOT EXISTS 显式补齐,确保新旧库一致。
missing_indexes = [
# sale_plans
'CREATE INDEX IF NOT EXISTS "idx_sale_plan_shipping_status" '
'ON "sale_plans" ("shipping_status")',
'CREATE INDEX IF NOT EXISTS "idx_sale_plan_shipping_status_source" '
'ON "sale_plans" ("shipping_status_source")',
'CREATE INDEX IF NOT EXISTS "idx_sale_plan_sort_group" ON "sale_plans" ("sort_group")',
'CREATE INDEX IF NOT EXISTS "idx_sale_plan_product_code" '
'ON "sale_plans" ("product_code")',
# certificates
'CREATE INDEX IF NOT EXISTS "idx_cert_created_at" ON "certificates" ("created_at")',
'CREATE INDEX IF NOT EXISTS "idx_cert_product_model" '
'ON "certificates" ("product_model")',
# print_logs(零索引表补齐)
'CREATE INDEX IF NOT EXISTS "idx_print_log_status" ON "print_logs" ("status")',
'CREATE INDEX IF NOT EXISTS "idx_print_log_certificate_id" '
'ON "print_logs" ("certificate_id")',
# unmatched_certificates(零索引表补齐)
'CREATE INDEX IF NOT EXISTS "idx_unmatched_status" '
'ON "unmatched_certificates" ("status")',
'CREATE INDEX IF NOT EXISTS "idx_unmatched_certificate_no" '
'ON "unmatched_certificates" ("certificate_no")',
]
with self.engine.begin() as conn:
for sql in missing_indexes:
conn.execute(text(sql))
def _seed_data(self) -> None:
"""初始化种子数据"""
from certflow.models.vba_mapping import VBAMapping
session = self.get_session()
try:
VBAMapping.seed(session)
finally:
session.close()
[文档]
def get_session(self) -> Session:
"""获取数据库会话
如果会话工厂未初始化则先初始化数据库,然后创建并返回新会话。
Returns:
Session: SQLAlchemy数据库会话实例,可用于执行数据库操作
Examples:
>>> db = DatabaseManager()
>>> session = db.get_session()
>>> try:
... # 执行数据库操作
... result = session.query(SalePlan).filter_by(id=1).first()
... finally:
... session.close()
"""
if self.SessionLocal is None:
self.init_db()
return self.SessionLocal()
[文档]
def close(self) -> None:
"""关闭数据库连接
释放数据库引擎资源,关闭所有活动连接。
Examples:
>>> db = DatabaseManager()
>>> db.init_db()
>>> # 使用数据库...
>>> db.close() # 程序退出前释放资源
"""
if self.engine:
self.engine.dispose()
logger.info("数据库连接已关闭")