zkxw/backend/scripts/backfill_party_customer_ids.py

#!/usr/bin/env python3
"""回填当事人父亲/子女/缺失母亲客户 ID。可重复执行。

用法（在 backend 目录、已配置 .env）:
  python -m scripts.backfill_party_customer_ids
  python -m scripts.backfill_party_customer_ids --dry-run
"""
from __future__ import annotations

import argparse
import os
import sys
from datetime import datetime
from typing import Any, Dict, Optional, Tuple

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

from app.core.database import SessionLocal
from app.crud import collection_protocol as cp_svc
from app.models.models import BizProtocolPartyCustomer
from app.models.research_donation import BizResearchDonationParty
from app.services.party_customer_sync import sync_party_customers


def _party_to_sync_dict(party: Any) -> Dict[str, Any]:
    return {
        "customer_id": (party.customer_id or "").strip(),
        "father_customer_id": (getattr(party, "father_customer_id", None) or "").strip(),
        "newborn_customer_id": (getattr(party, "newborn_customer_id", None) or "").strip(),
        "mother_name": (party.mother_name or "").strip(),
        "mother_phone": (party.mother_phone or "").strip(),
        "mother_id_card": (party.mother_id_card or "").strip(),
        "father_name": (party.father_name or "").strip(),
        "father_phone": (party.father_phone or "").strip(),
        "father_id_card": (party.father_id_card or "").strip(),
        "newborn_name": (party.newborn_name or "").strip(),
        "newborn_id_card": (party.newborn_id_card or "").strip(),
        "newborn_sex": (party.newborn_sex or "").strip(),
    }


def _ids(party: Any) -> Tuple[str, str, str]:
    return (
        (party.customer_id or "").strip(),
        (getattr(party, "father_customer_id", None) or "").strip(),
        (getattr(party, "newborn_customer_id", None) or "").strip(),
    )


def _needs_sync(party: Any) -> bool:
    mother_name = (party.mother_name or "").strip()
    father_name = (party.father_name or "").strip()
    newborn_name = (party.newborn_name or "").strip()
    cid, fid, nid = _ids(party)
    if mother_name and not cid:
        return True
    if father_name and not fid:
        return True
    if newborn_name and not nid:
        return True
    # 有姓名也允许再跑 sync，以便补证撞库；但默认只回填缺 ID 的，保持幂等少噪音
    return False


def _apply_ids(party: Any, synced: Dict[str, Any], user_id: Optional[str]) -> bool:
    before = _ids(party)
    after = (
        (synced.get("customer_id") or "").strip(),
        (synced.get("father_customer_id") or "").strip(),
        (synced.get("newborn_customer_id") or "").strip(),
    )
    if before == after:
        return False
    party.customer_id = after[0] or None
    party.father_customer_id = after[1] or None
    party.newborn_customer_id = after[2] or None
    party.update_by = user_id
    party.update_time = datetime.now()
    return True


def backfill_storage(db, *, dry_run: bool, user_id: str) -> Tuple[int, int]:
    updated = 0
    skipped = 0
    rows = (
        db.query(BizProtocolPartyCustomer)
        .filter(BizProtocolPartyCustomer.del_flag == 0)
        .order_by(BizProtocolPartyCustomer.create_time.asc())
        .all()
    )
    for party in rows:
        if not _needs_sync(party):
            skipped += 1
            continue
        synced = sync_party_customers(
            db,
            _party_to_sync_dict(party),
            tenant_no=party.tenant_no,
            user_id=user_id,
        )
        before = _ids(party)
        after = (
            (synced.get("customer_id") or "").strip(),
            (synced.get("father_customer_id") or "").strip(),
            (synced.get("newborn_customer_id") or "").strip(),
        )
        if before == after:
            skipped += 1
            continue
        if dry_run:
            updated += 1
            print(f"[dry-run] storage party {party.id} {before} -> {after}")
            continue
        cp_svc.upsert_party_customer(
            db,
            party.collection_protocol_id,
            party.tenant_no,
            user_id,
            **synced,
        )
        updated += 1
    return updated, skipped


def backfill_donation(db, *, dry_run: bool, user_id: str) -> Tuple[int, int]:
    updated = 0
    skipped = 0
    rows = (
        db.query(BizResearchDonationParty)
        .filter(BizResearchDonationParty.del_flag == 0)
        .order_by(BizResearchDonationParty.create_time.asc())
        .all()
    )
    for party in rows:
        if not _needs_sync(party):
            skipped += 1
            continue
        synced = sync_party_customers(
            db,
            _party_to_sync_dict(party),
            tenant_no=party.tenant_no,
            user_id=user_id,
        )
        if dry_run:
            before = _ids(party)
            after = (
                (synced.get("customer_id") or "").strip(),
                (synced.get("father_customer_id") or "").strip(),
                (synced.get("newborn_customer_id") or "").strip(),
            )
            if before != after:
                updated += 1
                print(
                    f"[dry-run] donation party {party.id} "
                    f"{before} -> {after}"
                )
            else:
                skipped += 1
            continue
        if _apply_ids(party, synced, user_id):
            db.add(party)
            db.commit()
            updated += 1
        else:
            skipped += 1
    return updated, skipped


def main() -> None:
    parser = argparse.ArgumentParser(description="回填当事人三角色客户 ID")
    parser.add_argument("--dry-run", action="store_true", help="只打印变更，不写库")
    parser.add_argument("--user-id", default="backfill", help="写入 update_by / create_by")
    args = parser.parse_args()

    db = SessionLocal()
    try:
        s_upd, s_skip = backfill_storage(db, dry_run=args.dry_run, user_id=args.user_id)
        d_upd, d_skip = backfill_donation(db, dry_run=args.dry_run, user_id=args.user_id)
        print(
            f"storage: updated={s_upd} skipped={s_skip}; "
            f"donation: updated={d_upd} skipped={d_skip}"
        )
        if args.dry_run:
            print("dry-run, no commit")
        else:
            print("done")
    finally:
        db.close()


if __name__ == "__main__":
    main()