87 lines
3.4 KiB
Python
87 lines
3.4 KiB
Python
"""波次5.4：业务必填列 NOT NULL 对齐（幂等，仅 PostgreSQL）。
|
||
|
||
背景：ORM 声明 NOT NULL、线上库却允许 NULL 的业务必填列（多为父记录外键），
|
||
与「严谨的业务逻辑关系」相悖。本脚本将这类列收紧为 NOT NULL，使库约束与 ORM 一致。
|
||
|
||
安全前提：先校验该列无 NULL 值（有 NULL 则跳过并上报），无 NULL 才执行 ALTER，
|
||
故 SET NOT NULL 不会因存量数据失败。相关外键均为 NO ACTION / CASCADE（无 SET NULL），
|
||
与 NOT NULL 兼容。
|
||
|
||
幂等：已 NOT NULL 的列跳过；可重复执行。
|
||
用法（backend/ 目录）：
|
||
python scripts/align_columns.py --check # 只读预览
|
||
python scripts/align_columns.py # 执行
|
||
"""
|
||
import argparse
|
||
import os
|
||
import sys
|
||
|
||
BACKEND_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
if BACKEND_DIR not in sys.path:
|
||
sys.path.insert(0, BACKEND_DIR)
|
||
|
||
from sqlalchemy import create_engine, text
|
||
|
||
from app.config.setting import settings
|
||
|
||
# ORM 声明必填（NOT NULL）、线上库却可空的业务列 → 收紧
|
||
COLUMNS = [
|
||
("bre_pedigree", "combination_id"),
|
||
("bre_pedigree", "child_code"),
|
||
("bre_selection_result", "tree_id"),
|
||
("bre_selection_rule", "target_id"),
|
||
("bre_prediction_value", "prediction_id"),
|
||
("bre_treatment", "trial_study_id"),
|
||
("bre_tree_evaluation", "tree_id"),
|
||
("bre_trial_study", "trial_id"),
|
||
("bre_trial_study_entry", "trial_study_id"),
|
||
("bre_trial_study_entry", "entry_number"),
|
||
("bre_trial_study_entry", "germplasm_id"),
|
||
]
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description="业务必填列 NOT NULL 对齐")
|
||
parser.add_argument("--check", action="store_true", help="仅打印将执行的 DDL，不落库")
|
||
args = parser.parse_args()
|
||
|
||
engine = create_engine(settings.DB_URI)
|
||
print(f"# 连接: {settings.DB_URI}")
|
||
|
||
# 每列用单事务：读锁（count）与 ALTER（ACCESS EXCLUSIVE）在同一事务内升级，
|
||
# 避免跨连接自锁（读连接持锁阻塞另一连接的 ALTER）。
|
||
for t, c in COLUMNS:
|
||
with engine.begin() as tx:
|
||
try:
|
||
row = tx.execute(text(
|
||
"SELECT is_nullable FROM information_schema.columns "
|
||
"WHERE table_name=:t AND column_name=:c"
|
||
), {"t": t, "c": c}).fetchone()
|
||
except Exception as e: # noqa: BLE001
|
||
print(f"# [跳过] {t}.{c}: 查列失败 {e!s}")
|
||
continue
|
||
if row is None:
|
||
print(f"# [跳过] {t}.{c}: 列不存在")
|
||
continue
|
||
if row[0] == "NO":
|
||
print(f"# [已存在] {t}.{c} 已 NOT NULL")
|
||
continue
|
||
nulls = tx.execute(text(f"SELECT count(*) FROM {t} WHERE {c} IS NULL")).scalar()
|
||
if nulls:
|
||
print(f"# [跳过] {t}.{c}: 存在 {nulls} 个 NULL 值，拒绝收紧")
|
||
continue
|
||
ddl = f"ALTER TABLE {t} ALTER COLUMN {c} SET NOT NULL"
|
||
print(" " + ddl)
|
||
if not args.check:
|
||
try:
|
||
tx.execute(text(ddl))
|
||
print("OK " + ddl)
|
||
except Exception as e: # noqa: BLE001
|
||
print("ERR " + ddl + " -> " + str(e)[:160])
|
||
|
||
print("# --check 模式：未执行任何变更。" if args.check else "# 完成。")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|