| 1 | """Durable journal backing the server package's in-memory work queues.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import json |
| 6 | import sqlite3 |
| 7 | import threading |
| 8 | from pathlib import Path |
| 9 | from typing import Any |
| 10 | |
| 11 | |
| 12 | class JobJournal: |
| 13 | """Persist queue snapshots without using SQLite as the scheduling queue.""" |
| 14 | |
| 15 | def __init__(self, path: Path) -> None: |
| 16 | self.path = path.expanduser().resolve() |
| 17 | self._lock = threading.RLock() |
| 18 | self._initialized = False |
| 19 | |
| 20 | def connect(self) -> sqlite3.Connection: |
| 21 | connection = sqlite3.connect(self.path, timeout=30) |
| 22 | connection.row_factory = sqlite3.Row |
| 23 | connection.execute("PRAGMA busy_timeout = 30000") |
| 24 | return connection |
| 25 | |
| 26 | def initialize(self) -> None: |
| 27 | with self._lock: |
| 28 | if self._initialized: |
| 29 | return |
| 30 | self.path.parent.mkdir(parents=True, exist_ok=True) |
| 31 | with self.connect() as connection: |
| 32 | connection.execute("PRAGMA journal_mode = WAL") |
| 33 | connection.execute( |
| 34 | """ |
| 35 | CREATE TABLE IF NOT EXISTS job_journal ( |
| 36 | kind TEXT NOT NULL, |
| 37 | version_id TEXT NOT NULL, |
| 38 | work_id TEXT NOT NULL, |
| 39 | agent_job_id TEXT, |
| 40 | job_json TEXT NOT NULL, |
| 41 | created_at TEXT NOT NULL, |
| 42 | updated_at TEXT NOT NULL, |
| 43 | PRIMARY KEY(kind, version_id) |
| 44 | ) |
| 45 | """ |
| 46 | ) |
| 47 | # job_id is generated by Director and is the idempotency key across |
| 48 | # retries. It is intentionally independent of the logical work_id. |
| 49 | connection.execute("DROP INDEX IF EXISTS idx_job_journal_idempotency") |
| 50 | connection.execute( |
| 51 | """ |
| 52 | CREATE UNIQUE INDEX IF NOT EXISTS idx_job_journal_idempotency |
| 53 | ON job_journal(kind, agent_job_id) |
| 54 | WHERE agent_job_id IS NOT NULL AND agent_job_id != '' |
| 55 | """ |
| 56 | ) |
| 57 | connection.execute( |
| 58 | "CREATE INDEX IF NOT EXISTS idx_job_journal_updated " |
| 59 | "ON job_journal(kind, updated_at)" |
| 60 | ) |
| 61 | self._initialized = True |
| 62 | |
| 63 | def load(self, kind: str) -> list[dict[str, Any]]: |
| 64 | with self._lock, self.connect() as connection: |
| 65 | rows = connection.execute( |
| 66 | "SELECT job_json FROM job_journal WHERE kind = ? " |
| 67 | "ORDER BY created_at, version_id", |
| 68 | (kind,), |
| 69 | ).fetchall() |
| 70 | return [json.loads(str(row["job_json"])) for row in rows] |
| 71 | |
| 72 | def get(self, kind: str, version_id: str) -> dict[str, Any] | None: |
| 73 | with self._lock, self.connect() as connection: |
| 74 | row = connection.execute( |
| 75 | "SELECT job_json FROM job_journal WHERE kind = ? AND version_id = ?", |
| 76 | (kind, version_id), |
| 77 | ).fetchone() |
| 78 | return json.loads(str(row["job_json"])) if row else None |
| 79 | |
| 80 | def get_by_agent_job_id(self, kind: str, agent_job_id: str) -> dict[str, Any] | None: |
| 81 | with self._lock, self.connect() as connection: |
| 82 | row = connection.execute( |
| 83 | "SELECT job_json FROM job_journal WHERE kind = ? AND agent_job_id = ?", |
| 84 | (kind, agent_job_id), |
| 85 | ).fetchone() |
| 86 | return json.loads(str(row["job_json"])) if row else None |
| 87 | |
| 88 | def save(self, kind: str, job: dict[str, Any]) -> None: |
| 89 | payload = json.dumps(job, ensure_ascii=False, separators=(",", ":")) |
| 90 | with self._lock, self.connect() as connection: |
| 91 | connection.execute( |
| 92 | """ |
| 93 | INSERT INTO job_journal ( |
| 94 | kind, version_id, work_id, agent_job_id, job_json, |
| 95 | created_at, updated_at |
| 96 | ) VALUES (?, ?, ?, ?, ?, ?, ?) |
| 97 | ON CONFLICT(kind, version_id) DO UPDATE SET |
| 98 | work_id = excluded.work_id, |
| 99 | agent_job_id = excluded.agent_job_id, |
| 100 | job_json = excluded.job_json, |
| 101 | updated_at = excluded.updated_at |
| 102 | """, |
| 103 | ( |
| 104 | kind, |
| 105 | job["version_id"], |
| 106 | job["work_id"], |
| 107 | job.get("agent_job_id"), |
| 108 | payload, |
| 109 | job["created_at"], |
| 110 | job["updated_at"], |
| 111 | ), |
| 112 | ) |
| 113 | |
| 114 | def delete(self, kind: str, version_id: str) -> None: |
| 115 | with self._lock, self.connect() as connection: |
| 116 | connection.execute( |
| 117 | "DELETE FROM job_journal WHERE kind = ? AND version_id = ?", |
| 118 | (kind, version_id), |
| 119 | ) |
| 120 |