shanxia2.0/backend/scripts/backfill_inventory_site_tokens.py

#!/usr/bin/env python3
"""为 site_access_token 为空的出入库单补 token。

Usage:
  cd backend && source venv/bin/activate
  python scripts/backfill_inventory_site_tokens.py
  python scripts/backfill_inventory_site_tokens.py --dry-run
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

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.models.db_models import DbInboundOutbound
from app.services.db_inventory_site_token import ensure_site_access_token


def backfill(*, dry_run: bool = False) -> int:
    db: Session = SessionLocal()
    updated = 0
    try:
        rows = (
            db.query(DbInboundOutbound)
            .filter(
                DbInboundOutbound.del_flag == 0,
                (DbInboundOutbound.site_access_token.is_(None))
                | (DbInboundOutbound.site_access_token == ""),
            )
            .all()
        )
        for row in rows:
            ensure_site_access_token(row)
            updated += 1
        if dry_run:
            db.rollback()
            print(f"[dry-run] would update {updated} inventory rows")
        else:
            db.commit()
            print(f"updated {updated} inventory rows")
        return updated
    finally:
        db.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="Backfill db_inbound_outbound site_access_token")
    parser.add_argument("--dry-run", action="store_true", help="Preview without committing")
    args = parser.parse_args()
    backfill(dry_run=args.dry_run)


if __name__ == "__main__":
    main()