gmv-analytics/prefect_gmv_scheduler.py

import argparse
import json
import os
from dataclasses import dataclass
from datetime import date, timedelta
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 project_bitable_common import (
    DailySyncConfig,
    build_daily_sync_context,
    get_base_url,
    get_default_end_date,
    get_default_start_date,
    get_feishu_chat_id,
    get_feishu_report_image_path,
    get_group_webhook_url,
    get_mapping_file,
    get_store_field_name,
    iter_dates,
    load_dotenv,
    log_step,
    run_daily_sync_for_date,
    send_missing_table_notification,
    send_sync_notification,
    validate_date,
)

BASE_DIR = Path(__file__).resolve().parent
DEFAULT_PREFECT_SYNC_STATE_FILE = BASE_DIR / "prefect_sync_state.json"


@dataclass(frozen=True)
class SyncConfig:
    start_date: str
    end_date: str
    base_url: str
    mapping_file: str
    store_field_name: str
    chat_id: str
    report_image_path: str
    group_webhook: str
    table_map_file: str
    dry_run: bool
    skip_missing_tables: bool


@dataclass(frozen=True)
class SyncWindow:
    start_date: date
    end_date: date


def load_sync_config(
    start_date: str | None = None,
    end_date: str | None = None,
    dry_run: bool = False,
    skip_missing_tables: bool = True,
    table_map_file: str = "",
) -> SyncConfig:
    load_dotenv()
    return SyncConfig(
        start_date=start_date or get_default_start_date(),
        end_date=end_date or get_default_end_date(),
        base_url=get_base_url(),
        mapping_file=str(get_mapping_file()),
        store_field_name=get_store_field_name(),
        chat_id=get_feishu_chat_id(),
        report_image_path=get_feishu_report_image_path(),
        group_webhook=get_group_webhook_url(),
        table_map_file=table_map_file,
        dry_run=dry_run,
        skip_missing_tables=skip_missing_tables,
    )


def get_prefect_sync_state_file() -> Path:
    return Path(os.getenv("PREFECT_SYNC_STATE_FILE", str(DEFAULT_PREFECT_SYNC_STATE_FILE)))


def load_prefect_sync_state() -> dict[str, str]:
    path = get_prefect_sync_state_file()
    if not path.exists():
        return {}
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except Exception as exc:
        log_step(f"读取 Prefect 同步状态失败，已忽略: path={path}, error={exc}")
        return {}
    return {str(key): str(value) for key, value in payload.items() if value is not None}


def save_prefect_sync_state(state: dict[str, str]) -> None:
    path = get_prefect_sync_state_file()
    path.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8")


def resolve_sync_window(
    config: SyncConfig,
    explicit_start_date: str | None,
    explicit_end_date: str | None,
    mode: str,
) -> SyncWindow:
    if explicit_start_date or explicit_end_date:
        start_date_value = validate_date(explicit_start_date or config.start_date)
        end_date_value = validate_date(explicit_end_date or config.end_date)
        if end_date_value < start_date_value:
            raise ValueError("end-date 不能早于 start-date")
        return SyncWindow(start_date=start_date_value, end_date=end_date_value)

    configured_start_date = validate_date(config.start_date)
    yesterday = validate_date(config.end_date)
    rolling_start_date = max(configured_start_date, yesterday - timedelta(days=4))

    if mode == "notify":
        state = load_prefect_sync_state()
        last_sync_start = state.get("last_sync_start_date")
        last_sync_end = state.get("last_sync_end_date")
        if last_sync_start and last_sync_end:
            start_date_value = validate_date(last_sync_start)
            end_date_value = validate_date(last_sync_end)
            if end_date_value < start_date_value:
                raise ValueError("prefect_sync_state.json 中的日期范围无效")
            return SyncWindow(start_date=start_date_value, end_date=end_date_value)
        return SyncWindow(start_date=rolling_start_date, end_date=yesterday)

    return SyncWindow(start_date=rolling_start_date, end_date=yesterday)


def mark_sync_window_completed(window: SyncWindow) -> None:
    state = load_prefect_sync_state()
    state["last_sync_start_date"] = window.start_date.isoformat()
    state["last_sync_end_date"] = window.end_date.isoformat()
    save_prefect_sync_state(state)


@task(name="notify-feishu")
def notify_feishu(context: SyncConfig, start_date_text: str, end_date_text: str) -> None:
    context_obj = build_daily_sync_context(
        DailySyncConfig(
            base_url=context.base_url,
            mapping_file=Path(context.mapping_file),
            store_field_name=context.store_field_name,
            table_map_file=context.table_map_file,
            dry_run=context.dry_run,
            skip_missing_tables=context.skip_missing_tables,
            chat_id=context.chat_id,
            report_image_path=context.report_image_path,
            group_webhook=context.group_webhook,
        )
    )
    send_sync_notification(context_obj, validate_date(start_date_text), validate_date(end_date_text))


