重构

prefect_inventory_scheduler.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+
+def bootstrap_prefect_env() -> None:
+    env_file = Path(__file__).resolve().parent / ".env"
+    if not env_file.exists():
+        return
+    for raw_line in env_file.read_text(encoding="utf-8").splitlines():
+        line = raw_line.strip()
+        if not line or line.startswith("#") or "=" not in line:
+            continue
+        key, value = line.split("=", 1)
+        key = key.strip()
+        if key.startswith("PREFECT_"):
+            os.environ.setdefault(key, value.strip())
+
+
+bootstrap_prefect_env()
+
+from prefect import flow, task
+
+from inventory_analytics.channel_inventory_sync.entrypoint import main as channel_inventory_sync_main
+from inventory_analytics.batch_inventory_report.cli import main as batch_inventory_report_main
+
+
+@task(name="run-channel-inventory-sync")
+def run_channel_inventory_sync_task() -> None:
+    original_argv = sys.argv[:]
+    try:
+        sys.argv = ["sync_channel_inventory.py", "--action", "all"]
+        exit_code = channel_inventory_sync_main()
+    finally:
+        sys.argv = original_argv
+    if exit_code != 0:
+        raise RuntimeError(f"渠道库存统计执行失败，退出码: {exit_code}")
+
+
+@task(name="run-batch-inventory-report")
+def run_batch_inventory_report_task(output: str = "", feishu_table: str = "批次库存", skip_feishu_sync: bool = False) -> None:
+    original_argv = sys.argv[:]
+    try:
+        sys.argv = ["generate_batch_inventory_report.py"]
+        if output:
+            sys.argv.extend(["--output", output])
+        if feishu_table:
+            sys.argv.extend(["--feishu-table", feishu_table])
+        if skip_feishu_sync:
+            sys.argv.append("--skip-feishu-sync")
+        exit_code = batch_inventory_report_main()
+    finally:
+        sys.argv = original_argv
+    if exit_code != 0:
+        raise RuntimeError(f"批次库存统计报表执行失败，退出码: {exit_code}")
+
+
+@flow(name="inventory-channel-inventory-sync")
+def inventory_channel_inventory_sync_flow() -> None:
+    run_channel_inventory_sync_task()
+
+
+@flow(name="inventory-batch-inventory-report")
+def inventory_batch_inventory_report_flow(
+    output: str = "",
+    feishu_table: str = "批次库存",
+    skip_feishu_sync: bool = False,
+) -> None:
+    run_batch_inventory_report_task(output=output, feishu_table=feishu_table, skip_feishu_sync=skip_feishu_sync)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Prefect 调度：每天 09:15 执行渠道库存统计和批次库存统计写飞书。")
+    parser.add_argument("--serve-channel-inventory-sync", action="store_true", help="启动渠道库存统计 flow 的本地 serve 调度")
+    parser.add_argument("--serve-batch-inventory-report", action="store_true", help="启动批次库存统计 flow 的本地 serve 调度")
+    parser.add_argument("--run-channel-inventory-sync", action="store_true", help="立即执行一次渠道库存统计 flow")
+    parser.add_argument("--run-batch-inventory-report", action="store_true", help="立即执行一次批次库存统计 flow")
+    parser.add_argument("--output", default="", help="批次库存统计输出路径，可选")
+    parser.add_argument("--feishu-table", default="批次库存", help="批次库存统计同步目标飞书表名")
+    parser.add_argument("--skip-feishu-sync", action="store_true", help="批次库存统计只生成本地文件，不写飞书")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    flow_kwargs = {
+        "output": args.output,
+        "feishu_table": args.feishu_table,
+        "skip_feishu_sync": args.skip_feishu_sync,
+    }
+    if args.serve_channel_inventory_sync:
+        inventory_channel_inventory_sync_flow.serve(
+            name="inventory-channel-inventory-sync-daily-0915",
+            cron="15 9 * * *",
+            tags=["inventory", "feishu", "channel-inventory-sync"],
+        )
+        return
+    if args.serve_batch_inventory_report:
+        inventory_batch_inventory_report_flow.serve(
+            name="inventory-batch-inventory-report-daily-0915",
+            cron="15 9 * * *",
+            parameters=flow_kwargs,
+            tags=["inventory", "feishu", "batch-inventory-report"],
+        )
+        return
+    if args.run_channel_inventory_sync:
+        inventory_channel_inventory_sync_flow()
+        return
+    if args.run_batch_inventory_report:
+        inventory_batch_inventory_report_flow(**flow_kwargs)
+        return
+    inventory_channel_inventory_sync_flow()
+    inventory_batch_inventory_report_flow(**flow_kwargs)
+
+
+if __name__ == "__main__":
+    main()