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)}