shanxia2.0/backend/scripts/migrate_pl_plan_to_multi_stop.py

#!/usr/bin/env python3
"""将历史运输计划迁移为「1 装货段 + 1 送货段」多段模型。

规则：
  - 无 pl_plan_load 的计划：supplier → 装货段，customer/destination → 送货段 seq=1
  - 复制 pl_plan_product 到装货/送货段产品行（plan_qty = actual_qty 或 quantity）
  - pl_plan_contract：supplier → side=load，demander → side=delivery，并写入 stop_id
  - 已归档计划：inbound_* 写入装货段产品，outbound_* 写入送货段产品

用法：
  cd backend && source venv/bin/activate
  python scripts/migrate_pl_plan_to_multi_stop.py --dry-run

  ENV_FILE=.env.production python scripts/migrate_pl_plan_to_multi_stop.py

环境变量：
  ENV_FILE / DATABASE_URL  目标库（由 app Settings 读取）
"""
from __future__ import annotations

import argparse
import os
import sys
from dataclasses import dataclass, field
from datetime import datetime
from decimal import Decimal
from pathlib import Path
from typing import Optional

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

from sqlalchemy.orm import Session

from app.core.database import SessionLocal
from app.core.utils import generate_uuid
from app.models.pl_models import (
    PlPlan,
    PlPlanContract,
    PlPlanDelivery,
    PlPlanDeliveryProduct,
    PlPlanLoad,
    PlPlanLoadProduct,
    PlPlanProduct,
)
from app.services.pl_plan_stop_rules import resolve_stop_dibang_defaults
from app.services.pl_plan_service import PLAN_CONTRACT_SIDE_DEMANDER, PLAN_CONTRACT_SIDE_SUPPLIER
from app.services.pl_plan_stop_service import PLAN_CONTRACT_SIDE_DELIVERY, PLAN_CONTRACT_SIDE_LOAD


@dataclass
class MigrationStats:
    plans_scanned: int = 0
    plans_migrated: int = 0
    loads_created: int = 0
    deliveries_created: int = 0
    load_products_created: int = 0
    delivery_products_created: int = 0
    contracts_updated: int = 0
    skipped_existing_load: int = 0
    warnings: list[str] = field(default_factory=list)


def _plan_qty(prod: PlPlanProduct) -> Optional[Decimal]:
    if prod.actual_qty is not None:
        return prod.actual_qty
    return prod.quantity


