from __future__ import annotations import argparse import logging import os import sys from pathlib import Path import time 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, get_run_logger, task from inventory_analytics.logging_utils import exception_fields from inventory_analytics.channel_inventory_sync.entrypoint import main as channel_inventory_sync_main from inventory_analytics.channel_inventory_sync.entrypoint import LOG_FILE as CHANNEL_INVENTORY_SYNC_LOG_FILE from inventory_analytics.batch_inventory_report.cli import main as batch_inventory_report_main from inventory_analytics.batch_inventory_report.service import LOG_FILE as BATCH_INVENTORY_REPORT_LOG_FILE def read_log_tail(log_file: Path, line_count: int = 80) -> str: try: lines = log_file.read_text(encoding="utf-8", errors="replace").splitlines() except FileNotFoundError: return f"日志文件不存在: {log_file.name}" except Exception as exc: # noqa: BLE001 return f"读取日志文件失败: {exc.__class__.__name__}: {exc}" return "\n".join(lines[-line_count:]) if lines else f"日志文件为空: {log_file.name}" @task(name="run-channel-inventory-sync") def run_channel_inventory_sync_task() -> None: logger = get_run_logger() started_at = time.monotonic() logger.info("渠道库存统计开始，action=all，安全日志模式=enabled") original_argv = sys.argv[:] original_safe_logs = os.environ.get("INVENTORY_SAFE_LOGS") original_prefect_log_bridge = os.environ.get("INVENTORY_PREFECT_LOG_BRIDGE") root_logger = logging.getLogger() original_handlers = root_logger.handlers[:] original_level = root_logger.level try: os.environ["INVENTORY_SAFE_LOGS"] = "1" os.environ["INVENTORY_PREFECT_LOG_BRIDGE"] = "1" sys.argv = ["sync_channel_inventory.py", "--action", "all"] exit_code = channel_inventory_sync_main() except Exception as exc: # noqa: BLE001 logger.error("渠道库存统计异常，%s", exception_fields(exc)) raise finally: sys.argv = original_argv if original_safe_logs is None: os.environ.pop("INVENTORY_SAFE_LOGS", None) else: os.environ["INVENTORY_SAFE_LOGS"] = original_safe_logs if original_prefect_log_bridge is None: os.environ.pop("INVENTORY_PREFECT_LOG_BRIDGE", None) else: os.environ["INVENTORY_PREFECT_LOG_BRIDGE"] = original_prefect_log_bridge root_logger.handlers = original_handlers root_logger.setLevel(original_level) if exit_code != 0: logger.error("渠道库存统计失败，退出码=%s", exit_code) logger.error("渠道库存统计日志尾部:\n%s", read_log_tail(CHANNEL_INVENTORY_SYNC_LOG_FILE)) raise RuntimeError(f"渠道库存统计执行失败，退出码: {exit_code}") logger.info("渠道库存统计完成，耗时毫秒=%s", int((time.monotonic() - started_at) * 1000)) @task(name="run-batch-inventory-report") def run_batch_inventory_report_task( output: str = "", feishu_table: str = "批次库存", skip_feishu_sync: bool = False, skip_report_output: bool = True, ) -> None: logger = get_run_logger() started_at = time.monotonic() logger.info( "批次库存统计开始，飞书同步=%s，本地报表输出=%s，安全日志模式=enabled", "disabled" if skip_feishu_sync else "enabled", "skipped" if skip_report_output else "enabled", ) original_argv = sys.argv[:] original_safe_logs = os.environ.get("INVENTORY_SAFE_LOGS") original_prefect_log_bridge = os.environ.get("INVENTORY_PREFECT_LOG_BRIDGE") root_logger = logging.getLogger() original_handlers = root_logger.handlers[:] original_level = root_logger.level try: os.environ["INVENTORY_SAFE_LOGS"] = "1" os.environ["INVENTORY_PREFECT_LOG_BRIDGE"] = "1" 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") if skip_report_output: sys.argv.append("--skip-report-output") sys.argv.append("--safe-logs") exit_code = batch_inventory_report_main() except Exception as exc: # noqa: BLE001 logger.error("批次库存统计异常，%s", exception_fields(exc)) raise finally: sys.argv = original_argv if original_safe_logs is None: os.environ.pop("INVENTORY_SAFE_LOGS", None) else: os.environ["INVENTORY_SAFE_LOGS"] = original_safe_logs if original_prefect_log_bridge is None: os.environ.pop("INVENTORY_PREFECT_LOG_BRIDGE", None) else: os.environ["INVENTORY_PREFECT_LOG_BRIDGE"] = original_prefect_log_bridge root_logger.handlers = original_handlers root_logger.setLevel(original_level) if exit_code != 0: logger.error("批次库存统计失败，退出码=%s", exit_code) logger.error("批次库存统计日志尾部:\n%s", read_log_tail(BATCH_INVENTORY_REPORT_LOG_FILE)) raise RuntimeError(f"批次库存统计报表执行失败，退出码: {exit_code}") logger.info("批次库存统计完成，耗时毫秒=%s", int((time.monotonic() - started_at) * 1000)) @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, skip_report_output: bool = True, ) -> None: run_batch_inventory_report_task( output=output, feishu_table=feishu_table, skip_feishu_sync=skip_feishu_sync, skip_report_output=skip_report_output, ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Prefect 调度：每天 09:15 和 19: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="批次库存统计只生成本地文件，不写飞书") parser.add_argument("--write-report-output", action="store_true", help="Prefect 执行批次库存统计时同时生成本地报表文件") 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, "skip_report_output": not args.write_report_output, } if args.serve_channel_inventory_sync: inventory_channel_inventory_sync_flow.serve( name="channel-inventory-sync-0915-1915", cron="15 9,19 * * *", tags=["inventory", "feishu", "channel-inventory-sync"], ) return if args.serve_batch_inventory_report: inventory_batch_inventory_report_flow.serve( name="batch-inventory-report-0915-1915", cron="15 9,19 * * *", 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()