182 lines
7.2 KiB
Python
182 lines
7.2 KiB
Python
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.batch_inventory_report.cli import main as batch_inventory_report_main
|
|
|
|
|
|
@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")
|
|
root_logger = logging.getLogger()
|
|
original_handlers = root_logger.handlers[:]
|
|
original_level = root_logger.level
|
|
try:
|
|
os.environ["INVENTORY_SAFE_LOGS"] = "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
|
|
root_logger.handlers = original_handlers
|
|
root_logger.setLevel(original_level)
|
|
if exit_code != 0:
|
|
logger.error("渠道库存统计失败，退出码=%s", exit_code)
|
|
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")
|
|
root_logger = logging.getLogger()
|
|
original_handlers = root_logger.handlers[:]
|
|
original_level = root_logger.level
|
|
try:
|
|
os.environ["INVENTORY_SAFE_LOGS"] = "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
|
|
root_logger.handlers = original_handlers
|
|
root_logger.setLevel(original_level)
|
|
if exit_code != 0:
|
|
logger.error("批次库存统计失败，退出码=%s", exit_code)
|
|
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 * * *", "15 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 * * *", "15 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()
|