shanxia2.0/backend/scripts/clear_pl_plan_instance_data.py

"""
清空运输计划（pl_plan）及关联实例数据，便于测试环境重新走流程。

用法（在 backend 目录、已配置 .env DATABASE_URL）：
  python scripts/clear_pl_plan_instance_data.py --yes
"""
from __future__ import annotations

import argparse
import os
import sys
from urllib.parse import urlparse

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

# 确保可加载 app 配置
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

TENANT_NO = "shanxia2.0"


def _engine_from_env() -> Engine:
    url = os.environ.get("DATABASE_URL")
    if not url:
        from app.core.config import settings

        url = settings.DATABASE_URL
    return create_engine(url, pool_pre_ping=True)


def _table_exists(conn, table: str) -> bool:
    row = conn.execute(
        text(
            "SELECT COUNT(*) FROM information_schema.tables "
            "WHERE table_schema = DATABASE() AND table_name = :t"
        ),
        {"t": table},
    ).scalar()
    return int(row or 0) > 0


def _count(conn, table: str, where: str = "", params: dict | None = None) -> int:
    if not _table_exists(conn, table):
        return 0
    sql = f"SELECT COUNT(*) FROM `{table}`"
    if where:
        sql += f" WHERE {where}"
    return int(conn.execute(text(sql), params or {}).scalar() or 0)


def _delete(conn, table: str, where: str, params: dict | None = None) -> int:
    if not _table_exists(conn, table):
        return 0
    sql = f"DELETE FROM `{table}` WHERE {where}"
    result = conn.execute(text(sql), params or {})
    return int(result.rowcount or 0)


