266 lines
8.8 KiB
Python
266 lines
8.8 KiB
Python
|
"""数据库自动初始化模块。
|
|||
|
|
|
|||
|
|
在应用启动时自动执行：
|
|||
|
|
1. 检查并创建数据库（如果不存在）
|
|||
|
|
2. 执行 Alembic 迁移，确保表结构最新
|
|||
|
|
3. 补全 ORM 模型中定义但数据库不存在的表
|
|||
|
|
4. 补全已有表中缺失的字段（只 ADD COLUMN，不 MODIFY/DELETE）
|
|||
|
|
5. 执行种子数据初始化
|
|||
|
|
|
|||
|
|
被调用方：backend.app.main 的 lifespan 函数。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
import subprocess
|
|||
|
|
import sys
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
from sqlalchemy import create_engine, text
|
|||
|
|
from sqlalchemy.exc import OperationalError
|
|||
|
|
|
|||
|
|
from backend.app.core.config import get_settings
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[3]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def auto_init_database() -> None:
|
|||
|
|
"""自动初始化数据库。
|
|||
|
|
|
|||
|
|
启动时调用，按顺序执行：
|
|||
|
|
1. 检查数据库是否存在，不存在则创建
|
|||
|
|
2. 执行 Alembic 迁移
|
|||
|
|
3. 检查并补全所有表（Alembic 迁移可能遗漏新增模型）
|
|||
|
|
4. 执行种子数据初始化
|
|||
|
|
"""
|
|||
|
|
settings = get_settings()
|
|||
|
|
|
|||
|
|
# 1. 确保数据库存在
|
|||
|
|
_ensure_database_exists(settings)
|
|||
|
|
|
|||
|
|
# 2. 执行 Alembic 迁移
|
|||
|
|
_run_alembic_migrations()
|
|||
|
|
|
|||
|
|
# 3. 补全检查：确保所有 ORM 模型对应的表都已创建
|
|||
|
|
_create_tables_directly()
|
|||
|
|
|
|||
|
|
# 4. 补全已有表的缺失字段
|
|||
|
|
_sync_columns()
|
|||
|
|
|
|||
|
|
# 5. 初始化种子数据
|
|||
|
|
_run_bootstrap_data()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _ensure_database_exists(settings) -> None:
|
|||
|
|
"""检查数据库是否存在，不存在则创建。"""
|
|||
|
|
# 先尝试连接目标数据库
|
|||
|
|
db_url = (
|
|||
|
|
f"mysql+pymysql://{settings.mysql_user}:{settings.mysql_password}"
|
|||
|
|
f"@{settings.mysql_host}:{settings.mysql_port}/{settings.mysql_database}"
|
|||
|
|
)
|
|||
|
|
engine = create_engine(db_url)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
with engine.connect() as conn:
|
|||
|
|
conn.execute(text("SELECT 1"))
|
|||
|
|
logger.info("数据库 '%s' 已存在，连接正常", settings.mysql_database)
|
|||
|
|
engine.dispose()
|
|||
|
|
return
|
|||
|
|
except OperationalError as e:
|
|||
|
|
# 如果是"未知数据库"错误，则创建数据库
|
|||
|
|
if "Unknown database" in str(e) or "1049" in str(e):
|
|||
|
|
logger.info("数据库 '%s' 不存在，正在创建...", settings.mysql_database)
|
|||
|
|
engine.dispose()
|
|||
|
|
_create_database(settings)
|
|||
|
|
else:
|
|||
|
|
logger.error("数据库连接失败: %s", e)
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _create_database(settings) -> None:
|
|||
|
|
"""创建数据库。"""
|
|||
|
|
# 连接到 MySQL 服务器（不指定数据库）
|
|||
|
|
server_url = (
|
|||
|
|
f"mysql+pymysql://{settings.mysql_user}:{settings.mysql_password}"
|
|||
|
|
f"@{settings.mysql_host}:{settings.mysql_port}"
|
|||
|
|
)
|
|||
|
|
engine = create_engine(server_url)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
with engine.connect() as conn:
|
|||
|
|
conn.execute(text(
|
|||
|
|
f"CREATE DATABASE IF NOT EXISTS `{settings.mysql_database}` "
|
|||
|
|
f"CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
|
|||
|
|
))
|
|||
|
|
conn.commit()
|
|||
|
|
logger.info("数据库 '%s' 创建成功", settings.mysql_database)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("创建数据库失败: %s", e)
|
|||
|
|
raise
|
|||
|
|
finally:
|
|||
|
|
engine.dispose()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _run_alembic_migrations() -> None:
|
|||
|
|
"""执行 Alembic 迁移，确保表结构最新。"""
|
|||
|
|
logger.info("正在执行数据库迁移...")
|
|||
|
|
|
|||
|
|
alembic_ini = PROJECT_ROOT / "backend" / "alembic.ini"
|
|||
|
|
if not alembic_ini.exists():
|
|||
|
|
logger.warning("未找到 alembic.ini，尝试使用 SQLAlchemy 直接创建表")
|
|||
|
|
_create_tables_directly()
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
result = subprocess.run(
|
|||
|
|
[sys.executable, "-m", "alembic", "upgrade", "head"],
|
|||
|
|
cwd=str(PROJECT_ROOT / "backend"),
|
|||
|
|
capture_output=True,
|
|||
|
|
text=True,
|
|||
|
|
timeout=60,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if result.returncode == 0:
|
|||
|
|
logger.info("数据库迁移完成")
|
|||
|
|
if result.stdout.strip():
|
|||
|
|
logger.debug("迁移输出: %s", result.stdout.strip())
|
|||
|
|
else:
|
|||
|
|
logger.error("数据库迁移失败，尝试直接创建表: %s", result.stderr[:200])
|
|||
|
|
_create_tables_directly()
|
|||
|
|
except subprocess.TimeoutExpired:
|
|||
|
|
logger.error("数据库迁移超时，尝试直接创建表")
|
|||
|
|
_create_tables_directly()
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("执行迁移异常，尝试直接创建表: %s", e)
|
|||
|
|
_create_tables_directly()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _create_tables_directly() -> None:
|
|||
|
|
"""检查并补全所有 ORM 模型对应的数据库表。
|
|||
|
|
|
|||
|
|
Base.metadata.create_all 默认 checkfirst=True，
|
|||
|
|
只创建不存在的表，已存在的表不会被修改或删除。
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
from sqlalchemy import inspect
|
|||
|
|
|
|||
|
|
from backend.app.db import Base, engine
|
|||
|
|
# 显式导入所有模型类，确保每个 ORM 类都注册到 Base.metadata
|
|||
|
|
from backend.app.models import ( # noqa: F401
|
|||
|
|
AIRecognitionLog,
|
|||
|
|
AuditLog,
|
|||
|
|
Customer,
|
|||
|
|
CustomerArrears,
|
|||
|
|
FileAttachment,
|
|||
|
|
LogisticsTask,
|
|||
|
|
LogisticsTrace,
|
|||
|
|
LogisticsWaybill,
|
|||
|
|
Menu,
|
|||
|
|
OrderSupplierTextLog,
|
|||
|
|
PerformanceStatCache,
|
|||
|
|
Product,
|
|||
|
|
ProductCategory,
|
|||
|
|
ProductPriceTier,
|
|||
|
|
ProductPricingRule,
|
|||
|
|
Role,
|
|||
|
|
RoleMenu,
|
|||
|
|
SalesOrder,
|
|||
|
|
SalesOrderApproveLog,
|
|||
|
|
SalesOrderItem,
|
|||
|
|
Supplier,
|
|||
|
|
SupplierProductCost,
|
|||
|
|
SystemConfig,
|
|||
|
|
SystemReminder,
|
|||
|
|
User,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
orm_tables = set(Base.metadata.tables.keys())
|
|||
|
|
logger.info("ORM 已注册 %d 张表: %s", len(orm_tables), sorted(orm_tables))
|
|||
|
|
|
|||
|
|
# 记录创建前已存在的表
|
|||
|
|
inspector = inspect(engine)
|
|||
|
|
before = set(inspector.get_table_names())
|
|||
|
|
|
|||
|
|
# create_all: 已存在的跳过，不存在的创建
|
|||
|
|
Base.metadata.create_all(engine, checkfirst=True)
|
|||
|
|
|
|||
|
|
# 对比前后差异，确认结果
|
|||
|
|
after = set(inspect(engine).get_table_names())
|
|||
|
|
newly_created = (orm_tables - before) & after
|
|||
|
|
still_missing = orm_tables - after
|
|||
|
|
|
|||
|
|
if newly_created:
|
|||
|
|
logger.info("新创建 %d 张表: %s", len(newly_created), sorted(newly_created))
|
|||
|
|
if still_missing:
|
|||
|
|
logger.error("以下表仍然缺失，请检查: %s", sorted(still_missing))
|
|||
|
|
else:
|
|||
|
|
logger.info("所有 %d 张表均已就绪", len(orm_tables))
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("检查/创建表失败: %s", e)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _sync_columns() -> None:
|
|||
|
|
"""检查已有表是否缺少 ORM 模型中定义的字段，自动补全。
|
|||
|
|
|
|||
|
|
只做 ADD COLUMN，不做 MODIFY/DELETE，确保安全。
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
from sqlalchemy import inspect, text
|
|||
|
|
|
|||
|
|
from backend.app.db import Base, engine
|
|||
|
|
|
|||
|
|
inspector = inspect(engine)
|
|||
|
|
existing_tables = set(inspector.get_table_names())
|
|||
|
|
added_count = 0
|
|||
|
|
|
|||
|
|
for table_name, table in Base.metadata.tables.items():
|
|||
|
|
if table_name not in existing_tables:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
db_columns = {col["name"] for col in inspector.get_columns(table_name)}
|
|||
|
|
|
|||
|
|
for col in table.columns:
|
|||
|
|
if col.name not in db_columns:
|
|||
|
|
# 生成 ALTER TABLE ADD COLUMN 语句
|
|||
|
|
col_type = col.type.compile(dialect=engine.dialect)
|
|||
|
|
nullable = "NULL" if col.nullable else "NOT NULL"
|
|||
|
|
default = ""
|
|||
|
|
if col.default is not None:
|
|||
|
|
default_val = col.default.arg
|
|||
|
|
if callable(default_val):
|
|||
|
|
# 函数默认值（如 now()）跳过，让数据库处理
|
|||
|
|
default = ""
|
|||
|
|
elif isinstance(default_val, str):
|
|||
|
|
default = f"DEFAULT '{default_val}'"
|
|||
|
|
else:
|
|||
|
|
default = f"DEFAULT {default_val}"
|
|||
|
|
|
|||
|
|
sql = f"ALTER TABLE `{table_name}` ADD COLUMN `{col.name}` {col_type} {nullable} {default}"
|
|||
|
|
logger.info("补字段: %s.%s", table_name, col.name)
|
|||
|
|
|
|||
|
|
with engine.connect() as conn:
|
|||
|
|
conn.execute(text(sql))
|
|||
|
|
conn.commit()
|
|||
|
|
added_count += 1
|
|||
|
|
|
|||
|
|
if added_count:
|
|||
|
|
logger.info("共补全 %d 个缺失字段", added_count)
|
|||
|
|
else:
|
|||
|
|
logger.info("所有表字段均已同步")
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("同步表字段失败: %s", e)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _run_bootstrap_data() -> None:
|
|||
|
|
"""执行种子数据初始化。"""
|
|||
|
|
logger.info("正在初始化种子数据...")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
from backend.scripts.bootstrap_data import main as bootstrap_main
|
|||
|
|
bootstrap_main()
|
|||
|
|
logger.info("种子数据初始化完成")
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("种子数据初始化失败: %s", e)
|
|||
|
|
# 种子数据失败不阻止启动
|