market_sync/bin/task_watch.py

"""自适应任务跟踪监视器。

用法:
  .venv/bin/python bin/task_watch.py          # 单次检查
  .venv/bin/python bin/task_watch.py --watch   # 持续监视（自动调整检查间隔）

设计：
  1. 查询 dataset_registry 中 status='running' 的任务
  2. 从 config 表中解析历史运行耗时（解析 lastMessage 中的 elapsed_sec）
  3. 估算剩余时间 = 历史平均耗时 - 已用时间
  4. 自适应下次检查间隔 = clamp(剩余 * 0.15, 15s, 300s)
     无历史数据时默认 60s
  5. 持续监视模式下，每次检查后动态调整下一次的 sleep 时间
"""
from __future__ import annotations

import argparse
import json
import re
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

_PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

from sqlalchemy import create_engine, text

from app.core.config import settings
from app.core.utils.logging import setup_logging, get_logger

setup_logging()
logger = get_logger("task_watch")

# ── 历史耗时解析 ─────────────────────────────────────────────────────

_ELAPSED_RE = re.compile(r"elapsed[=_ ]sec[= ]?(\d+\.?\d*)", re.IGNORECASE)
_ELAPSED_SUFFIX_RE = re.compile(r"(\d+\.?\d*)s")


def _parse_elapsed_sec(message: str) -> float | None:
    for pattern in (_ELAPSED_RE,):
        m = pattern.search(message)
        if m:
            return float(m.group(1))
    return None


# ── 历史耗时数据库（从 config 表 lastMessage 解析）───────────────────


def _load_history(conn: Any) -> dict[str, list[float]]:
    history: dict[str, list[float]] = {}
    rows = conn.execute(
        text("""
            SELECT key, value
            FROM market_data.config
            WHERE category = 'schedule'
        """)
    ).fetchall()
    for key, value_json in rows:
        try:
            val = json.loads(value_json)
        except (json.JSONDecodeError, TypeError):
            continue
        job = val.get("job") or key.removeprefix("schedule_")
        msg = val.get("lastMessage") or ""
        elapsed = _parse_elapsed_sec(msg)
        if elapsed and elapsed > 0:
            history.setdefault(job, []).append(elapsed)
    return history


# ── 活跃任务查询 ──────────────────────────────────────────────────────


def _fetch_running_tasks(conn: Any) -> list[dict[str, Any]]:
    rows = conn.execute(
        text("""
            SELECT dataset_id, status, started_at, message
            FROM market_data.dataset_registry
            WHERE status = 'running'
            ORDER BY started_at NULLS LAST
        """)
    ).fetchall()
    tasks = []
    for row in rows:
        tasks.append({
            "dataset_id": row[0],
            "status": row[1],
            "started_at": row[2],
            "message": row[3] or "",
        })
    return tasks


# ── 估算 ──────────────────────────────────────────────────────────────


def _estimate(
    task: dict[str, Any],
    history: dict[str, list[float]],
) -> tuple[float | None, float | None, float | None]:
    """返回 (历史平均耗时秒, 已用秒, 预估剩余秒)。"""
    tid = task["dataset_id"]
    starts = task["started_at"]
    if starts is None:
        return None, None, None

    started_wall = starts.replace(tzinfo=None)
    elapsed = (datetime.now() - started_wall).total_seconds()
    elapsed = max(round(elapsed, 1), 0)

    durations = history.get(tid, [])
    if not durations:
        return None, elapsed, None

    avg_duration = sum(durations) / len(durations)
    remaining = max(round(avg_duration - elapsed, 1), 0)
    return round(avg_duration, 1), elapsed, remaining


# ── 自适应检查间隔 ────────────────────────────────────────────────────


def _next_interval(remaining: float | None) -> float:
    if remaining is None:
        return 60.0
    interval = remaining * 0.15
    return max(15.0, min(interval, 300.0))


# ── 输出 ──────────────────────────────────────────────────────────────


def _fmt_sec(sec: float | None) -> str:
    if sec is None:
        return "N/A"
    if sec < 60:
        return f"{sec:.0f}s"
    if sec < 3600:
        return f"{sec/60:.1f}min"
    return f"{sec/3600:.1f}h"


def _print_report(
    tasks: list[dict[str, Any]],
    history: dict[str, list[float]],
    next_interval: float,
) -> None:
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    header = f"── Task Watch [{now}] ──"
    print(header)
    if not tasks:
        print("  无运行中的任务")
        print(f"  下次检查: {_fmt_sec(next_interval)}")
        return

    lines = []
    for t in tasks:
        avg, elapsed, remaining = _estimate(t, history)
        durations = history.get(t["dataset_id"], [])
        progress = ""
        if avg and elapsed:
            pct = min(elapsed / avg * 100, 99.9)
            progress = f" {pct:.0f}%"
        line = (
            f"  {t['dataset_id']:20s}"
            f" elapsed={_fmt_sec(elapsed):>8s}"
            f" avg={_fmt_sec(avg):>8s}"
            f" remain={_fmt_sec(remaining):>8s}"
            f"{progress}"
            f"  (历史 {len(durations)} 次)"
        )
        lines.append(line)

    for l in lines:
        print(l)
    print(f"  下次检查: {_fmt_sec(next_interval)}")
    print("─" * len(header))


# ── 主循环 ────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(description="自适应任务跟踪监视器")
    parser.add_argument("--watch", action="store_true", help="持续监视模式（自动调整检查间隔）")
    args = parser.parse_args()

    engine = create_engine(settings.pg_sqlalchemy_url())

    while True:
        with engine.connect() as conn:
            history = _load_history(conn)
            tasks = _fetch_running_tasks(conn)

        interval = 60.0
        if tasks:
            intervals = []
            for t in tasks:
                _, _, remaining = _estimate(t, history)
                intervals.append(_next_interval(remaining))
            interval = min(intervals) if intervals else 60.0

        _print_report(tasks, history, interval)

        if not args.watch:
            break

        time.sleep(interval)


if __name__ == "__main__":
    main()