def main() -> None:
    parser = argparse.ArgumentParser(description="清空运输计划实例数据")
    parser.add_argument(
        "--yes",
        action="store_true",
        help="确认执行删除（未带此参数仅预览计数）",
    )
    parser.add_argument("--tenant", default=TENANT_NO, help="租户编号，默认 shanxia2.0")
    args = parser.parse_args()

    engine = _engine_from_env()
    parsed = urlparse(str(engine.url))
    db_name = parsed.path.lstrip("/")
    print(f"目标库: {db_name} @ {parsed.hostname}:{parsed.port or 3306}")
    print(f"租户: {args.tenant}")
    print("---")

    with engine.connect() as conn:
        plan_count = _count(conn, "pl_plan", "tenant_no = :t", {"t": args.tenant})
        print(f"pl_plan 行数: {plan_count}")

        if not args.yes:
            print("预览模式，未删除。请加 --yes 执行。")
            return

    stats: dict[str, int] = {}

    with engine.begin() as conn:
        tenant = args.tenant

        def del_by_plan(table: str, extra: str = "") -> None:
            cond = f"plan_id IN (SELECT id FROM pl_plan WHERE tenant_no = :t)"
            if extra:
                cond = f"{cond} AND {extra}"
            n = _delete(conn, table, cond, {"t": tenant})
            if n:
                stats[table] = stats.get(table, 0) + n

        # 先收集运输单 id（删 pl_plan_transport 后无法再关联）
        transport_ids: list[str] = []
        if _table_exists(conn, "pl_plan_transport"):
            rows = conn.execute(
                text(
                    "SELECT DISTINCT transport_id FROM pl_plan_transport "
                    "WHERE tenant_no = :t AND plan_id IN (SELECT id FROM pl_plan WHERE tenant_no = :t)"
                ),
                {"t": tenant},
            ).fetchall()
            transport_ids = [str(r[0]) for r in rows if r[0]]

        def del_by_transport_ids(table: str, col: str = "transport_id") -> None:
            if not transport_ids:
                return
            placeholders = ",".join([f":tid{i}" for i in range(len(transport_ids))])
            params = {f"tid{i}": v for i, v in enumerate(transport_ids)}
            n = _delete(conn, table, f"{col} IN ({placeholders})", params)
            if n:
                stats[table] = stats.get(table, 0) + n

        def del_transports() -> None:
            if not transport_ids:
                return
            placeholders = ",".join([f":tid{i}" for i in range(len(transport_ids))])
            params = {f"tid{i}": v for i, v in enumerate(transport_ids)}
            n = _delete(conn, "pl_transport", f"id IN ({placeholders})", params)
            if n:
                stats["pl_transport"] = n

        def del_by_transport(table: str) -> None:
            del_by_transport_ids(table, "transport_id")

        def del_inbound_children() -> None:
            if not _table_exists(conn, "pl_inbound_order"):
                return
            cond = "order_id IN (SELECT id FROM pl_inbound_order WHERE plan_id IN (SELECT id FROM pl_plan WHERE tenant_no = :t))"
            n = _delete(
                conn,
                "pl_inbound_order_item_brand_qty",
                cond,
                {"t": tenant},
            )
            if n:
                stats["pl_inbound_order_item_brand_qty"] = n
            n = _delete(conn, "pl_inbound_order_item", cond, {"t": tenant})
            if n:
                stats["pl_inbound_order_item"] = n
            n = _delete(conn, "pl_inbound_order", "plan_id IN (SELECT id FROM pl_plan WHERE tenant_no = :t)", {"t": tenant})
            if n:
                stats["pl_inbound_order"] = n

        def del_outbound_children() -> None:
            if not _table_exists(conn, "pl_outbound_order"):
                return
            cond = "order_id IN (SELECT id FROM pl_outbound_order WHERE plan_id IN (SELECT id FROM pl_plan WHERE tenant_no = :t))"
            n = _delete(
                conn,
                "pl_outbound_order_item_brand_qty",
                cond,
                {"t": tenant},
            )
            if n:
                stats["pl_outbound_order_item_brand_qty"] = n
            n = _delete(conn, "pl_outbound_order_item", cond, {"t": tenant})
            if n:
                stats["pl_outbound_order_item"] = n
            n = _delete(conn, "pl_outbound_order", "plan_id IN (SELECT id FROM pl_plan WHERE tenant_no = :t)", {"t": tenant})
            if n:
                stats["pl_outbound_order"] = n

        # 地磅：先收集关联 id
        db_pp_ids: list[str] = []
        db_inv_ids: list[str] = []
        if _table_exists(conn, "pl_transport_db_link"):
            rows = conn.execute(
                text(
                    "SELECT DISTINCT db_purchase_plan_id, db_inventory_id FROM pl_transport_db_link "
                    "WHERE tenant_no = :t AND plan_id IN (SELECT id FROM pl_plan WHERE tenant_no = :t)"
                ),
                {"t": tenant},
            ).fetchall()
            for pp, inv in rows:
                if pp and str(pp).strip():
                    db_pp_ids.append(str(pp))
                if inv and str(inv).strip():
                    db_inv_ids.append(str(inv))

        # 运输单子表
        del_by_transport("pl_transport_product_brand_qty")
        del_by_transport("pl_transport_product_feedback")
        del_by_plan("pl_transport_db_link")
        del_by_plan("pl_transport_driver_stop_feedback")
        del_by_plan("pl_transport_site_stop_feedback")
        del_by_plan("pl_transport_site_link")
        del_inbound_children()
        del_outbound_children()
        del_by_transport("pl_transport_approval")
        del_by_transport("pl_freight_settlement_transport")
        del_by_plan("pl_plan_transport")
        del_transports()

        def del_orphan_transports() -> None:
            """计划关联已删后仍残留的运输单（租户级实例数据）"""
            if not _table_exists(conn, "pl_transport"):
                return
            n = _delete(conn, "pl_transport", "tenant_no = :t", {"t": tenant})
            if n:
                stats["pl_transport(orphan)"] = n

        def del_freight_settlements() -> None:
            if not _table_exists(conn, "pl_freight_settlement"):
                return
            sid_cond = "settlement_id IN (SELECT id FROM pl_freight_settlement WHERE tenant_no = :t)"
            n = _delete(conn, "pl_freight_settlement_approval", sid_cond, {"t": tenant})
            if n:
                stats["pl_freight_settlement_approval"] = n
            n = _delete(conn, "pl_freight_settlement_transport", sid_cond, {"t": tenant})
            if n:
                stats["pl_freight_settlement_transport(extra)"] = n
            n = _delete(conn, "pl_freight_settlement", "tenant_no = :t", {"t": tenant})
            if n:
                stats["pl_freight_settlement"] = n

        del_orphan_transports()
        del_freight_settlements()

        # 计划子表
        del_by_plan("pl_plan_load_product")
        del_by_plan("pl_plan_delivery_product")
        del_by_plan("pl_plan_load")
        del_by_plan("pl_plan_delivery")
        del_by_plan("pl_plan_contract")
        del_by_plan("pl_plan_product_brand_qty")
        del_by_plan("pl_plan_product")
        del_by_plan("pl_plan_change")

        # 地磅出入库与计划
        if _table_exists(conn, "db_inbound_outbound"):
            if db_inv_ids:
                placeholders = ",".join([":i" + str(i) for i in range(len(db_inv_ids))])
                params = {"t": tenant, **{f"i{i}": v for i, v in enumerate(db_inv_ids)}}
                n = conn.execute(
                    text(f"DELETE FROM db_inbound_outbound WHERE id IN ({placeholders})"),
                    params,
                ).rowcount
                if n:
                    stats["db_inbound_outbound(by_inv)"] = int(n)
            if db_pp_ids:
                placeholders = ",".join([":p" + str(i) for i in range(len(db_pp_ids))])
                params = {f"p{i}": v for i, v in enumerate(db_pp_ids)}
                n = conn.execute(
                    text(f"DELETE FROM db_inbound_outbound WHERE purchase_plan_id IN ({placeholders})"),
                    params,
                ).rowcount
                if n:
                    stats["db_inbound_outbound(by_pp)"] = int(n)
            n = _delete(
                conn,
                "db_inbound_outbound",
                "plan_no IN (SELECT plan_no FROM pl_plan WHERE tenant_no = :t)",
                {"t": tenant},
            )
            if n:
                stats["db_inbound_outbound(by_plan_no)"] = n

        if _table_exists(conn, "db_purchase_plan"):
            if db_pp_ids:
                placeholders = ",".join([":p" + str(i) for i in range(len(db_pp_ids))])
                params = {f"p{i}": v for i, v in enumerate(db_pp_ids)}
                n = conn.execute(
                    text(f"DELETE FROM db_purchase_plan WHERE id IN ({placeholders})"),
                    params,
                ).rowcount
                if n:
                    stats["db_purchase_plan(by_link)"] = int(n)
            n = _delete(
                conn,
                "db_purchase_plan",
                "tenant_no = :t AND remark LIKE '来源运输计划%'",
                {"t": tenant},
            )
            if n:
                stats["db_purchase_plan(by_remark)"] = n
            n = _delete(
                conn,
                "db_purchase_plan",
                "tenant_no = :t AND plan_no IN (SELECT plan_no FROM pl_plan WHERE tenant_no = :t)",
                {"t": tenant},
            )
            if n:
                stats["db_purchase_plan(by_plan_no)"] = n

        n = _delete(conn, "pl_plan", "tenant_no = :t", {"t": tenant})
        stats["pl_plan"] = n

        remaining = _count(conn, "pl_plan", "tenant_no = :t", {"t": tenant})

    print("已删除行数:")
    for k, v in sorted(stats.items()):
        print(f"  {k}: {v}")
    print(f"剩余 pl_plan: {remaining}")
    print("完成。")


if __name__ == "__main__":
    main()