383f704b53
Cleaned up various files by removing unused or redundant import statements to improve code readability and reduce potential confusion. No functional changes were made.
78 lines
3.1 KiB
Python
78 lines
3.1 KiB
Python
# mmc/src/common/database/db_migration.py
|
|
|
|
from sqlalchemy import inspect, text
|
|
from src.common.database.sqlalchemy_models import Base, get_engine
|
|
from src.common.logger import get_logger
|
|
|
|
logger = get_logger("db_migration")
|
|
|
|
|
|
def check_and_migrate_database():
|
|
"""
|
|
检查数据库结构并自动迁移（添加缺失的表和列）。
|
|
"""
|
|
logger.info("正在检查数据库结构并执行自动迁移...")
|
|
engine = get_engine()
|
|
inspector = inspect(engine)
|
|
|
|
# 1. 获取数据库中所有已存在的表名
|
|
db_table_names = set(inspector.get_table_names())
|
|
|
|
# 2. 遍历所有在代码中定义的模型
|
|
for table_name, table in Base.metadata.tables.items():
|
|
logger.debug(f"正在检查表: {table_name}")
|
|
|
|
# 3. 如果表不存在，则创建它
|
|
if table_name not in db_table_names:
|
|
logger.info(f"表 '{table_name}' 不存在，正在创建...")
|
|
try:
|
|
table.create(engine)
|
|
logger.info(f"表 '{table_name}' 创建成功。")
|
|
except Exception as e:
|
|
logger.error(f"创建表 '{table_name}' 失败: {e}")
|
|
continue
|
|
|
|
# 4. 如果表已存在，则检查并添加缺失的列
|
|
db_columns = {col["name"] for col in inspector.get_columns(table_name)}
|
|
model_columns = {col.name for col in table.c}
|
|
|
|
missing_columns = model_columns - db_columns
|
|
if not missing_columns:
|
|
logger.debug(f"表 '{table_name}' 结构一致，无需修改。")
|
|
continue
|
|
|
|
logger.info(f"在表 '{table_name}' 中发现缺失的列: {', '.join(missing_columns)}")
|
|
with engine.connect() as connection:
|
|
trans = connection.begin()
|
|
try:
|
|
for column_name in missing_columns:
|
|
column = table.c[column_name]
|
|
|
|
# 构造并执行 ALTER TABLE 语句
|
|
try:
|
|
column_type = column.type.compile(engine.dialect)
|
|
sql = f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}"
|
|
|
|
# 添加默认值和非空约束的处理
|
|
if column.default is not None:
|
|
default_value = column.default.arg
|
|
if isinstance(default_value, str):
|
|
sql += f" DEFAULT '{default_value}'"
|
|
else:
|
|
sql += f" DEFAULT {default_value}"
|
|
|
|
if not column.nullable:
|
|
sql += " NOT NULL"
|
|
|
|
connection.execute(text(sql))
|
|
logger.info(f"成功向表 '{table_name}' 添加列 '{column_name}'。")
|
|
except Exception as e:
|
|
logger.error(f"向表 '{table_name}' 添加列 '{column_name}' 失败: {e}")
|
|
|
|
trans.commit()
|
|
except Exception as e:
|
|
logger.error(f"在表 '{table_name}' 添加列时发生错误，事务已回滚: {e}")
|
|
trans.rollback()
|
|
|
|
logger.info("数据库结构检查与自动迁移完成。")
|