zkxw/backend/scripts/migrate_qc_material_stock_check_requisition.py

#!/usr/bin/env python3
"""执行质检物料盘点登记相关 DB 迁移（requisition_status + usage_id/batch_id）。

用法（在 backend 目录）：
  source venv/bin/activate
  python scripts/migrate_qc_material_stock_check_requisition.py
"""
from __future__ import annotations

from sqlalchemy import text

from app.core.database import engine

MIGRATIONS = [
    (
        "biz_qc_material_usage.requisition_status",
        """ALTER TABLE `biz_qc_material_usage`
  ADD COLUMN `requisition_status` varchar(16) NOT NULL DEFAULT 'open'
    COMMENT '领用状态：open-待登记/有剩余 closed-已使用完'
    AFTER `purpose`""",
        """UPDATE `biz_qc_material_usage`
SET `requisition_status` = 'open'
WHERE `source_type` = 'manual_out'
  AND `change_type` = 'use'
  AND (`requisition_status` IS NULL OR `requisition_status` = '')""",
    ),
    (
        "biz_qc_material_stock_check_unused_line.usage_id",
        """ALTER TABLE `biz_qc_material_stock_check_unused_line`
  ADD COLUMN `usage_id` varchar(64) DEFAULT NULL COMMENT '关联领用记录ID(biz_qc_material_usage.id)' AFTER `check_id`""",
        None,
    ),
    (
        "biz_qc_material_stock_check_unused_line.batch_id",
        """ALTER TABLE `biz_qc_material_stock_check_unused_line`
  ADD COLUMN `batch_id` varchar(64) DEFAULT NULL COMMENT '关联入库批次ID(手工登记批次检索)' AFTER `usage_id`""",
        None,
    ),
    (
        "idx_qc_check_unused_line_usage",
        """ALTER TABLE `biz_qc_material_stock_check_unused_line`
  ADD KEY `idx_qc_check_unused_line_usage` (`usage_id`, `del_flag`)""",
        None,
    ),
]


def _run(conn, sql: str, label: str) -> None:
    try:
        conn.execute(text(sql))
        print(f"OK  {label}")
    except Exception as e:
        msg = str(e)
        if "Duplicate column name" in msg or "1060" in msg:
            print(f"SKIP {label} (column exists)")
        elif "Duplicate key name" in msg or "1061" in msg:
            print(f"SKIP {label} (index exists)")
        else:
            raise


def main() -> None:
    with engine.begin() as conn:
        for label, ddl, follow_up in MIGRATIONS:
            _run(conn, ddl, label)
            if follow_up:
                _run(conn, follow_up, f"{label} data backfill")
    print("Migration completed.")


if __name__ == "__main__":
    main()