shanxia2.0/backend/scripts/backfill_name_pinyin_abbr.py

#!/usr/bin/env python3
"""回填 crm_product / crm_customer 的 name_pinyin_abbr（仅空字段）。

用法：
  cd backend
  venv/bin/python scripts/backfill_name_pinyin_abbr.py --dry-run
  venv/bin/python scripts/backfill_name_pinyin_abbr.py
  venv/bin/python scripts/backfill_name_pinyin_abbr.py --env-file .env.production
"""
from __future__ import annotations

import argparse
import os
import sys
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))


def load_settings_env(env_file: Optional[str]) -> None:
    if env_file:
        os.environ["ENV_FILE"] = env_file
    from dotenv import load_dotenv

    candidates = []
    if env_file:
        p = Path(env_file)
        candidates.append(p if p.is_absolute() else BACKEND_ROOT / p)
    candidates.extend(
        [
            BACKEND_ROOT / ".env.local",
            BACKEND_ROOT / ".env",
            BACKEND_ROOT / ".env.production",
        ]
    )
    for path in candidates:
        if path.is_file():
            load_dotenv(path, override=True)
            break


def _empty_pinyin_filter(model):
    from sqlalchemy import or_

    return or_(model.name_pinyin_abbr.is_(None), model.name_pinyin_abbr == "")


def backfill_model(db, model, name_attr: str, *, dry_run: bool, batch_size: int) -> int:
    from app.utils.pinyin_abbr import build_name_pinyin_abbr

    q = db.query(model).filter(model.del_flag == 0).filter(_empty_pinyin_filter(model))
    total = q.count()
    print(f"{model.__tablename__}: empty={total}")
    if total == 0:
        return 0

    updated = 0
    dry_printed = 0
    while True:
        rows = q.limit(batch_size).all()
        if not rows:
            break
        for row in rows:
            try:
                name = getattr(row, name_attr) or ""
                abbr = build_name_pinyin_abbr(name)
                if not abbr:
                    continue
                if dry_run:
                    if dry_printed < 20:
                        print(f"  would set {row.id}: {abbr}")
                        dry_printed += 1
                    updated += 1
                else:
                    row.name_pinyin_abbr = abbr
                    updated += 1
            except Exception as exc:  # noqa: BLE001 — batch continue
                print(f"  skip {getattr(row, 'id', '?')}: {exc}")
        if dry_run:
            if updated > dry_printed:
                print(f"  ... and {updated - dry_printed} more")
            break
        db.commit()
        # filled rows leave the empty filter; no offset needed
    print(f"{model.__tablename__}: updated={updated}")
    return updated


def main() -> int:
    parser = argparse.ArgumentParser(description="Backfill CRM name_pinyin_abbr (empty only)")
    parser.add_argument("--dry-run", action="store_true", help="只统计/预览，不写库")
    parser.add_argument("--env-file", default=None, help="指定 dotenv 文件")
    parser.add_argument("--batch-size", type=int, default=200)
    args = parser.parse_args()

    load_settings_env(args.env_file)

    from app.core.database import SessionLocal
    from app.models.crm_models import CrmCustomer, CrmProduct

    db = SessionLocal()
    try:
        backfill_model(
            db, CrmProduct, "product_name", dry_run=args.dry_run, batch_size=args.batch_size
        )
        backfill_model(
            db, CrmCustomer, "customer_name", dry_run=args.dry_run, batch_size=args.batch_size
        )
    finally:
        db.close()
    return 0


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