301 lines
10 KiB
Python
301 lines
10 KiB
Python
"""SQLite persistence for runs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
import json
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
|
|
ACTIVE_STATUSES = ("queued", "running")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Run:
|
|
id: int
|
|
target_kind: str
|
|
target_path: str
|
|
queue_group: str
|
|
trigger: str
|
|
status: str
|
|
created_at: str
|
|
started_at: str | None
|
|
finished_at: str | None
|
|
exit_code: int | None
|
|
cancel_requested_at: str | None
|
|
parent_run_id: int | None
|
|
log_path: str | None
|
|
message: str | None
|
|
input_json: str
|
|
|
|
|
|
def now() -> str:
|
|
return datetime.now(UTC).isoformat()
|
|
|
|
|
|
class RunStore:
|
|
def __init__(self, database_path: Path) -> None:
|
|
self.database_path = database_path
|
|
|
|
def initialize(self) -> None:
|
|
self.database_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with self._connect() as connection:
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS runs (
|
|
id INTEGER PRIMARY KEY,
|
|
target_kind TEXT NOT NULL CHECK(target_kind IN ('task', 'group')),
|
|
target_path TEXT NOT NULL,
|
|
queue_group TEXT NOT NULL,
|
|
trigger TEXT NOT NULL CHECK(trigger IN ('manual', 'schedule')),
|
|
status TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
started_at TEXT,
|
|
finished_at TEXT,
|
|
exit_code INTEGER,
|
|
cancel_requested_at TEXT,
|
|
parent_run_id INTEGER REFERENCES runs(id),
|
|
log_path TEXT,
|
|
message TEXT,
|
|
input_json TEXT NOT NULL DEFAULT '{}'
|
|
)
|
|
"""
|
|
)
|
|
columns = {
|
|
row["name"]
|
|
for row in connection.execute("PRAGMA table_info(runs)").fetchall()
|
|
}
|
|
if "input_json" not in columns:
|
|
connection.execute(
|
|
"ALTER TABLE runs ADD COLUMN input_json TEXT NOT NULL DEFAULT '{}'"
|
|
)
|
|
connection.execute(
|
|
"""
|
|
CREATE INDEX IF NOT EXISTS runs_queue_index
|
|
ON runs(queue_group, status, created_at)
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"""
|
|
CREATE INDEX IF NOT EXISTS runs_target_index
|
|
ON runs(target_kind, target_path, created_at DESC)
|
|
"""
|
|
)
|
|
|
|
def mark_interrupted(self) -> None:
|
|
with self._connect() as connection:
|
|
connection.execute(
|
|
"""
|
|
UPDATE runs
|
|
SET status = 'interrupted',
|
|
finished_at = ?,
|
|
message = 'Application restarted before this run completed.'
|
|
WHERE status IN ('queued', 'running')
|
|
""",
|
|
(now(),),
|
|
)
|
|
|
|
def enqueue(
|
|
self,
|
|
*,
|
|
target_kind: str,
|
|
target_path: str,
|
|
queue_group: str,
|
|
trigger: str = "manual",
|
|
input_data: Mapping[str, str | int | list[str]] | None = None,
|
|
) -> tuple[Run, bool]:
|
|
input_json = json.dumps(input_data or {}, sort_keys=True, separators=(",", ":"))
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
existing = connection.execute(
|
|
"""
|
|
SELECT * FROM runs
|
|
WHERE target_kind = ? AND target_path = ?
|
|
AND parent_run_id IS NULL
|
|
AND status IN ('queued', 'running')
|
|
ORDER BY id DESC
|
|
LIMIT 1
|
|
""",
|
|
(target_kind, target_path),
|
|
).fetchone()
|
|
if existing is not None:
|
|
return self._run(existing), False
|
|
cursor = connection.execute(
|
|
"""
|
|
INSERT INTO runs (
|
|
target_kind, target_path, queue_group, trigger, status, created_at, input_json
|
|
) VALUES (?, ?, ?, ?, 'queued', ?, ?)
|
|
""",
|
|
(target_kind, target_path, queue_group, trigger, now(), input_json),
|
|
)
|
|
return self.get(cursor.lastrowid, connection=connection), True
|
|
|
|
def set_input(
|
|
self,
|
|
run_id: int,
|
|
input_data: Mapping[str, str | int | list[str]],
|
|
) -> Run:
|
|
input_json = json.dumps(input_data, sort_keys=True, separators=(",", ":"))
|
|
with self._connect() as connection:
|
|
connection.execute(
|
|
"UPDATE runs SET input_json = ? WHERE id = ?",
|
|
(input_json, run_id),
|
|
)
|
|
run = self.get(run_id, connection=connection)
|
|
assert run is not None
|
|
return run
|
|
|
|
def claim_next(self, queue_group: str) -> Run | None:
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
row = connection.execute(
|
|
"""
|
|
SELECT * FROM runs
|
|
WHERE queue_group = ? AND parent_run_id IS NULL AND status = 'queued'
|
|
ORDER BY created_at, id
|
|
LIMIT 1
|
|
""",
|
|
(queue_group,),
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
run_id = row["id"]
|
|
connection.execute(
|
|
"UPDATE runs SET status = 'running', started_at = ? WHERE id = ?",
|
|
(now(), run_id),
|
|
)
|
|
return self.get(run_id, connection=connection)
|
|
|
|
def create_child(self, parent_run: Run, task_path: str) -> Run:
|
|
with self._connect() as connection:
|
|
cursor = connection.execute(
|
|
"""
|
|
INSERT INTO runs (
|
|
target_kind, target_path, queue_group, trigger, status, created_at,
|
|
started_at, parent_run_id, input_json
|
|
) VALUES ('task', ?, ?, ?, 'running', ?, ?, ?, '{}')
|
|
""",
|
|
(
|
|
task_path,
|
|
parent_run.queue_group,
|
|
parent_run.trigger,
|
|
now(),
|
|
now(),
|
|
parent_run.id,
|
|
),
|
|
)
|
|
return self.get(cursor.lastrowid, connection=connection)
|
|
|
|
def finish(
|
|
self,
|
|
run_id: int,
|
|
*,
|
|
status: str,
|
|
exit_code: int | None = None,
|
|
message: str | None = None,
|
|
log_path: str | None = None,
|
|
) -> Run:
|
|
with self._connect() as connection:
|
|
connection.execute(
|
|
"""
|
|
UPDATE runs
|
|
SET status = ?, exit_code = ?, finished_at = ?, message = ?,
|
|
log_path = COALESCE(?, log_path)
|
|
WHERE id = ?
|
|
""",
|
|
(status, exit_code, now(), message, log_path, run_id),
|
|
)
|
|
return self.get(run_id, connection=connection)
|
|
|
|
def request_cancel(self, run_id: int) -> Run | None:
|
|
with self._connect() as connection:
|
|
run = self.get(run_id, connection=connection)
|
|
if run is None or run.status not in ACTIVE_STATUSES:
|
|
return run
|
|
if run.status == "queued":
|
|
connection.execute(
|
|
"""
|
|
UPDATE runs
|
|
SET status = 'cancelled', cancel_requested_at = ?, finished_at = ?,
|
|
message = 'Cancelled before execution.'
|
|
WHERE id = ?
|
|
""",
|
|
(now(), now(), run_id),
|
|
)
|
|
elif run.cancel_requested_at is None:
|
|
connection.execute(
|
|
"UPDATE runs SET cancel_requested_at = ? WHERE id = ?",
|
|
(now(), run_id),
|
|
)
|
|
return self.get(run_id, connection=connection)
|
|
|
|
def is_cancel_requested(self, run_id: int) -> bool:
|
|
run = self.get(run_id)
|
|
return run is not None and run.cancel_requested_at is not None
|
|
|
|
def active_child(self, parent_run_id: int) -> Run | None:
|
|
with self._connect() as connection:
|
|
row = connection.execute(
|
|
"""
|
|
SELECT * FROM runs
|
|
WHERE parent_run_id = ? AND status = 'running'
|
|
ORDER BY id DESC
|
|
LIMIT 1
|
|
""",
|
|
(parent_run_id,),
|
|
).fetchone()
|
|
return self._run(row) if row else None
|
|
|
|
def get(self, run_id: int, *, connection: sqlite3.Connection | None = None) -> Run | None:
|
|
if connection is not None:
|
|
row = connection.execute("SELECT * FROM runs WHERE id = ?", (run_id,)).fetchone()
|
|
return self._run(row) if row else None
|
|
with self._connect() as owned_connection:
|
|
return self.get(run_id, connection=owned_connection)
|
|
|
|
def last_run(self, target_kind: str, target_path: str) -> Run | None:
|
|
with self._connect() as connection:
|
|
row = connection.execute(
|
|
"""
|
|
SELECT * FROM runs
|
|
WHERE target_kind = ? AND target_path = ? AND parent_run_id IS NULL
|
|
ORDER BY id DESC
|
|
LIMIT 1
|
|
""",
|
|
(target_kind, target_path),
|
|
).fetchone()
|
|
return self._run(row) if row else None
|
|
|
|
def history(self, target_kind: str, target_path: str) -> list[Run]:
|
|
with self._connect() as connection:
|
|
rows = connection.execute(
|
|
"""
|
|
SELECT * FROM runs
|
|
WHERE target_kind = ? AND target_path = ? AND parent_run_id IS NULL
|
|
ORDER BY id DESC
|
|
LIMIT 50
|
|
""",
|
|
(target_kind, target_path),
|
|
).fetchall()
|
|
return [self._run(row) for row in rows]
|
|
|
|
def children(self, parent_run_id: int) -> list[Run]:
|
|
with self._connect() as connection:
|
|
rows = connection.execute(
|
|
"SELECT * FROM runs WHERE parent_run_id = ? ORDER BY id",
|
|
(parent_run_id,),
|
|
).fetchall()
|
|
return [self._run(row) for row in rows]
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
connection = sqlite3.connect(self.database_path, timeout=30)
|
|
connection.row_factory = sqlite3.Row
|
|
return connection
|
|
|
|
@staticmethod
|
|
def _run(row: sqlite3.Row) -> Run:
|
|
return Run(**dict(row))
|