demand-agent-new/examples/demand/run_hive_gap_mysql.py

"""Run DemandAgent from Hive supply-gap rows and write only MySQL demand_content.

This entrypoint restores the old business starting point (ODPS/Hive supply
gap) without restoring prepare/run_mining or any Hive writes.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Any


sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from examples.demand.pg_pattern_repository import normalize_platform


def _load_project_env() -> None:
    try:
        from dotenv import load_dotenv
    except ImportError:
        return
    load_dotenv(Path(__file__).parent.parent.parent / ".env")


def _safe_path_segment(value: str) -> str:
    return "".join(ch if ch.isalnum() or ch in {"-", "_"} else "_" for ch in value).strip("_") or "run"


def _prepare_run_output_paths(execution_id: int, run_label: str) -> None:
    run_root = Path(__file__).parent / "test_output_data" / "mysql_runs" / _safe_path_segment(run_label)
    result_base_dir = run_root / "intermediate" / "result"
    os.environ["DEMAND_RESULT_BASE_DIR"] = str(result_base_dir)
    os.environ["DEMAND_OUTPUT_BASE_DIR"] = str(run_root / "output")
    os.environ["DEMAND_TRACE_STORE_PATH"] = str(run_root / ".trace")
    os.environ["DEMAND_WEIGHT_DATA_DIR"] = str(run_root / "intermediate" / "data")

    demand_items_path = result_base_dir / str(execution_id) / f"execution_id_{execution_id}_demand_items.json"
    if demand_items_path.exists():
        demand_items_path.unlink()


def _configure_mysql_env(run_label: str) -> None:
    os.environ["DEMAND_OUTPUT_MODE"] = "mysql_demand_content"
    os.environ["DEMAND_MYSQL_ENTRYPOINT"] = "run_hive_gap_mysql"
    os.environ["DEMAND_RUN_LABEL"] = run_label


def _validate_execution_success(execution_id: int) -> None:
    from examples.demand.db_manager import query_execution_for_evidence

    execution = query_execution_for_evidence(execution_id)
    if not execution:
        raise ValueError(f"execution_id={execution_id} 不存在，Hive gap 入口只允许跑已有 PG execution")
    if str(execution.get("status") or "").strip().lower() != "success":
        raise ValueError(
            f"execution_id={execution_id} status={execution.get('status')}，不是 success，拒绝生成需求"
        )


def _build_scope(row: dict[str, Any], execution_id: int, count: int) -> dict[str, Any]:
    return {
        "scope_source": "odps_gap",
        "merge_leve2": row["cluster_name"],
        "platform": normalize_platform(row.get("platform_type") or "piaoquan"),
        "gap_dt": row.get("gap_dt"),
        "requested_count": int(count),
        "lack_count": row.get("lack_count"),
        "pattern_execution_id": int(execution_id),
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run DemandAgent from Hive gap rows and write only MySQL demand_content."
    )
    parser.add_argument("--execution-id", type=int, default=None, help="Existing PG pattern_mining_execution.id")
    parser.add_argument("--skip-categories", type=int, default=0, help="Skip the first N Hive gap categories")
    parser.add_argument("--max-categories", type=int, default=None, help="Only run the first N gap categories")
    parser.add_argument("--max-total-count", type=int, default=None, help="Cap total requested demand count")
    parser.add_argument("--per-category-count", type=int, default=None, help="Override each selected gap row count")
    parser.add_argument("--run-label-prefix", default=None, help="Prefix stored in ext_data.run_label")
    parser.add_argument("--dry-run", action="store_true", help="Print gap plan only; do not call the agent")
    parser.add_argument(
        "--fail-on-partition-not-ready",
        action="store_true",
        help="Exit non-zero when yesterday's Hive source partition is not ready; useful for systemd retry.",
    )
    return parser.parse_args()


async def async_main() -> dict[str, Any]:
    _load_project_env()
    args = parse_args()

    from examples.demand.data_query_tools import (
        get_demand_gap_source_partition_status,
        get_demand_merge_level2_names,
    )
    from examples.demand.db_manager import query_latest_success_execution_id

    execution_id = args.execution_id or query_latest_success_execution_id()
    if not execution_id:
        raise ValueError("未找到 PG Pattern V2 success execution")
    _validate_execution_success(int(execution_id))

    partition_status = get_demand_gap_source_partition_status()
    if not partition_status.get("ready"):
        return {
            "execution_id": int(execution_id),
            "dry_run": bool(args.dry_run),
            "skipped_reason": "hive_gap_source_partition_not_ready",
            "partition_status": partition_status,
            "planned": [],
            "results": [],
            "_exit_code": 75 if args.fail_on_partition_not_ready else 0,
        }

    gap_rows = get_demand_merge_level2_names(day=str(partition_status["dt"]))
    if not gap_rows:
        return {
            "execution_id": int(execution_id),
            "dry_run": bool(args.dry_run),
            "skipped_reason": "hive_gap_rows_empty",
            "partition_status": partition_status,
            "planned": [],
            "results": [],
        }
    if args.skip_categories:
        gap_rows = gap_rows[max(int(args.skip_categories), 0):]
    if args.max_categories is not None:
        gap_rows = gap_rows[: max(int(args.max_categories), 0)]

    remaining = int(args.max_total_count) if args.max_total_count is not None else None
    prefix = args.run_label_prefix or f"hive_gap_mysql_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    results: list[dict[str, Any]] = []
    planned: list[dict[str, Any]] = []

    for index, row in enumerate(gap_rows, start=1):
        requested = int(row.get("count") or 0)
        if args.per_category_count is not None:
            requested = max(int(args.per_category_count), 0)
        if remaining is not None:
            if remaining <= 0:
                break
            requested = min(requested, remaining)
        if requested <= 0:
            continue

        cluster_name = str(row["cluster_name"])
        platform_type = str(row.get("platform_type") or "piaoquan")
        run_label = f"{prefix}_{index:02d}_{_safe_path_segment(cluster_name)}"
        scope = _build_scope(row, int(execution_id), requested)
        planned.append({"run_label": run_label, "count": requested, "demand_scope": scope})

        if args.dry_run:
            if remaining is not None:
                remaining -= requested
            continue

        from examples.demand.run import main as run_main

        _configure_mysql_env(run_label)
        _prepare_run_output_paths(int(execution_id), run_label)
        result = await run_main(
            cluster_name=cluster_name,
            platform_type=platform_type,
            count=requested,
            execution_id=int(execution_id),
            task_id=None,
            demand_scope=scope,
        )
        result["run_label"] = run_label
        result["demand_scope"] = scope
        results.append(result)
        if remaining is not None:
            remaining -= requested

    return {
        "execution_id": int(execution_id),
        "dry_run": bool(args.dry_run),
        "partition_status": partition_status,
        "planned": planned,
        "results": results,
    }


def main() -> None:
    result = asyncio.run(async_main())
    exit_code = int(result.pop("_exit_code", 0) or 0)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    if exit_code:
        raise SystemExit(exit_code)


if __name__ == "__main__":
    main()