34 lines
1.5 KiB
Python
34 lines
1.5 KiB
Python
"""老库补列：create_all 只会建新表，不会给已存在的表加字段。
|
||
|
||
模型新增可空字段后（如定开项目的验收时间 / 开票 / 回款），启动时自动
|
||
`ALTER TABLE ... ADD COLUMN` 补齐，不需要手写迁移脚本。
|
||
|
||
用法：在 `Base.metadata.create_all()` 之后调用 `ensure_columns(engine, Base.metadata)`。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from sqlalchemy import inspect, text
|
||
from sqlalchemy.engine import Engine
|
||
from sqlalchemy.sql.schema import MetaData
|
||
|
||
|
||
def ensure_columns(engine: Engine, metadata: MetaData) -> list[str]:
|
||
"""给已存在的表补上模型里新增的可空字段，返回补过的列名（表.列）。"""
|
||
insp = inspect(engine)
|
||
added: list[str] = []
|
||
for table in metadata.tables.values():
|
||
if not insp.has_table(table.name):
|
||
continue # 表还不存在 → create_all 会按最新模型建，这里不用管
|
||
existing = {c["name"] for c in insp.get_columns(table.name)}
|
||
for col in table.columns:
|
||
if col.name in existing:
|
||
continue
|
||
if not col.nullable:
|
||
print(f"[warn] 跳过高字段 {table.name}.{col.name}：新增列需允许为空")
|
||
continue
|
||
ddl = col.type.compile(engine.dialect)
|
||
with engine.begin() as conn:
|
||
conn.execute(text(f"ALTER TABLE {table.name} ADD COLUMN {col.name} {ddl} NULL"))
|
||
added.append(f"{table.name}.{col.name}")
|
||
return added
|