zkxw/backend/scripts/seed_processes.py

#!/usr/bin/env python3
"""批量初始化工艺库测试数据（biz_process），已存在工序代码则跳过。

用法（在 backend 目录）：
  source venv/bin/activate
  python scripts/seed_processes.py
  python scripts/seed_processes.py --tenant-no zkxw

依赖：biz_process_category 中已有对应分类编码（可先运行 reset_process_categories.py）。
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Dict, Optional

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from app.core.database import SessionLocal
from app.crud import process as process_svc
from app.models.models import BizProcessCategory

ADMIN_ID = "d3000000000000000000000000000001"

# (process_code, process_name, category_code, description, equipment_name, estimated_duration_h, order_num, remark)
SAMPLE_PROCESSES = [
    (
        "PR-YBJS-001",
        "样本接收登记",
        "PC-YBJS",
        "核对样本信息、登记条码并分配接收批次",
        "样本接收台",
        "0.5",
        1,
        "样本接收首道工序",
    ),
    (
        "PR-YBJS-002",
        "样本初检",
        "PC-YBJS",
        "外观检查、温度记录及初筛不合格样本隔离",
        "生物安全柜",
        "0.75",
        2,
        None,
    ),
    (
        "PR-XBYP-001",
        "外周血分离",
        "PC-XBYP",
        "密度梯度离心分离单个核细胞",
        "台式离心机",
        "1",
        3,
        None,
    ),
    (
        "PR-XBYP-002",
        "细胞培养扩增",
        "PC-XBYP",
        "接种至培养体系并进行扩增培养",
        "CO2培养箱",
        "8",
        4,
        "长周期工序",
    ),
    (
        "PR-XBYP-003",
        "换液传代",
        "PC-XBYP",
        "更换培养基并按密度传代",
        "生物安全柜",
        "1.5",
        5,
        None,
    ),
    (
        "PR-ZJJC-001",
        "细胞活性检测",
        "PC-ZJJC",
        "台盼蓝染色法检测细胞活率",
        "倒置显微镜",
        "0.5",
        6,
        None,
    ),
    (
        "PR-ZJJC-002",
        "支原体检测",
        "PC-ZJJC",
        "PCR法检测支原体污染",
        "PCR仪",
        "3",
        7,
        None,
    ),
    (
        "PR-DCGL-001",
        "程序降温冻存",
        "PC-DCGL",
        "按程序降温曲线完成细胞冻存",
        "程序降温仪",
        "2.5",
        8,
        None,
    ),
    (
        "PR-DCGL-002",
        "冻存管入库",
        "PC-DCGL",
        "扫码登记并转入液氮罐指定位置",
        "扫码枪",
        "0.33",
        9,
        None,
    ),
    (
        "PR-BZZF-001",
        "成品包装发运",
        "PC-BZZF",
        "成品核对、冷链包装及发运交接",
        "包装台",
        "0.75",
        10,
        "末道工序",
    ),
]


def _load_category_map(db, tenant_no: str) -> Dict[str, str]:
    rows = (
        db.query(BizProcessCategory)
        .filter(BizProcessCategory.tenant_no == tenant_no, BizProcessCategory.del_flag == 0)
        .all()
    )
    return {row.category_code: row.id for row in rows}


def seed(tenant_no: str = "zkxw", user_id: str = ADMIN_ID) -> None:
    db = SessionLocal()
    created = 0
    skipped = 0
    missing_categories: set[str] = set()
    try:
        category_map = _load_category_map(db, tenant_no)
        for (
            code,
            name,
            category_code,
            description,
            equipment_name,
            estimated_duration,
            order_num,
            remark,
        ) in SAMPLE_PROCESSES:
            if process_svc.process_code_exists(db, code, tenant_no):
                skipped += 1
                continue
            category_id: Optional[str] = category_map.get(category_code)
            if category_code and not category_id:
                missing_categories.add(category_code)
            process_svc.create_process(
                db,
                process_code=code,
                process_name=name,
                category_id=category_id,
                description=description,
                equipment_ids=None,
                estimated_duration=estimated_duration,
                tenant_no=tenant_no,
                order_num=order_num,
                remark=remark,
                user_id=user_id,
            )
            created += 1
        print(
            f"工艺库测试数据初始化完成（租户 {tenant_no}）："
            f"新增 {created} 条，跳过 {skipped} 条（工序代码已存在）。"
        )
        if missing_categories:
            print(
                "以下工序分类编码未找到，对应工艺已写入但 category_id 为空："
                f" {', '.join(sorted(missing_categories))}"
            )
            print("可先执行：python scripts/reset_process_categories.py")
    finally:
        db.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="初始化 biz_process 测试数据")
    parser.add_argument("--tenant-no", default="zkxw", help="租户编号，默认 zkxw")
    parser.add_argument("--user-id", default=ADMIN_ID, help="create_by / update_by 用户 id")
    args = parser.parse_args()
    seed(tenant_no=args.tenant_no, user_id=args.user_id)


if __name__ == "__main__":
    main()