vent_agent_2.0/db/schedule_store.py

"""定时任务定义与执行记录的存取。"""

from __future__ import annotations

from typing import Any

from db import db_conn
from obs import get_logger

logger = get_logger(__name__)


def create_schedule(
    user_name: str,
    title: str,
    cron: str,
    prompt: str,
    mode: str = "full",
    enabled: bool = True,
) -> dict[str, Any]:
    """新建定时任务定义。"""
    schedule_id = db_conn.new_id()
    stamp = db_conn.fmt_dt()
    with db_conn.transaction():
        db_conn.execute(
            "INSERT INTO schedules (id, user_name, title, cron, prompt, mode, enabled, created_at, updated_at)"
            " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                schedule_id,
                user_name,
                title,
                cron,
                prompt,
                mode,
                1 if enabled else 0,
                stamp,
                stamp,
            ),
        )
    logger.info("新建定时任务 id=%s cron=%s user=%s", schedule_id, cron, user_name)
    return get_schedule(schedule_id) or {}


def get_schedule(schedule_id: str) -> dict[str, Any] | None:
    """按 ID 取任务定义。"""
    row = db_conn.query_one("SELECT * FROM schedules WHERE id = ?", (schedule_id,))
    return _hydrate(row) if row else None


def list_schedules(
    user_name: str | None = None, enabled_only: bool = False, limit: int = 200
) -> list[dict[str, Any]]:
    """列出任务定义。``user_name=None`` 表示不限用户（调度器装载用）。"""
    clauses: list[str] = []
    params: list[Any] = []
    if user_name is not None:
        clauses.append("user_name = ?")
        params.append(user_name)
    if enabled_only:
        clauses.append("enabled = 1")
    where = f" WHERE {' AND '.join(clauses)}" if clauses else ""
    params.append(limit)
    rows = db_conn.query(
        f"SELECT * FROM schedules{where} ORDER BY created_at DESC LIMIT ?", tuple(params)
    )
    return [_hydrate(row) for row in rows]


def update_schedule(schedule_id: str, **fields: Any) -> dict[str, Any] | None:
    """更新任务定义的指定字段（白名单过滤，未知字段忽略）。"""
    allowed = {"title", "cron", "prompt", "mode", "enabled"}
    updates = {k: v for k, v in fields.items() if k in allowed and v is not None}
    if not updates:
        return get_schedule(schedule_id)
    if "enabled" in updates:
        updates["enabled"] = 1 if updates["enabled"] else 0
    assignments = ", ".join(f"{key} = ?" for key in updates)
    params = list(updates.values()) + [db_conn.fmt_dt(), schedule_id]
    db_conn.execute(
        f"UPDATE schedules SET {assignments}, updated_at = ? WHERE id = ?", tuple(params)
    )
    return get_schedule(schedule_id)


def set_enabled(schedule_id: str, enabled: bool) -> dict[str, Any] | None:
    """启停任务。"""
    return update_schedule(schedule_id, enabled=enabled)


def delete_schedule(schedule_id: str) -> bool:
    """删除任务定义及其执行记录。"""
    with db_conn.transaction():
        db_conn.execute("DELETE FROM schedule_runs WHERE schedule_id = ?", (schedule_id,))
        affected = db_conn.execute("DELETE FROM schedules WHERE id = ?", (schedule_id,))
    return affected > 0


# ============================================================
# 执行记录
# ============================================================

def start_run(schedule_id: str, thread_id: str) -> str:
    """登记一次执行开始，返回 run_id。"""
    run_id = db_conn.new_id()
    db_conn.execute(
        "INSERT INTO schedule_runs (id, schedule_id, thread_id, status, started_at)"
        " VALUES (?, ?, ?, 'running', ?)",
        (run_id, schedule_id, thread_id, db_conn.fmt_dt()),
    )
    return run_id


def finish_run(
    run_id: str, status: str, error: str | None = None, output: str | None = None
) -> None:
    """登记一次执行结束。status: success | waiting_input | failed | cancelled。"""
    db_conn.execute(
        "UPDATE schedule_runs SET status = ?, finished_at = ?, error = ?, output = ? WHERE id = ?",
        (status, None if status == "waiting_input" else db_conn.fmt_dt(), error, output, run_id),
    )


def list_runs(schedule_id: str, limit: int = 20) -> list[dict[str, Any]]:
    """列出某任务的执行记录（最近优先）。"""
    return db_conn.query(
        "SELECT * FROM schedule_runs WHERE schedule_id = ?"
        " ORDER BY started_at DESC, id DESC LIMIT ?",
        (schedule_id, limit),
    )


def latest_run(schedule_id: str) -> dict[str, Any] | None:
    """取最近一次执行记录。"""
    runs = list_runs(schedule_id, limit=1)
    return runs[0] if runs else None


def _hydrate(row: dict[str, Any]) -> dict[str, Any]:
    schedule = dict(row)
    schedule["enabled"] = bool(schedule.get("enabled"))
    return schedule