zkxw/backend/scripts/migrate_process_equipment_id.py

#!/usr/bin/env python3
"""按 equipment_name 精确匹配回填 equipment_ids。用法: python scripts/migrate_process_equipment_id.py [--dry-run]"""
import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.core.database import SessionLocal
from app.models.models import BizEquipment, BizProcess, BizProcessRouteStep, BizProductionTaskStep

TABLES = [
    ("biz_process", BizProcess),
    ("biz_process_route_step", BizProcessRouteStep),
    ("biz_production_task_step", BizProductionTaskStep),
]


def _split_names(name: str) -> list[str]:
    text = (name or "").strip()
    if not text:
        return []
    for sep in ("、", "，", ";", "；", "|"):
        text = text.replace(sep, ",")
    return [x.strip() for x in text.split(",") if x.strip()]


def migrate_table(db, model, dry_run: bool) -> int:
    updated = 0
    rows = db.query(model).filter(model.del_flag == 0, model.equipment_ids.is_(None)).all()
    for row in rows:
        names = _split_names(row.equipment_name)
        if not names:
            continue
        matched_ids: list[str] = []
        matched_names: list[str] = []
        for name in names:
            eq = (
                db.query(BizEquipment)
                .filter(
                    BizEquipment.tenant_no == row.tenant_no,
                    BizEquipment.del_flag == 0,
                    BizEquipment.equipment_name == name,
                )
                .first()
            )
            if not eq:
                continue
            matched_ids.append(eq.id)
            matched_names.append(eq.equipment_name or name)
        if not matched_ids:
            continue
        if not dry_run:
            row.equipment_ids = ",".join(matched_ids)
            row.equipment_name = ",".join(matched_names)
        updated += 1
    return updated


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    db = SessionLocal()
    try:
        for label, model in TABLES:
            n = migrate_table(db, model, args.dry_run)
            print(f"{label}: matched {n}")
        if not args.dry_run:
            db.commit()
            print("committed")
        else:
            print("dry-run, no commit")
    finally:
        db.close()


if __name__ == "__main__":
    main()