xiaobai-review/xiaobai-datahub/datahub/scheduler.py

from __future__ import annotations

import threading
from collections.abc import Callable
from datetime import datetime, time
from typing import Any

from datahub.db import HubDB
from datahub.logutil import get_logger
from datahub.pipeline import Pipeline
from datahub.timeutil import isoformat, now_shanghai, yyyymmdd

LOGGER = get_logger()

JobFn = Callable[[str], Any]


def is_open_day(db: HubDB, day: str) -> bool:
    row = db.fetchone(
        "SELECT is_open FROM trade_calendar WHERE exchange = 'SSE' AND cal_date = ?",
        (day,),
    )
    if row is None:
        return True  # unknown calendar: do not skip reference refresh
    return int(row["is_open"]) == 1


class Scheduler:
    """Calendar-driven in-process scheduler. Non-trading days skip EOD fetches."""

    def __init__(self, db: HubDB, pipeline: Pipeline, jobs: dict[str, JobFn] | None = None) -> None:
        self.db = db
        self.pipeline = pipeline
        self.jobs = jobs or {
            "precheck": self._precheck,
            "eod_a": self._eod_a,
            "eod_b": self._eod_b,
            "cleanup": self._cleanup,
            "backup": self._backup,
            "history_backfill": self._history_backfill,
        }
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._fired: set[tuple[str, str, str]] = set()

    def start(self, interval_seconds: float = 30.0) -> None:
        if self._thread and self._thread.is_alive():
            return

        def loop() -> None:
            while not self._stop.wait(interval_seconds):
                try:
                    self.tick()
                except Exception:
                    LOGGER.exception("scheduler tick failed")

        self._thread = threading.Thread(target=loop, name="datahub-scheduler", daemon=True)
        self._thread.start()

    def stop(self, timeout: float = 5.0) -> None:
        self._stop.set()
        if self._thread and self._thread is not threading.current_thread():
            self._thread.join(timeout)

    def tick(self, clock: datetime | None = None) -> list[str]:
        now = clock or now_shanghai()
        day = yyyymmdd(now)
        current = now.timetz() if False else now.time()
        ran: list[str] = []
        plan = [
            ("precheck", time(8, 45)),
            ("eod_a", time(15, 5)),
            ("eod_b", time(15, 10)),
            ("cleanup", time(0, 30)),
            ("backup", time(0, 40)),
        ]
        open_day = is_open_day(self.db, day)
        for job_id, at in plan:
            if current < at:
                continue
            key = (job_id, day, at.strftime("%H%M"))
            if key in self._fired:
                continue
            if job_id in {"eod_a", "eod_b"} and not open_day:
                self._fired.add(key)
                continue
            self._fired.add(key)
            self.run_job(job_id, day)
            ran.append(job_id)
        return ran

    def run_job(self, job_id: str, trade_date: str) -> dict[str, Any]:
        fn = self.jobs.get(job_id)
        if fn is None:
            raise KeyError(job_id)
        started = isoformat()
        run_id = None
        with self.db.write() as connection:
            cur = connection.execute(
                "INSERT INTO job_runs(job_id, state, started_at, attempt) VALUES (?,?,?,1)",
                (job_id, "running", started),
            )
            run_id = cur.lastrowid
        try:
            result = fn(trade_date) or {}
            with self.db.write() as connection:
                connection.execute(
                    "UPDATE job_runs SET state=?, finished_at=?, rows_out=?, detail=? WHERE id=?",
                    ("ok", isoformat(), result.get("rows") if isinstance(result, dict) else None, str(result)[:2000], run_id),
                )
            return {"job_id": job_id, "result": result, "state": "ok"}
        except Exception as exc:
            with self.db.write() as connection:
                connection.execute(
                    "UPDATE job_runs SET state=?, finished_at=?, error=? WHERE id=?",
                    ("failed", isoformat(), str(exc), run_id),
                )
            raise

    def _precheck(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.ingest_reference(trade_date)

    def _eod_a(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.run_eod_batch_a(trade_date)

    def _eod_b(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.run_eod_batch_b(trade_date)

    def _history_backfill(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.backfill_history(trade_date)

    def _cleanup(self, trade_date: str) -> dict[str, Any]:
        result = self.pipeline.cleanup()
        if now_shanghai().weekday() == 6:
            self.pipeline.db.vacuum()
            result["vacuum"] = True
        return result

    def _backup(self, trade_date: str) -> dict[str, Any]:
        from pathlib import Path

        dest_dir = Path(self.pipeline.settings.backup_dir)
        dest = dest_dir / f"datahub-{trade_date}.db"
        self.pipeline.db.backup_to(dest)
        keep = int(self.pipeline.settings.quality.get("backup_retain") or 14)
        backups = sorted(dest_dir.glob("datahub-*.db"))
        for old in backups[:-keep]:
            old.unlink(missing_ok=True)
        return {"path": str(dest.name), "kept": min(len(backups), keep)}