shanxia2.0/backend/scripts/migrate_dibang_data.py

#!/usr/bin/env python3
"""Migrate legacy dibang business data into shanxia2.0 db_* tables.

Reads from the original three tables (purchase_plans, inbound_outbound,
weighing_slips), preserves legacy int IDs as VARCHAR primary keys (e.g. 12 → "12"),
remaps foreign keys, and writes portal public fields (tenant_no, del_flag, etc.).

Usage:
  cd backend && source venv/bin/activate
  python scripts/migrate_dibang_data.py --dry-run \\
    --source-url mysql+pymysql://user:pass@127.0.0.1:3306/dibang

  python scripts/migrate_dibang_data.py \\
    --source-url mysql+pymysql://user:pass@127.0.0.1:3306/dibang \\
    --mapping-out ./dibang-id-map.json

Environment:
  DIBANG_SOURCE_DATABASE_URL  default for --source-url
  DATABASE_URL                target DB (from app settings / .env)
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import date, datetime
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple

from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine

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

from app.core.config import settings

SOURCE_TABLES = {
    "purchase_plans": "purchase_plans",
    "inbound_outbound": "inbound_outbound",
    "weighing_slips": "weighing_slips",
}

TARGET_TABLES = {
    "purchase_plans": "db_purchase_plan",
    "inbound_outbound": "db_inbound_outbound",
    "weighing_slips": "db_weighing_slip",
}


@dataclass
class MigrationStats:
    source_counts: Dict[str, int] = field(default_factory=dict)
    inserted: Dict[str, int] = field(default_factory=lambda: defaultdict(int))
    skipped: Dict[str, int] = field(default_factory=lambda: defaultdict(int))
    warnings: List[str] = field(default_factory=list)
    orphan_fks: List[str] = field(default_factory=list)


def parse_date(value: Any) -> Optional[date]:
    if value is None or value == "":
        return None
    if isinstance(value, date) and not isinstance(value, datetime):
        return value
    text_value = str(value).strip()
    if not text_value:
        return None
    for fmt in ("%Y-%m-%d", "%Y/%m/%d"):
        try:
            return datetime.strptime(text_value[:10], fmt).date()
        except ValueError:
            continue
    raise ValueError(f"无法解析日期: {value!r}")


def parse_datetime(value: Any) -> Optional[datetime]:
    if value is None or value == "":
        return None
    if isinstance(value, datetime):
        return value
    text_value = str(value).strip()
    if not text_value:
        return None
    formats: Iterable[Tuple[str, int]] = (
        ("%Y-%m-%d %H:%M:%S", 19),
        ("%Y-%m-%d %H:%M", 16),
        ("%Y-%m-%d", 10),
        ("%Y/%m/%d %H:%M:%S", 19),
        ("%Y/%m/%d", 10),
    )
    for fmt, length in formats:
        try:
            return datetime.strptime(text_value[:length], fmt)
        except ValueError:
            continue
    parsed = parse_date(text_value)
    return datetime.combine(parsed, datetime.min.time()) if parsed else None


def clip_text(value: Any, max_len: int, label: str, stats: MigrationStats) -> Optional[str]:
    if value is None:
        return None
    text_value = str(value)
    if len(text_value) <= max_len:
        return text_value
    stats.warnings.append(f"{label} 截断至 {max_len} 字符")
    return text_value[:max_len]


def fetch_rows(engine: Engine, table: str) -> List[dict]:
    with engine.connect() as conn:
        result = conn.execute(text(f"SELECT * FROM `{table}` ORDER BY `id`"))
        return [dict(row._mapping) for row in result]


def count_rows(engine: Engine, table: str) -> int:
    with engine.connect() as conn:
        return int(conn.execute(text(f"SELECT COUNT(*) FROM `{table}`")).scalar_one())


def map_id(
    table_key: str,
    legacy_id: Any,
    id_maps: Dict[str, Dict[int, str]],
    stats: MigrationStats,
    label: str,
) -> Optional[str]:
    if legacy_id is None:
        return None
    try:
        legacy_int = int(legacy_id)
    except (TypeError, ValueError):
        stats.orphan_fks.append(f"{label}: 非法 legacy id {legacy_id!r}")
        return None
    mapped = id_maps.get(table_key, {}).get(legacy_int)
    if mapped is None:
        stats.orphan_fks.append(f"{label}: legacy id {legacy_int} 无映射")
    return mapped


def legacy_id_to_str(legacy_id: int) -> str:
    """原地磅 int 主键原样保留，写入门户 VARCHAR(64) 字段。"""
    return str(int(legacy_id))


def build_id_maps(rows_by_table: Dict[str, List[dict]]) -> Dict[str, Dict[int, str]]:
    id_maps: Dict[str, Dict[int, str]] = {}
    for table_key, rows in rows_by_table.items():
        id_maps[table_key] = {
            int(row["id"]): legacy_id_to_str(int(row["id"])) for row in rows
        }
    return id_maps


def transform_plan(
    row: dict,
    new_id: str,
    tenant_no: str,
    stats: MigrationStats,
    *,
    last_weighing_slip_id: Optional[str] = None,
) -> dict:
    create_time = parse_datetime(row.get("created_at")) or datetime.now()
    plan_date = parse_date(row.get("plan_date")) or parse_date(row.get("start_date")) or create_time.date()
    return {
        "id": new_id,
        "plan_no": row.get("plan_no"),
        "plan_date": plan_date,
        "start_date": parse_date(row.get("start_date")),
        "end_date": parse_date(row.get("end_date")),
        "business_type": row["business_type"],
        "product": row["product"],
        "specification": row.get("specification"),
        "planned_quantity": float(row["planned_quantity"]),
        "vehicle_no": row.get("vehicle_no") or "",
        "customer_name": row.get("customer_name"),
        "execution_status": row.get("execution_status"),
        "executed_at": parse_datetime(row.get("executed_at")),
        "last_weighing_slip_id": last_weighing_slip_id,
        "executed_quantity": float(row.get("executed_quantity") or 0),
        "auto_generate_enabled": int(row.get("auto_generate_enabled") or 0),
        "auto_generate_end_month": row.get("auto_generate_end_month"),
        "auto_generate_last_month": row.get("auto_generate_last_month"),
        "tenant_no": tenant_no,
        "order_num": None,
        "remark": clip_text(row.get("remarks"), 512, "plan.remark", stats),
        "create_by": None,
        "create_dept_id": "",
        "create_time": create_time,
        "update_by": None,
        "update_time": create_time,
        "del_flag": int(row.get("is_deleted") or 0),
    }


def transform_inventory(
    row: dict,
    new_id: str,
    tenant_no: str,
    id_maps: Dict[str, Dict[int, str]],
    stats: MigrationStats,
) -> dict:
    create_time = parse_datetime(row.get("created_at")) or datetime.now()
    return {
        "id": new_id,
        "date": parse_date(row.get("date")) or create_time.date(),
        "partner_company": row["partner_company"],
        "business_type": row["business_type"],
        "product": row["product"],
        "quantity": float(row["quantity"]),
        "plan_no": row.get("plan_no"),
        "batch_no": row.get("batch_no"),
        "vehicle_no": row["vehicle_no"],
        "status": row.get("status"),
        "purchase_plan_id": map_id(
            "purchase_plans",
            row.get("purchase_plan_id"),
            id_maps,
            stats,
            f"inventory#{row.get('id')} purchase_plan_id",
        ),
        "tenant_no": tenant_no,
        "order_num": None,
        "remark": clip_text(row.get("remarks"), 512, "inventory.remark", stats),
        "create_by": None,
        "create_dept_id": "",
        "create_time": create_time,
        "update_by": None,
        "update_time": create_time,
        "del_flag": int(row.get("is_deleted") or 0),
    }


def transform_weighing(
    row: dict,
    new_id: str,
    tenant_no: str,
    id_maps: Dict[str, Dict[int, str]],
    stats: MigrationStats,
) -> dict:
    create_time = parse_datetime(row.get("created_at")) or datetime.now()
    gross_time = parse_datetime(row.get("gross_weight_time")) or create_time
    tare_time = parse_datetime(row.get("tare_weight_time"))
    return {
        "id": new_id,
        "vehicle_no": row["vehicle_no"],
        "plan_no": row.get("plan_no"),
        "cargo_name": row["cargo_name"],
        "business_type": row["business_type"],
        "batch_no": row.get("batch_no"),
        "supplier": row["supplier"],
        "consignee": row["consignee"],
        "gross_weight": float(row["gross_weight"]),
        "gross_weight_time": gross_time,
        "tare_weight": float(row["tare_weight"]),
        "tare_weight_time": tare_time,
        "deduction": float(row.get("deduction") or 0),
        "net_weight": float(row["net_weight"]),
        "planned_quantity": row.get("planned_quantity"),
        "executed_quantity": row.get("executed_quantity"),
        "purchase_plan_id": map_id(
            "purchase_plans",
            row.get("purchase_plan_id"),
            id_maps,
            stats,
            f"weighing#{row.get('id')} purchase_plan_id",
        ),
        "inbound_outbound_id": map_id(
            "inbound_outbound",
            row.get("inbound_outbound_id"),
            id_maps,
            stats,
            f"weighing#{row.get('id')} inbound_outbound_id",
        ),
        "weighing_point": row.get("weighing_point") or "山峡地磅",
        "current_quantity": row.get("current_quantity"),
        "is_voided": int(row.get("is_voided") or 0),
        "tenant_no": tenant_no,
        "order_num": None,
        "remark": clip_text(row.get("remarks"), 512, "weighing.remark", stats),
        "create_by": None,
        "create_dept_id": "",
        "create_time": create_time,
        "update_by": None,
        "update_time": create_time,
        "del_flag": 0,
    }


INSERT_SQL = {
    "purchase_plans": """
        INSERT INTO db_purchase_plan (
            id, plan_no, plan_date, start_date, end_date, business_type, product, specification,
            planned_quantity, vehicle_no, customer_name, execution_status, executed_at,
            last_weighing_slip_id, executed_quantity, auto_generate_enabled,
            auto_generate_end_month, auto_generate_last_month,
            tenant_no, order_num, remark, create_by, create_dept_id, create_time,
            update_by, update_time, del_flag
        ) VALUES (
            :id, :plan_no, :plan_date, :start_date, :end_date, :business_type, :product, :specification,
            :planned_quantity, :vehicle_no, :customer_name, :execution_status, :executed_at,
            :last_weighing_slip_id, :executed_quantity, :auto_generate_enabled,
            :auto_generate_end_month, :auto_generate_last_month,
            :tenant_no, :order_num, :remark, :create_by, :create_dept_id, :create_time,
            :update_by, :update_time, :del_flag
        )
    """,
    "inbound_outbound": """
        INSERT INTO db_inbound_outbound (
            id, date, partner_company, business_type, product, quantity, plan_no, batch_no,
            vehicle_no, status, purchase_plan_id, tenant_no, order_num, remark,
            create_by, create_dept_id, create_time, update_by, update_time, del_flag
        ) VALUES (
            :id, :date, :partner_company, :business_type, :product, :quantity, :plan_no, :batch_no,
            :vehicle_no, :status, :purchase_plan_id, :tenant_no, :order_num, :remark,
            :create_by, :create_dept_id, :create_time, :update_by, :update_time, :del_flag
        )
    """,
    "weighing_slips": """
        INSERT INTO db_weighing_slip (
            id, vehicle_no, plan_no, cargo_name, business_type, batch_no, supplier, consignee,
            gross_weight, gross_weight_time, tare_weight, tare_weight_time, deduction, net_weight,
            planned_quantity, executed_quantity, purchase_plan_id, inbound_outbound_id,
            weighing_point, current_quantity, is_voided, tenant_no, order_num, remark,
            create_by, create_dept_id, create_time, update_by, update_time, del_flag
        ) VALUES (
            :id, :vehicle_no, :plan_no, :cargo_name, :business_type, :batch_no, :supplier, :consignee,
            :gross_weight, :gross_weight_time, :tare_weight, :tare_weight_time, :deduction, :net_weight,
            :planned_quantity, :executed_quantity, :purchase_plan_id, :inbound_outbound_id,
            :weighing_point, :current_quantity, :is_voided, :tenant_no, :order_num, :remark,
            :create_by, :create_dept_id, :create_time, :update_by, :update_time, :del_flag
        )
    """,
}

UPDATE_PLAN_LAST_WEIGHING_SQL = """
    UPDATE db_purchase_plan
    SET last_weighing_slip_id = :last_weighing_slip_id, update_time = :update_time
    WHERE id = :id