@flow(name="gmv-daily-data-sync")
def gmv_daily_data_sync_flow(
    start_date: str | None = None,
    end_date: str | None = None,
    dry_run: bool = False,
    skip_missing_tables: bool = True,
    table_map_file: str = "",
) -> None:
    config = load_sync_config(
        start_date=start_date,
        end_date=end_date,
        dry_run=dry_run,
        skip_missing_tables=skip_missing_tables,
        table_map_file=table_map_file,
    )
    sync_window = resolve_sync_window(config, start_date, end_date, mode="sync")
    context = build_daily_sync_context(
        DailySyncConfig(
            base_url=config.base_url,
            mapping_file=Path(config.mapping_file),
            store_field_name=config.store_field_name,
            table_map_file=config.table_map_file,
            dry_run=config.dry_run,
            skip_missing_tables=config.skip_missing_tables,
            chat_id=config.chat_id,
            report_image_path=config.report_image_path,
            group_webhook=config.group_webhook,
        )
    )

    log_step(f"Prefect 数据同步开始: {sync_window.start_date.isoformat()} -> {sync_window.end_date.isoformat()}")
    for stat_date in iter_dates(sync_window.start_date, sync_window.end_date):
        log_step(f"Prefect 开始处理日期: {stat_date.isoformat()}")
        run_daily_sync_for_date(context, stat_date)
    if not config.dry_run:
        mark_sync_window_completed(sync_window)
    if context.missing_tables:
        log_step(f"Prefect 数据同步发现缺失表: {len(context.missing_tables)} 个")
        send_missing_table_notification(
            context,
            context.missing_tables,
            sync_window.start_date,
            sync_window.end_date,
        )
    log_step("Prefect 数据同步完成")


@flow(name="gmv-daily-notify")
def gmv_daily_notify_flow(
    start_date: str | None = None,
    end_date: str | None = None,
    dry_run: bool = False,
    skip_missing_tables: bool = True,
    table_map_file: str = "",
) -> None:
    config = load_sync_config(
        start_date=start_date,
        end_date=end_date,
        dry_run=dry_run,
        skip_missing_tables=skip_missing_tables,
        table_map_file=table_map_file,
    )
    sync_window = resolve_sync_window(config, start_date, end_date, mode="notify")
    log_step(f"Prefect 飞书通知开始: {sync_window.start_date.isoformat()} -> {sync_window.end_date.isoformat()}")
    notify_feishu(config, sync_window.start_date.isoformat(), sync_window.end_date.isoformat())
    log_step("Prefect 飞书通知完成")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Prefect 调度：每天 8 点同步前 5 天 GMV 数据，9 点发送飞书通知。")
    parser.add_argument("--start-date", default="", help="开始日期，格式 YYYY-MM-DD，默认读取 .env 中的 SYNC_START_DATE")
    parser.add_argument("--end-date", default="", help="结束日期，格式 YYYY-MM-DD，默认昨天")
    parser.add_argument("--table-map-file", default="", help="可选，项目到表名的映射 JSON 文件路径")
    parser.add_argument("--dry-run", action="store_true", help="只输出将要更新的内容，不调用飞书写接口")
    parser.add_argument(
        "--skip-missing-tables",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="找不到对应月份的飞书表时跳过并在同步结束后推送提醒（默认开启）",
    )
    parser.add_argument("--serve", action="store_true", help="启动 Prefect deployments，按 8 点同步、9 点通知 定时执行")
    parser.add_argument(
        "--mode",
        choices=("sync", "notify", "all"),
        default="all",
        help="单次执行模式：sync=只跑数据同步，notify=只发飞书通知，all=先同步再通知",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    flow_kwargs = {
        "start_date": args.start_date or None,
        "end_date": args.end_date or None,
        "dry_run": args.dry_run,
        "skip_missing_tables": args.skip_missing_tables,
        "table_map_file": args.table_map_file,
    }
    if args.serve:
        gmv_daily_data_sync_flow.serve(
            name="gmv-daily-data-sync-8am",
            cron="0 8 * * *",
            parameters=flow_kwargs,
            tags=["gmv", "feishu", "bitable"],
        )
        gmv_daily_notify_flow.serve(
            name="gmv-daily-notify-9am",
            cron="0 9 * * *",
            parameters=flow_kwargs,
            tags=["gmv", "feishu", "notify"],
        )
        return
    if args.mode == "sync":
        gmv_daily_data_sync_flow(**flow_kwargs)
        return
    if args.mode == "notify":
        gmv_daily_notify_flow(**flow_kwargs)
        return
    gmv_daily_data_sync_flow(**flow_kwargs)
    gmv_daily_notify_flow(**flow_kwargs)


if __name__ == "__main__":
    main()