def _migrate_plan(db: Session, plan: PlPlan, stats: MigrationStats, dry_run: bool) -> None:
    existing_load = (
        db.query(PlPlanLoad)
        .filter(
            PlPlanLoad.plan_id == plan.id,
            PlPlanLoad.tenant_no == plan.tenant_no,
            PlPlanLoad.del_flag == 0,
        )
        .first()
    )
    if existing_load:
        stats.skipped_existing_load += 1
        return

    products = (
        db.query(PlPlanProduct)
        .filter(
            PlPlanProduct.plan_id == plan.id,
            PlPlanProduct.del_flag == 0,
        )
        .order_by(PlPlanProduct.order_num.asc(), PlPlanProduct.create_time.asc())
        .all()
    )

    now = datetime.now()
    archived = (plan.archive_status or "").strip() == "archived"

    load_customer_name = (plan.supplier_name or "").strip() or None
    default_gen, default_dir = resolve_stop_dibang_defaults(load_customer_name, "load")
    load_id = generate_uuid()

    load_row = PlPlanLoad(
        id=load_id,
        plan_id=plan.id,
        customer_id=plan.supplier_id,
        customer_name=plan.supplier_name,
        address=plan.pickup_address,
        lng=plan.pickup_lng,
        lat=plan.pickup_lat,
        load_time=plan.pickup_time,
        contact=plan.pickup_contact,
        notes=plan.pickup_notes,
        generate_dibang=1 if default_gen else 0,
        dibang_direction=default_dir or "none",
        tenant_no=plan.tenant_no,
        order_num=0,
        remark=None,
        create_by=plan.create_by,
        create_dept_id=plan.create_dept_id or "",
        create_time=now,
        update_by=plan.update_by,
        update_time=now,
        del_flag=0,
    )

    delivery_customer_name = (plan.customer_name or plan.destination_name or "").strip() or None
    del_gen, del_dir = resolve_stop_dibang_defaults(delivery_customer_name, "delivery")
    delivery_id = generate_uuid()

    delivery_row = PlPlanDelivery(
        id=delivery_id,
        plan_id=plan.id,
        seq=1,
        customer_id=plan.customer_id or plan.destination_id,
        customer_name=plan.customer_name or plan.destination_name,
        address=plan.destination_address,
        lng=plan.destination_lng,
        lat=plan.destination_lat,
        delivery_time=plan.delivery_time,
        contact=plan.delivery_contact,
        notes=plan.destination_notes,
        generate_dibang=1 if del_gen else 0,
        dibang_direction=del_dir or "none",
        tenant_no=plan.tenant_no,
        order_num=0,
        remark=None,
        create_by=plan.create_by,
        create_dept_id=plan.create_dept_id or "",
        create_time=now,
        update_by=plan.update_by,
        update_time=now,
        del_flag=0,
    )

    load_products: list[PlPlanLoadProduct] = []
    delivery_products: list[PlPlanDeliveryProduct] = []

    for idx, prod in enumerate(products):
        load_prod = PlPlanLoadProduct(
            id=generate_uuid(),
            load_id=load_id,
            plan_id=plan.id,
            plan_product_id=prod.id,
            product_id=prod.product_id,
            product_no=prod.product_no,
            product_name=prod.product_name,
            model=prod.model,
            unit=prod.unit or "吨",
            package_spec=prod.package_spec,
            delivery_method=prod.delivery_method,
            plan_qty=_plan_qty(prod),
            inbound_qty=prod.inbound_qty if archived else None,
            inbound_price=prod.inbound_price if archived else None,
            inbound_amount=prod.inbound_amount if archived else None,
            tenant_no=plan.tenant_no,
            order_num=idx,
            remark=prod.remark,
            create_by=plan.create_by,
            create_dept_id=plan.create_dept_id or "",
            create_time=now,
            update_by=plan.update_by,
            update_time=now,
            del_flag=0,
        )
        load_products.append(load_prod)

        delivery_prod = PlPlanDeliveryProduct(
            id=generate_uuid(),
            delivery_id=delivery_id,
            plan_id=plan.id,
            plan_product_id=prod.id,
            product_id=prod.product_id,
            product_no=prod.product_no,
            product_name=prod.product_name,
            model=prod.model,
            unit=prod.unit or "吨",
            package_spec=prod.package_spec,
            delivery_method=prod.delivery_method,
            plan_qty=_plan_qty(prod),
            outbound_qty=prod.outbound_qty if archived else None,
            outbound_price=prod.outbound_price if archived else None,
            outbound_amount=prod.outbound_amount if archived else None,
            tenant_no=plan.tenant_no,
            order_num=idx,
            remark=prod.remark,
            create_by=plan.create_by,
            create_dept_id=plan.create_dept_id or "",
            create_time=now,
            update_by=plan.update_by,
            update_time=now,
            del_flag=0,
        )
        delivery_products.append(delivery_prod)

    contract_rows = (
        db.query(PlPlanContract)
        .filter(
            PlPlanContract.plan_id == plan.id,
            PlPlanContract.tenant_no == plan.tenant_no,
            PlPlanContract.del_flag == 0,
        )
        .all()
    )

    contracts_to_update: list[PlPlanContract] = []
    for row in contract_rows:
        if (row.stop_id or "").strip():
            continue
        side = (row.side or "").strip()
        if side == PLAN_CONTRACT_SIDE_SUPPLIER:
            row.side = PLAN_CONTRACT_SIDE_LOAD
            row.stop_id = load_id
            contracts_to_update.append(row)
        elif side == PLAN_CONTRACT_SIDE_DEMANDER:
            row.side = PLAN_CONTRACT_SIDE_DELIVERY
            row.stop_id = delivery_id
            contracts_to_update.append(row)

    if not dry_run:
        db.add(load_row)
        db.add(delivery_row)
        for lp in load_products:
            db.add(lp)
        for dp in delivery_products:
            db.add(dp)
        for row in contracts_to_update:
            row.update_by = plan.update_by
            row.update_time = now

    stats.plans_migrated += 1
    stats.loads_created += 1
    stats.deliveries_created += 1
    stats.load_products_created += len(load_products)
    stats.delivery_products_created += len(delivery_products)
    stats.contracts_updated += len(contracts_to_update)


def run_migration(tenant_no: Optional[str], dry_run: bool) -> MigrationStats:
    stats = MigrationStats()
    db = SessionLocal()
    try:
        query = db.query(PlPlan).filter(PlPlan.del_flag == 0)
        if tenant_no:
            query = query.filter(PlPlan.tenant_no == tenant_no)
        plans = query.order_by(PlPlan.create_time.asc()).all()
        stats.plans_scanned = len(plans)

        for plan in plans:
            _migrate_plan(db, plan, stats, dry_run)

        if dry_run:
            db.rollback()
        else:
            db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()
    return stats


def print_stats(stats: MigrationStats, dry_run: bool) -> None:
    mode = "DRY-RUN" if dry_run else "COMMITTED"
    print(f"=== migrate_pl_plan_to_multi_stop ({mode}) ===")
    print(f"plans_scanned:           {stats.plans_scanned}")
    print(f"plans_migrated:          {stats.plans_migrated}")
    print(f"skipped_existing_load:   {stats.skipped_existing_load}")
    print(f"loads_created:           {stats.loads_created}")
    print(f"deliveries_created:      {stats.deliveries_created}")
    print(f"load_products_created:   {stats.load_products_created}")
    print(f"delivery_products_created: {stats.delivery_products_created}")
    print(f"contracts_updated:       {stats.contracts_updated}")
    if stats.warnings:
        print("warnings:")
        for w in stats.warnings:
            print(f"  - {w}")


def main() -> None:
    parser = argparse.ArgumentParser(description="迁移运输计划至装货/送货多段模型")
    parser.add_argument("--dry-run", action="store_true", help="仅统计不写库")
    parser.add_argument("--tenant-no", default=None, help="仅迁移指定租户")
    parser.add_argument(
        "--env-file",
        default=os.environ.get("ENV_FILE"),
        help="环境配置文件路径（如 .env.production）",
    )
    args = parser.parse_args()

    if args.env_file:
        os.environ["ENV_FILE"] = args.env_file

    stats = run_migration(args.tenant_no, args.dry_run)
    print_stats(stats, args.dry_run)


if __name__ == "__main__":
    main()