优化迁移脚本速度

scripts/migration/migrate_runtime.py

@@ -5,17 +5,19 @@
 
 工作流:
     1. 读 mapping.yaml + cards.csv
-    2. 连奇成查:卡元数据(card_meta) + 累计流量(card_usage) + 代理佣金/主钱包余额
+    2. 连奇成查:累计流量(card_usage) + 套餐生命周期 + 代理佣金/主钱包余额
     3. 卡 → 当前生效套餐 → mapping.packages → target_package_id
     4. 生成 step2_*.sql
 
 用法:
     python3 migrate_runtime.py [--config-dir config] [--resources-dir resources] [--output-dir output]
+                               [--batch-size 1000]
 """
 from __future__ import annotations
 
 import argparse
 import sys
+import time
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
@@ -28,6 +30,7 @@ def main() -> int:
     parser.add_argument("--config-dir", default="config", help="yaml 配置目录")
     parser.add_argument("--resources-dir", default="resources", help="CSV 资源目录")
     parser.add_argument("--output-dir", default="output", help="SQL/CSV 输出目录")
+    parser.add_argument("--batch-size", type=int, default=legacy_query.BATCH_SIZE, help="老库批量查询大小")
     args = parser.parse_args()
 
     base = Path(__file__).resolve().parent
@@ -35,9 +38,13 @@ def main() -> int:
     resources_dir = (base / args.resources_dir).resolve()
     output_dir = (base / args.output_dir).resolve()
     output_dir.mkdir(parents=True, exist_ok=True)
+    batch_size = args.batch_size
 
+    t0 = time.perf_counter()
     mapping = mapping_loader.load_mapping(config_dir)
     cards, devices, errors = csv_loader.load_assets(resources_dir, mapping)
+    print(f"[migrate] 加载 CSV: cards={len(cards)} devices={len(devices)} elapsed={time.perf_counter()-t0:.1f}s")
+
     if errors:
         sql_builder.write_runtime_input_errors(output_dir, errors)
         print(f"cards.csv/devices.csv 存在 {len(errors)} 个基础错误,已写入 {output_dir / 'errors.csv'}", file=sys.stderr)
@@ -46,27 +53,39 @@ def main() -> int:
     iccid_pairs = [(c.iccid_19, c.iccid_20, c.allow_iccid_19_lookup) for c in cards]
     agent_ids = [a.legacy_agent_id for a in mapping.agents.values() if (a.target_shop_code or "").strip()]
 
+    t1 = time.perf_counter()
     dsn = mapping_loader.load_legacy_dsn(config_dir)
     with legacy_query.connect_readonly(dsn) as conn:
-        card_metas, _dup_iccids = legacy_query.fetch_card_meta(conn, iccid_pairs)
-        usage_snapshots = legacy_query.fetch_card_usage(conn, iccid_pairs)
-        package_lifecycles = legacy_query.fetch_package_lifecycles(conn, iccid_pairs)
-        commission_snapshots = legacy_query.fetch_commission_accounts(conn, agent_ids) if agent_ids else {}
-        agent_balances = legacy_query.fetch_agent_balances(conn, agent_ids) if agent_ids else {}
+        print(f"[migrate] 连接老库: elapsed={time.perf_counter()-t1:.1f}s")
 
+        t2 = time.perf_counter()
+        usage_snapshots = legacy_query.fetch_card_usage(conn, iccid_pairs, batch_size=batch_size)
+        print(f"[migrate] fetch_card_usage: cards={len(usage_snapshots)} elapsed={time.perf_counter()-t2:.1f}s")
+
+        t3 = time.perf_counter()
+        package_lifecycles = legacy_query.fetch_package_lifecycles(conn, iccid_pairs, batch_size=batch_size)
+        print(f"[migrate] fetch_package_lifecycles: cards={len(package_lifecycles)} elapsed={time.perf_counter()-t3:.1f}s")
+
+        t4 = time.perf_counter()
+        commission_snapshots = legacy_query.fetch_commission_accounts(conn, agent_ids, batch_size=batch_size) if agent_ids else {}
+        agent_balances = legacy_query.fetch_agent_balances(conn, agent_ids, batch_size=batch_size) if agent_ids else {}
+        print(f"[migrate] fetch_agent_data: agents={len(agent_ids)} elapsed={time.perf_counter()-t4:.1f}s")
+
+    t5 = time.perf_counter()
     warnings = sql_builder.write_step2(
         output_dir,
         cards=cards,
         devices=devices,
         mapping=mapping,
-        card_metas=card_metas,
         usage_snapshots=usage_snapshots,
         package_lifecycles=package_lifecycles,
         commission_snapshots=commission_snapshots,
         agent_balances=agent_balances,
     )
+    print(f"[migrate] SQL 生成: elapsed={time.perf_counter()-t5:.1f}s")
 
-    print(f"完成:卡 {len(cards)} 张,代理 {len(agent_ids)} 个,警告 {len(warnings)} 条")
+    total = time.perf_counter() - t0
+    print(f"完成:卡 {len(cards)} 张,代理 {len(agent_ids)} 个,警告 {len(warnings)} 条,总耗时 {total:.1f}s")
     print(f"SQL 输出: {output_dir}")
     return 0