xiaobai-review/xiaobai-datahub/datahub/cli.py

"""Command-line entry for one-shot datahub operations."""

from __future__ import annotations

import argparse
import json
import sys

from datahub.hub import build_hub
from datahub.pipeline import OFFICIAL_DATASETS
from datahub.settings import load_settings
from datahub.timeutil import yyyymmdd


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="xiaobai-datahub CLI")
    sub = parser.add_subparsers(dest="command", required=True)
    history = sub.add_parser("history-backfill", help="回补 2016 年起交易日历和网站所用指数日 K")
    history.add_argument("--calendar-start", default=None, help="日历起点，默认配置 calendar_start")
    history.add_argument("--index-days", type=int, default=None, help="指数回补交易日数量，默认 260")
    history.add_argument("--force", action="store_true", help="覆盖已发布的指数日期")
    refresh = sub.add_parser("eod-refresh", help="对指定交易日补跑盘后正式数据（跳过已发布数据集，仍走质量门禁）")
    refresh.add_argument("--trade-date", default=None, help="交易日 YYYYMMDD，默认今天")
    refresh.add_argument(
        "--force", action="store_true",
        help="对 --dataset 指定的数据集强制重取重发（生成新批次，保留上一批次可回滚）",
    )
    refresh.add_argument(
        "--dataset", default=None,
        help="配合 --force 使用：只强制重发该数据集（如 valuation）",
    )
    stocks_refresh = sub.add_parser("stocks-refresh", help="刷新股票主档并发布正式快照（幂等：无变化则跳过）")
    stocks_refresh.add_argument("--trade-date", default=None, help="交易日 YYYYMMDD，默认今天")
    stocks_refresh.add_argument("--force", action="store_true", help="即使快照无变化也重新发布")
    moneyflow_backfill = sub.add_parser(
        "moneyflow-backfill", help="回补资金流历史（默认覆盖网站所需的最近 N 个交易日，跳过已发布日期）",
    )
    moneyflow_backfill.add_argument("--end-date", default=None, help="截止交易日 YYYYMMDD，默认今天")
    moneyflow_backfill.add_argument("--trading-days", type=int, default=None, help="回补交易日数量，默认配置 moneyflow_history_trading_days")
    moneyflow_backfill.add_argument("--force", action="store_true", help="覆盖已发布的资金流日期")
    args = parser.parse_args(argv)

    settings = load_settings()
    hub = build_hub(settings)
    if args.command == "history-backfill":
        result = hub.pipeline.backfill_history(
            calendar_start=args.calendar_start,
            index_days=args.index_days,
            force=args.force,
        )
        json.dump(result, sys.stdout, ensure_ascii=False, indent=2, default=str)
        sys.stdout.write("\n")
        return 0 if result.get("ok") else 1
    if args.command == "eod-refresh":
        day = yyyymmdd(args.trade_date) if args.trade_date else yyyymmdd()
        if args.force:
            datasets = tuple(sorted({args.dataset} & OFFICIAL_DATASETS)) if args.dataset else ()
            if args.dataset and not datasets:
                parser.error(f"unknown dataset: {args.dataset}")
            if not datasets:
                parser.error("--force requires --dataset (e.g. --dataset valuation)")
            result = {}
            for dataset in datasets:
                result[dataset] = hub.pipeline.run_dataset(dataset, day)
        else:
            result = hub.pipeline.run_eod_missing(day)
        hub.pipeline.audit("cli", "eod-refresh", f"eod:{day}", json.dumps(
            {"force": bool(args.force), "dataset": args.dataset,
             **{name: item.get("state") for name, item in result.items() if isinstance(item, dict)}},
            ensure_ascii=False,
        ))
        if args.force:
            payload = {"trade_date": day, "datasets": result}
            json.dump(payload, sys.stdout, ensure_ascii=False, indent=2, default=str)
            sys.stdout.write("\n")
            return 0
        missing = hub.pipeline.missing_official_datasets(day)
        payload = {"trade_date": day, "datasets": result, "missing_after": missing}
        json.dump(payload, sys.stdout, ensure_ascii=False, indent=2, default=str)
        sys.stdout.write("\n")
        return 0 if not missing else 1
    if args.command == "stocks-refresh":
        day = yyyymmdd(args.trade_date) if args.trade_date else yyyymmdd()
        result = hub.pipeline.refresh_stocks(day, force=args.force)
        hub.pipeline.audit("cli", "stocks-refresh", f"stocks:{day}", json.dumps(
            {"force": bool(args.force), "state": result.get("state"), "batch_id": result.get("batch_id")},
            ensure_ascii=False,
        ))
        json.dump(result, sys.stdout, ensure_ascii=False, indent=2, default=str)
        sys.stdout.write("\n")
        return 0 if result.get("state") != "failed" else 1
    if args.command == "moneyflow-backfill":
        result = hub.pipeline.backfill_moneyflow_history(
            end_date=args.end_date,
            trading_days=args.trading_days,
            force=args.force,
        )
        hub.pipeline.audit("cli", "moneyflow-backfill", f"moneyflow:{result.get('end')}", json.dumps(
            {"published": len(result.get("published") or []), "skipped": len(result.get("skipped") or []),
             "failed": len(result.get("failed") or [])},
            ensure_ascii=False,
        ))
        json.dump(result, sys.stdout, ensure_ascii=False, indent=2, default=str)
        sys.stdout.write("\n")
        return 0 if result.get("ok") else 1
    parser.error(f"unknown command: {args.command}")
    return 2


if __name__ == "__main__":
    raise SystemExit(main())