"""


def truncate_target_tables(conn, tenant_no: str) -> None:
    for target in TARGET_TABLES.values():
        conn.execute(
            text(f"DELETE FROM `{target}` WHERE tenant_no = :tenant_no"),
            {"tenant_no": tenant_no},
        )


def export_mapping(id_maps: Dict[str, Dict[int, str]], path: Path) -> None:
    payload = {
        table_key: {str(old_id): new_id for old_id, new_id in mapping.items()}
        for table_key, mapping in id_maps.items()
    }
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")


def print_report(stats: MigrationStats, dry_run: bool) -> None:
    mode = "DRY-RUN" if dry_run else "EXECUTED"
    print(f"\n=== 地磅历史数据迁移 ({mode}) ===")
    print("\n源表行数:")
    for key, count in stats.source_counts.items():
        print(f"  {SOURCE_TABLES[key]}: {count}")

    if not dry_run:
        print("\n已写入目标表:")
        for key, count in stats.inserted.items():
            print(f"  {TARGET_TABLES[key]}: {count}")

    if stats.warnings:
        print(f"\n警告 ({len(stats.warnings)}):")
        for item in stats.warnings[:20]:
            print(f"  - {item}")
        if len(stats.warnings) > 20:
            print(f"  ... 另有 {len(stats.warnings) - 20} 条")

    if stats.orphan_fks:
        print(f"\n关联 ID 缺失 ({len(stats.orphan_fks)}):")
        for item in stats.orphan_fks[:20]:
            print(f"  - {item}")
        if len(stats.orphan_fks) > 20:
            print(f"  ... 另有 {len(stats.orphan_fks) - 20} 条")
    else:
        print("\n关联 ID 校验: 通过（无 orphan FK）")


def migrate(
    source_engine: Engine,
    target_engine: Engine,
    *,
    tenant_no: str,
    dry_run: bool,
    truncate_target: bool,
    include_deleted: bool,
    mapping_out: Optional[Path],
) -> MigrationStats:
    stats = MigrationStats()
    rows_by_table: Dict[str, List[dict]] = {}

    for key, source_table in SOURCE_TABLES.items():
        rows = fetch_rows(source_engine, source_table)
        if not include_deleted:
            rows = [row for row in rows if int(row.get("is_deleted") or 0) == 0]
        rows_by_table[key] = rows
        stats.source_counts[key] = len(rows)

    id_maps = build_id_maps(rows_by_table)

    transformed: Dict[str, List[dict]] = {
        "purchase_plans": [
            transform_plan(
                row,
                id_maps["purchase_plans"][int(row["id"])],
                tenant_no,
                stats,
                last_weighing_slip_id=None,
            )
            for row in rows_by_table["purchase_plans"]
        ],
        "inbound_outbound": [
            transform_inventory(
                row,
                id_maps["inbound_outbound"][int(row["id"])],
                tenant_no,
                id_maps,
                stats,
            )
            for row in rows_by_table["inbound_outbound"]
        ],
        "weighing_slips": [
            transform_weighing(
                row,
                id_maps["weighing_slips"][int(row["id"])],
                tenant_no,
                id_maps,
                stats,
            )
            for row in rows_by_table["weighing_slips"]
        ],
    }

    plan_last_weighing_updates: List[Tuple[str, Optional[str], datetime]] = []
    for row in rows_by_table["purchase_plans"]:
        legacy_last = row.get("last_weighing_slip_id")
        if legacy_last is None:
            continue
        plan_new_id = id_maps["purchase_plans"][int(row["id"])]
        weighing_new_id = map_id(
            "weighing_slips",
            legacy_last,
            id_maps,
            stats,
            f"plan#{row.get('id')} last_weighing_slip_id",
        )
        update_time = parse_datetime(row.get("executed_at")) or parse_datetime(row.get("created_at")) or datetime.now()
        plan_last_weighing_updates.append((plan_new_id, weighing_new_id, update_time))

    if dry_run:
        for key, rows in transformed.items():
            stats.inserted[key] = len(rows)
        print_report(stats, dry_run=True)
        return stats

    with target_engine.begin() as conn:
        if truncate_target:
            truncate_target_tables(conn, tenant_no)

        for key in ("purchase_plans", "inbound_outbound", "weighing_slips"):
            for payload in transformed[key]:
                conn.execute(text(INSERT_SQL[key]), payload)
            stats.inserted[key] = len(transformed[key])

        for plan_id, weighing_id, update_time in plan_last_weighing_updates:
            conn.execute(
                text(UPDATE_PLAN_LAST_WEIGHING_SQL),
                {
                    "id": plan_id,
                    "last_weighing_slip_id": weighing_id,
                    "update_time": update_time,
                },
            )

    if mapping_out:
        export_mapping(id_maps, mapping_out)
        print(f"\nID 映射已写入: {mapping_out}")

    print_report(stats, dry_run=False)
    return stats


def main() -> int:
    parser = argparse.ArgumentParser(description="Migrate legacy dibang data into db_* tables")
    parser.add_argument(
        "--source-url",
        default="",
        help="Source MySQL URL (legacy dibang DB). Falls back to DIBANG_SOURCE_DATABASE_URL",
    )
    parser.add_argument(
        "--target-url",
        default="",
        help="Target MySQL URL (portal DB). Defaults to app DATABASE_URL",
    )
    parser.add_argument(
        "--tenant-no",
        default="shanxia2.0",
        help="Tenant number written to db_* rows",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate and print counts without writing",
    )
    parser.add_argument(
        "--truncate-target",
        action="store_true",
        help="Delete existing db_* rows for tenant before import (use with care)",
    )
    parser.add_argument(
        "--include-deleted",
        action="store_true",
        help="Include source rows with is_deleted=1 (maps to del_flag=1)",
    )
    parser.add_argument(
        "--mapping-out",
        default="",
        help="Write legacy_id→new_id mapping JSON to this path (identity when IDs preserved)",
    )
    args = parser.parse_args()

    source_url = args.source_url or os.environ.get("DIBANG_SOURCE_DATABASE_URL", "")
    target_url = args.target_url or settings.DATABASE_URL

    if not source_url:
        print("错误: 请通过 --source-url 或环境变量 DIBANG_SOURCE_DATABASE_URL 指定源地磅库连接", file=sys.stderr)
        return 2

    if args.truncate_target and args.dry_run:
        print("错误: --truncate-target 不能与 --dry-run 同时使用", file=sys.stderr)
        return 2

    source_engine = create_engine(source_url, pool_pre_ping=True)
    target_engine = create_engine(target_url, pool_pre_ping=True)

    mapping_out = Path(args.mapping_out) if args.mapping_out else None

    try:
        migrate(
            source_engine,
            target_engine,
            tenant_no=args.tenant_no,
            dry_run=args.dry_run,
            truncate_target=args.truncate_target,
            include_deleted=args.include_deleted,
            mapping_out=mapping_out,
        )
    except Exception as exc:
        print(f"迁移失败: {exc}", file=sys.stderr)
        raise

    return 0


if __name__ == "__main__":
    raise SystemExit(main())