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()
|