Integrate task runner into python tools
This commit is contained in:
ajp_anton
@@ -0,0 +1,271 @@
|
||||
"""SQLite persistence for runs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
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
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
"""
|
||||
)
|
||||
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",
|
||||
) -> tuple[Run, bool]:
|
||||
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
|
||||
) VALUES (?, ?, ?, ?, 'queued', ?)
|
||||
""",
|
||||
(target_kind, target_path, queue_group, trigger, now()),
|
||||
)
|
||||
return self.get(cursor.lastrowid, connection=connection), True
|
||||
|
||||
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
|
||||
) 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))
|
||||
Reference in New Issue
Block a user