crabcode 0.1.6__py3-none-win_amd64.whl

crabcode_core/schedule/store.py

@@ -0,0 +1,558 @@
+"""SQLite schedule store — persistent storage for scheduled jobs and their runs."""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from typing import Any
+
+from crabcode_core.paths import get_config_home
+from crabcode_core.schedule.models import JobRun, ScheduleJob
+
+
+def _db_path() -> Path:
+    """Database path below the configured CrabCode home."""
+    return get_config_home() / "schedules.db"
+
+
+_SCHEMA = """\
+CREATE TABLE IF NOT EXISTS schedules (
+    id TEXT PRIMARY KEY,
+    name TEXT NOT NULL,
+    prompt TEXT NOT NULL,
+    schedule TEXT NOT NULL,
+    schedule_type TEXT NOT NULL,
+    cwd TEXT,
+    enabled INTEGER NOT NULL DEFAULT 1,
+    status TEXT NOT NULL DEFAULT 'active',
+    last_run TEXT,
+    next_run TEXT,
+    run_count INTEGER NOT NULL DEFAULT 0,
+    max_runs INTEGER,
+    created_at TEXT NOT NULL,
+    session_id TEXT,
+    description TEXT NOT NULL DEFAULT '',
+    tags TEXT NOT NULL DEFAULT '[]',
+    timeout INTEGER,
+    model_profile TEXT,
+    extra TEXT NOT NULL DEFAULT '{}',
+    claimed_by TEXT,
+    claimed_until TEXT
+);
+
+CREATE TABLE IF NOT EXISTS schedule_runs (
+    id TEXT PRIMARY KEY,
+    job_id TEXT NOT NULL,
+    status TEXT NOT NULL DEFAULT 'pending',
+    started_at TEXT,
+    finished_at TEXT,
+    duration_seconds REAL,
+    session_id TEXT,
+    exit_code INTEGER,
+    error_message TEXT,
+    result_summary TEXT NOT NULL DEFAULT '',
+    tokens_used INTEGER NOT NULL DEFAULT 0,
+    created_at TEXT NOT NULL,
+    FOREIGN KEY (job_id) REFERENCES schedules(id) ON DELETE CASCADE
+);
+
+CREATE INDEX IF NOT EXISTS idx_schedule_runs_job
+    ON schedule_runs(job_id, created_at DESC);
+
+CREATE INDEX IF NOT EXISTS idx_schedules_pending
+    ON schedules(next_run)
+    WHERE status = 'active' AND enabled = 1;
+"""
+
+_MIGRATIONS: list[str] = [
+    "ALTER TABLE schedules ADD COLUMN claimed_by TEXT",
+    "ALTER TABLE schedules ADD COLUMN claimed_until TEXT",
+]
+
+
+def _get_conn(db_path: Path) -> sqlite3.Connection:
+    """Open (or create) the SQLite database and ensure the schema exists."""
+    db_path.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(str(db_path))
+    conn.execute("PRAGMA busy_timeout=5000")
+    conn.execute("PRAGMA journal_mode=WAL")
+    conn.execute("PRAGMA foreign_keys=ON")
+    conn.executescript(_SCHEMA)
+    for stmt in _MIGRATIONS:
+        try:
+            conn.execute(stmt)
+            conn.commit()
+        except sqlite3.OperationalError:
+            pass  # column/table already exists
+    return conn
+
+
+def _row_to_schedule(row: tuple[Any, ...], cols: list[str]) -> dict[str, Any]:
+    """Convert a raw DB row into a dict, deserialising JSON columns."""
+    d = dict(zip(cols, row))
+    # Deserialize JSON columns
+    for key in ("tags", "extra"):
+        val = d.get(key)
+        if isinstance(val, str):
+            try:
+                d[key] = json.loads(val)
+            except (json.JSONDecodeError, TypeError):
+                d[key] = [] if key == "tags" else {}
+    # Convert SQLite integer booleans
+    if "enabled" in d and isinstance(d["enabled"], int):
+        d["enabled"] = bool(d["enabled"])
+    return d
+
+
+def _row_to_run(row: tuple[Any, ...], cols: list[str]) -> dict[str, Any]:
+    """Convert a raw schedule_runs row into a dict."""
+    d = dict(zip(cols, row))
+    return d
+
+
+class ScheduleStore:
+    """SQLite-backed persistence for :class:`ScheduleJob` and :class:`JobRun`."""
+
+    def __init__(self, db_path: Path | None = None):
+        self._db_path = db_path or _db_path()
+        self._conn: sqlite3.Connection | None = None
+
+    # ------------------------------------------------------------------
+    # Connection lifecycle
+    # ------------------------------------------------------------------
+
+    def _conn_or_create(self) -> sqlite3.Connection:
+        if self._conn is None:
+            self._conn = _get_conn(self._db_path)
+        return self._conn
+
+    def close(self) -> None:
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+
+    # ------------------------------------------------------------------
+    # Schedule CRUD
+    # ------------------------------------------------------------------
+
+    def upsert_schedule(self, job: ScheduleJob) -> None:
+        """Insert or update a :class:`ScheduleJob` row."""
+        conn = self._conn_or_create()
+        conn.execute(
+            """INSERT INTO schedules
+               (id, name, prompt, schedule, schedule_type, cwd, enabled, status,
+                last_run, next_run, run_count, max_runs, created_at, session_id,
+                description, tags, timeout, model_profile, extra)
+               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+               ON CONFLICT(id) DO UPDATE SET
+                 name = excluded.name,
+                 prompt = excluded.prompt,
+                 schedule = excluded.schedule,
+                 schedule_type = excluded.schedule_type,
+                 cwd = excluded.cwd,
+                 enabled = excluded.enabled,
+                 status = excluded.status,
+                 last_run = excluded.last_run,
+                 next_run = excluded.next_run,
+                 run_count = excluded.run_count,
+                 max_runs = excluded.max_runs,
+                 created_at = excluded.created_at,
+                 session_id = excluded.session_id,
+                 description = excluded.description,
+                 tags = excluded.tags,
+                 timeout = excluded.timeout,
+                 model_profile = excluded.model_profile,
+                 extra = excluded.extra,
+                 claimed_by = NULL,
+                 claimed_until = NULL""",
+            (
+                job.id,
+                job.name,
+                job.prompt,
+                job.schedule,
+                job.schedule_type.value,
+                job.cwd,
+                1 if job.enabled else 0,
+                job.status.value,
+                job.last_run,
+                job.next_run,
+                job.run_count,
+                job.max_runs,
+                job.created_at,
+                job.session_id,
+                job.description,
+                json.dumps(job.tags, ensure_ascii=False),
+                job.timeout,
+                job.model_profile,
+                json.dumps(job.extra, ensure_ascii=False),
+            ),
+        )
+        conn.commit()
+
+    def get_schedule(self, job_id: str) -> dict[str, Any] | None:
+        """Return a single schedule as a dict, or ``None`` if not found."""
+        conn = self._conn_or_create()
+        row = conn.execute(
+            "SELECT * FROM schedules WHERE id = ?", (job_id,)
+        ).fetchone()
+        if not row:
+            return None
+        cols = [d[0] for d in conn.execute("SELECT * FROM schedules LIMIT 0").description]
+        return _row_to_schedule(row, cols)
+
+    def resolve_schedule(self, selector: str) -> dict[str, Any] | None:
+        """Resolve an exact job ID or one unambiguous ID prefix."""
+        value = str(selector or "").strip()
+        if not value:
+            return None
+        exact = self.get_schedule(value)
+        if exact is not None:
+            return exact
+        conn = self._conn_or_create()
+        rows = conn.execute(
+            "SELECT * FROM schedules WHERE id LIKE ? ORDER BY created_at DESC LIMIT 2",
+            (f"{value}%",),
+        ).fetchall()
+        if len(rows) != 1:
+            return None
+        cols = [d[0] for d in conn.execute("SELECT * FROM schedules LIMIT 0").description]
+        return _row_to_schedule(rows[0], cols)
+
+    def list_schedules(
+        self,
+        *,
+        status: str | None = None,
+        schedule_type: str | None = None,
+        enabled: bool | None = None,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        """List schedules with optional filtering, most recently created first."""
+        conn = self._conn_or_create()
+        clauses: list[str] = []
+        params: list[Any] = []
+        if status is not None:
+            clauses.append("status = ?")
+            params.append(status)
+        if schedule_type is not None:
+            clauses.append("schedule_type = ?")
+            params.append(schedule_type)
+        if enabled is not None:
+            clauses.append("enabled = ?")
+            params.append(1 if enabled else 0)
+
+        where = ""
+        if clauses:
+            where = "WHERE " + " AND ".join(clauses)
+
+        rows = conn.execute(
+            f"SELECT * FROM schedules {where} ORDER BY created_at DESC LIMIT ?",
+            (*params, limit),
+        ).fetchall()
+        cols = [d[0] for d in conn.execute("SELECT * FROM schedules LIMIT 0").description]
+        return [_row_to_schedule(r, cols) for r in rows]
+
+    def delete_schedule(self, job_id: str) -> None:
+        """Delete a schedule and all its associated runs (cascade)."""
+        conn = self._conn_or_create()
+        # Delete runs first (defensive — FK ON DELETE CASCADE should handle it)
+        conn.execute("DELETE FROM schedule_runs WHERE job_id = ?", (job_id,))
+        conn.execute("DELETE FROM schedules WHERE id = ?", (job_id,))
+        conn.commit()
+
+    def claim_due_schedules(
+        self,
+        owner: str,
+        *,
+        now_iso: str | None = None,
+        limit: int = 1,
+        lease_seconds: int = 660,
+    ) -> list[dict[str, Any]]:
+        """Atomically lease due jobs to one scheduler instance.
+
+        The lease is stored in SQLite so multiple CrabCode processes can share
+        the default schedule database without executing the same occurrence.
+        """
+        conn = self._conn_or_create()
+        now = now_iso or datetime.now(timezone.utc).isoformat()
+        lease_until = (
+            datetime.fromisoformat(now) + timedelta(seconds=max(1, lease_seconds))
+        ).isoformat()
+        safe_limit = max(1, int(limit))
+        try:
+            conn.execute("BEGIN IMMEDIATE")
+            rows = conn.execute(
+                """SELECT * FROM schedules
+                   WHERE status = 'active'
+                     AND enabled = 1
+                     AND next_run IS NOT NULL
+                     AND next_run <= ?
+                     AND (max_runs IS NULL OR run_count < max_runs)
+                     AND (claimed_until IS NULL OR claimed_until <= ?)
+                   ORDER BY next_run ASC
+                   LIMIT ?""",
+                (now, now, safe_limit),
+            ).fetchall()
+            cols = [
+                d[0]
+                for d in conn.execute("SELECT * FROM schedules LIMIT 0").description
+            ]
+            claimed: list[dict[str, Any]] = []
+            for row in rows:
+                data = _row_to_schedule(row, cols)
+                changed = conn.execute(
+                    """UPDATE schedules
+                       SET claimed_by = ?, claimed_until = ?
+                       WHERE id = ?
+                         AND status = 'active'
+                         AND enabled = 1
+                         AND next_run IS NOT NULL
+                         AND next_run <= ?
+                         AND (claimed_until IS NULL OR claimed_until <= ?)""",
+                    (owner, lease_until, data["id"], now, now),
+                ).rowcount
+                if changed:
+                    data["claimed_by"] = owner
+                    data["claimed_until"] = lease_until
+                    claimed.append(data)
+            conn.commit()
+            return claimed
+        except BaseException:
+            conn.rollback()
+            raise
+
+    def release_claim(self, job_id: str, owner: str) -> None:
+        """Release a lease only when it is still owned by *owner*."""
+        conn = self._conn_or_create()
+        conn.execute(
+            """UPDATE schedules
+               SET claimed_by = NULL, claimed_until = NULL
+               WHERE id = ? AND claimed_by = ?""",
+            (job_id, owner),
+        )
+        conn.commit()
+
+    def claim_schedule(
+        self,
+        job_id: str,
+        owner: str,
+        *,
+        lease_seconds: int,
+        include_terminal: bool = False,
+    ) -> bool:
+        """Lease a job for a manual run, optionally including terminal jobs."""
+        conn = self._conn_or_create()
+        now = datetime.now(timezone.utc)
+        status_clause = (
+            "status IN ('active', 'error', 'disabled')"
+            if include_terminal
+            else "status = 'active' AND enabled = 1"
+        )
+        changed = conn.execute(
+            f"""UPDATE schedules
+               SET claimed_by = ?, claimed_until = ?
+               WHERE id = ?
+                 AND {status_clause}
+                 AND (claimed_until IS NULL OR claimed_until <= ?)""",
+            (
+                owner,
+                (now + timedelta(seconds=max(1, lease_seconds))).isoformat(),
+                job_id,
+                now.isoformat(),
+            ),
+        ).rowcount
+        conn.commit()
+        return bool(changed)
+
+    def renew_claim(self, job_id: str, owner: str, lease_seconds: int) -> bool:
+        """Extend an existing lease for a bounded execution."""
+        conn = self._conn_or_create()
+        claimed_until = (
+            datetime.now(timezone.utc) + timedelta(seconds=max(1, lease_seconds))
+        ).isoformat()
+        changed = conn.execute(
+            """UPDATE schedules
+               SET claimed_until = ?
+               WHERE id = ? AND claimed_by = ?""",
+            (claimed_until, job_id, owner),
+        ).rowcount
+        conn.commit()
+        return bool(changed)
+
+    def release_all_claims(self, owner: str) -> None:
+        conn = self._conn_or_create()
+        conn.execute(
+            """UPDATE schedules
+               SET claimed_by = NULL, claimed_until = NULL
+               WHERE claimed_by = ?""",
+            (owner,),
+        )
+        conn.commit()
+
+    # ------------------------------------------------------------------
+    # JobRun operations
+    # ------------------------------------------------------------------
+
+    def record_run(self, run: JobRun) -> None:
+        """Persist a :class:`JobRun` record (insert only — runs are immutable)."""
+        conn = self._conn_or_create()
+        conn.execute(
+            """INSERT INTO schedule_runs
+               (id, job_id, status, started_at, finished_at, duration_seconds,
+                session_id, exit_code, error_message, result_summary,
+                tokens_used, created_at)
+               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+            (
+                run.id,
+                run.job_id,
+                run.status.value,
+                run.started_at,
+                run.finished_at,
+                run.duration_seconds,
+                run.session_id,
+                run.exit_code,
+                run.error_message,
+                run.result_summary,
+                run.tokens_used,
+                run.created_at,
+            ),
+        )
+        conn.commit()
+
+    def update_run(self, run: JobRun) -> None:
+        """Update an existing :class:`JobRun` record (e.g. status transition)."""
+        conn = self._conn_or_create()
+        conn.execute(
+            """UPDATE schedule_runs
+               SET status = ?, started_at = ?, finished_at = ?,
+                   duration_seconds = ?, session_id = ?, exit_code = ?,
+                   error_message = ?, result_summary = ?, tokens_used = ?
+               WHERE id = ?""",
+            (
+                run.status.value,
+                run.started_at,
+                run.finished_at,
+                run.duration_seconds,
+                run.session_id,
+                run.exit_code,
+                run.error_message,
+                run.result_summary,
+                run.tokens_used,
+                run.id,
+            ),
+        )
+        conn.commit()
+
+    def list_runs(
+        self,
+        job_id: str,
+        *,
+        status: str | None = None,
+        limit: int = 50,
+    ) -> list[dict[str, Any]]:
+        """List runs for a given job, newest first."""
+        conn = self._conn_or_create()
+        if status is not None:
+            rows = conn.execute(
+                "SELECT * FROM schedule_runs WHERE job_id = ? AND status = ? "
+                "ORDER BY created_at DESC LIMIT ?",
+                (job_id, status, limit),
+            ).fetchall()
+        else:
+            rows = conn.execute(
+                "SELECT * FROM schedule_runs WHERE job_id = ? "
+                "ORDER BY created_at DESC LIMIT ?",
+                (job_id, limit),
+            ).fetchall()
+        cols = [d[0] for d in conn.execute("SELECT * FROM schedule_runs LIMIT 0").description]
+        return [_row_to_run(r, cols) for r in rows]
+
+    def get_run(self, run_id: str) -> dict[str, Any] | None:
+        """Return a single run as a dict, or ``None`` if not found."""
+        conn = self._conn_or_create()
+        row = conn.execute(
+            "SELECT * FROM schedule_runs WHERE id = ?", (run_id,)
+        ).fetchone()
+        if not row:
+            return None
+        cols = [d[0] for d in conn.execute("SELECT * FROM schedule_runs LIMIT 0").description]
+        return _row_to_run(row, cols)
+
+    # ------------------------------------------------------------------
+    # Restart recovery
+    # ------------------------------------------------------------------
+
+    def get_pending_schedules(self, now_iso: str | None = None) -> list[dict[str, Any]]:
+        """Return schedules that are due for execution (for restart recovery).
+
+        A schedule is *pending* when all of the following hold:
+        - ``status`` is ``active``
+        - ``enabled`` is true
+        - ``next_run`` is not ``NULL`` and ≤ *now*
+        - ``max_runs`` has not been reached (or is unlimited)
+
+        This is the primary method used by the scheduler engine on startup
+        to discover jobs that fired while the process was down.
+        """
+        conn = self._conn_or_create()
+        now = now_iso or datetime.now(timezone.utc).isoformat()
+
+        rows = conn.execute(
+            """SELECT * FROM schedules
+               WHERE status = 'active'
+                 AND enabled = 1
+                 AND next_run IS NOT NULL
+                 AND next_run <= ?
+                 AND (max_runs IS NULL OR run_count < max_runs)
+               ORDER BY next_run ASC""",
+            (now,),
+        ).fetchall()
+        cols = [d[0] for d in conn.execute("SELECT * FROM schedules LIMIT 0").description]
+        return [_row_to_schedule(r, cols) for r in rows]
+
+    def get_next_wakeup(self, now_iso: str | None = None) -> str | None:
+        """Return when a runnable job next needs polling.
+
+        A live lease delays the wakeup until it expires; otherwise ``next_run``
+        is used.  This avoids both 30-second jitter for near-future jobs and a
+        busy loop around jobs currently owned by another process.
+        """
+        conn = self._conn_or_create()
+        now = now_iso or datetime.now(timezone.utc).isoformat()
+        rows = conn.execute(
+            """SELECT next_run, claimed_until FROM schedules
+               WHERE status = 'active'
+                 AND enabled = 1
+                 AND next_run IS NOT NULL
+                 AND (max_runs IS NULL OR run_count < max_runs)"""
+        ).fetchall()
+        candidates = [
+            claimed_until
+            if claimed_until is not None and claimed_until > now
+            else next_run
+            for next_run, claimed_until in rows
+        ]
+        return min(candidates) if candidates else None
+
+    def get_stale_running_runs(self, now_iso: str | None = None) -> list[dict[str, Any]]:
+        """Return runs that are still marked as ``running`` (orphaned after a crash).
+
+        These should be transitioned to ``failed`` on startup.
+        """
+        conn = self._conn_or_create()
+        now = now_iso or datetime.now(timezone.utc).isoformat()
+        rows = conn.execute(
+            """SELECT schedule_runs.*
+               FROM schedule_runs
+               LEFT JOIN schedules ON schedules.id = schedule_runs.job_id
+               WHERE schedule_runs.status = 'running'
+                 AND (
+                   schedules.id IS NULL
+                   OR schedules.claimed_until IS NULL
+                   OR schedules.claimed_until <= ?
+                 )""",
+            (now,),
+        ).fetchall()
+        cols = [d[0] for d in conn.execute("SELECT * FROM schedule_runs LIMIT 0").description]
+        return [_row_to_run(r, cols) for r in rows]