inventory_analytics/prefect_inventory_scheduler.py

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()