384 lines
15 KiB
Python
384 lines
15 KiB
Python
"""Durable, private SQLite state for the autonomous-release Candidate."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sqlite3
|
|
from typing import Any, Iterator, Mapping, Sequence
|
|
|
|
from .errors import StateConflict
|
|
from .models import StepState, TaskState, canonical_json
|
|
|
|
|
|
SCHEMA_VERSION = 2
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class StoredStep:
|
|
task_id: str
|
|
step_id: str
|
|
state: StepState
|
|
attempt: int
|
|
max_attempts: int
|
|
worker_id: str | None
|
|
session_id: str | None
|
|
lease_expires_at: datetime | None
|
|
result_digest: str | None
|
|
updated_at: datetime
|
|
last_error: str | None
|
|
|
|
|
|
class StateStore:
|
|
"""SQLite store with WAL, full sync, schema checks and serialized CAS writes."""
|
|
|
|
def __init__(self, path: str | Path) -> None:
|
|
self.path = Path(path)
|
|
self._prepare_path()
|
|
self._initialize()
|
|
os.chmod(self.path, 0o600)
|
|
|
|
def _prepare_path(self) -> None:
|
|
if self.path.exists() and self.path.is_symlink():
|
|
raise StateConflict("STATE_PATH_SYMLINK", "state database must not be a symlink")
|
|
parent = self.path.parent
|
|
parent.mkdir(parents=True, exist_ok=True, mode=0o750)
|
|
if parent.is_symlink():
|
|
raise StateConflict("STATE_DIRECTORY_SYMLINK", "state directory must not be a symlink")
|
|
current_mode = parent.stat().st_mode & 0o777
|
|
os.chmod(parent, current_mode & ~0o027)
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
if self.path.exists() and self.path.is_symlink():
|
|
raise StateConflict("STATE_PATH_SYMLINK", "state database must not be a symlink")
|
|
connection = sqlite3.connect(self.path, timeout=30.0)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA foreign_keys=ON")
|
|
connection.execute("PRAGMA journal_mode=WAL")
|
|
connection.execute("PRAGMA synchronous=FULL")
|
|
connection.execute("PRAGMA busy_timeout=30000")
|
|
return connection
|
|
|
|
@contextmanager
|
|
def transaction(self) -> Iterator[sqlite3.Connection]:
|
|
connection = self._connect()
|
|
try:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
yield connection
|
|
connection.commit()
|
|
except Exception:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|
|
|
|
def _initialize(self) -> None:
|
|
with self.transaction() as connection:
|
|
connection.execute(
|
|
"CREATE TABLE IF NOT EXISTS metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL)"
|
|
)
|
|
row = connection.execute(
|
|
"SELECT value FROM metadata WHERE key='schema_version'"
|
|
).fetchone()
|
|
version = int(row["value"]) if row is not None else None
|
|
if version not in (None, 1, SCHEMA_VERSION):
|
|
raise StateConflict(
|
|
"SCHEMA_VERSION_UNSUPPORTED",
|
|
"state schema is not supported",
|
|
{"found": version, "expected": SCHEMA_VERSION},
|
|
)
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS tasks (
|
|
task_id TEXT PRIMARY KEY,
|
|
plan_digest TEXT NOT NULL,
|
|
plan_json TEXT NOT NULL,
|
|
state TEXT NOT NULL,
|
|
revision INTEGER NOT NULL DEFAULT 0,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS steps (
|
|
task_id TEXT NOT NULL,
|
|
step_id TEXT NOT NULL,
|
|
state TEXT NOT NULL,
|
|
attempt INTEGER NOT NULL DEFAULT 0,
|
|
max_attempts INTEGER NOT NULL DEFAULT 1,
|
|
worker_id TEXT,
|
|
session_id TEXT,
|
|
lease_expires_at TEXT,
|
|
result_digest TEXT,
|
|
last_error TEXT,
|
|
updated_at TEXT NOT NULL,
|
|
PRIMARY KEY(task_id, step_id),
|
|
FOREIGN KEY(task_id) REFERENCES tasks(task_id) ON DELETE CASCADE
|
|
);
|
|
CREATE TABLE IF NOT EXISTS events (
|
|
sequence INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
task_id TEXT NOT NULL,
|
|
event_type TEXT NOT NULL,
|
|
payload_json TEXT NOT NULL,
|
|
previous_digest TEXT,
|
|
digest TEXT NOT NULL UNIQUE,
|
|
created_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS consumed_nonces (
|
|
scope TEXT NOT NULL,
|
|
nonce TEXT NOT NULL,
|
|
consumed_at TEXT NOT NULL,
|
|
PRIMARY KEY(scope, nonce)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS credentials (
|
|
credential_id TEXT PRIMARY KEY,
|
|
task_id TEXT NOT NULL,
|
|
reference TEXT NOT NULL UNIQUE,
|
|
state TEXT NOT NULL,
|
|
metadata_json TEXT NOT NULL,
|
|
issued_at TEXT NOT NULL,
|
|
expires_at TEXT NOT NULL,
|
|
revoked_at TEXT
|
|
);
|
|
CREATE TABLE IF NOT EXISTS idempotency (
|
|
namespace TEXT NOT NULL,
|
|
operation_id TEXT NOT NULL,
|
|
result_json TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
PRIMARY KEY(namespace, operation_id)
|
|
);
|
|
"""
|
|
)
|
|
if version == 1:
|
|
existing = {
|
|
item["name"] for item in connection.execute("PRAGMA table_info(steps)")
|
|
}
|
|
additions = {
|
|
"max_attempts": "INTEGER NOT NULL DEFAULT 1",
|
|
"lease_expires_at": "TEXT",
|
|
"result_digest": "TEXT",
|
|
}
|
|
for column, declaration in additions.items():
|
|
if column not in existing:
|
|
connection.execute(f"ALTER TABLE steps ADD COLUMN {column} {declaration}")
|
|
connection.execute(
|
|
"INSERT INTO metadata(key, value) VALUES('schema_version', ?) "
|
|
"ON CONFLICT(key) DO UPDATE SET value=excluded.value",
|
|
(str(SCHEMA_VERSION),),
|
|
)
|
|
|
|
@staticmethod
|
|
def _now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
def create_task(
|
|
self,
|
|
task_id: str,
|
|
plan_digest: str,
|
|
plan: Mapping[str, Any],
|
|
step_ids: tuple[str, ...],
|
|
step_limits: Mapping[str, int] | None = None,
|
|
) -> None:
|
|
now = self._now()
|
|
limits = step_limits or {}
|
|
with self.transaction() as connection:
|
|
connection.execute(
|
|
"INSERT INTO tasks(task_id, plan_digest, plan_json, state, created_at, updated_at) "
|
|
"VALUES(?, ?, ?, ?, ?, ?)",
|
|
(task_id, plan_digest, canonical_json(plan), TaskState.PLANNED.value, now, now),
|
|
)
|
|
connection.executemany(
|
|
"INSERT INTO steps(task_id, step_id, state, max_attempts, updated_at) "
|
|
"VALUES(?, ?, ?, ?, ?)",
|
|
[
|
|
(task_id, step_id, StepState.PENDING.value, int(limits.get(step_id, 1)), now)
|
|
for step_id in step_ids
|
|
],
|
|
)
|
|
|
|
def get_task_state(self, task_id: str) -> tuple[TaskState, int]:
|
|
with self._connect() as connection:
|
|
row = connection.execute(
|
|
"SELECT state, revision FROM tasks WHERE task_id=?", (task_id,)
|
|
).fetchone()
|
|
if row is None:
|
|
raise KeyError(task_id)
|
|
return TaskState(row["state"]), int(row["revision"])
|
|
|
|
def compare_and_set_task_state(
|
|
self, task_id: str, expected_revision: int, new_state: TaskState
|
|
) -> int:
|
|
with self.transaction() as connection:
|
|
cursor = connection.execute(
|
|
"UPDATE tasks SET state=?, revision=revision+1, updated_at=? "
|
|
"WHERE task_id=? AND revision=?",
|
|
(new_state.value, self._now(), task_id, expected_revision),
|
|
)
|
|
if cursor.rowcount != 1:
|
|
raise StateConflict(
|
|
"TASK_REVISION_CONFLICT",
|
|
"task state changed concurrently",
|
|
{"task_id": task_id, "expected_revision": expected_revision},
|
|
)
|
|
row = connection.execute(
|
|
"SELECT revision FROM tasks WHERE task_id=?", (task_id,)
|
|
).fetchone()
|
|
return int(row["revision"])
|
|
|
|
def transition_step(
|
|
self,
|
|
task_id: str,
|
|
step_id: str,
|
|
expected: Sequence[StepState],
|
|
state: StepState,
|
|
*,
|
|
worker_id: str | None = None,
|
|
session_id: str | None = None,
|
|
lease_expires_at: datetime | None = None,
|
|
result_digest: str | None = None,
|
|
last_error: str | None = None,
|
|
increment_attempt: bool = False,
|
|
clear_assignment: bool = False,
|
|
) -> StoredStep:
|
|
with self.transaction() as connection:
|
|
row = connection.execute(
|
|
"SELECT * FROM steps WHERE task_id=? AND step_id=?", (task_id, step_id)
|
|
).fetchone()
|
|
if row is None:
|
|
raise KeyError(f"{task_id}:{step_id}")
|
|
if StepState(row["state"]) not in set(expected):
|
|
raise StateConflict(
|
|
"STEP_STATE_CONFLICT",
|
|
"step state changed or transition is invalid",
|
|
{"task_id": task_id, "step_id": step_id, "state": row["state"]},
|
|
)
|
|
if increment_attempt and int(row["attempt"]) >= int(row["max_attempts"]):
|
|
raise StateConflict(
|
|
"STEP_ATTEMPTS_EXHAUSTED",
|
|
"step has no remaining attempts",
|
|
{"task_id": task_id, "step_id": step_id},
|
|
)
|
|
next_worker = None if clear_assignment else (worker_id or row["worker_id"])
|
|
next_session = None if clear_assignment else (session_id or row["session_id"])
|
|
connection.execute(
|
|
"UPDATE steps SET state=?, attempt=attempt+?, worker_id=?, session_id=?, "
|
|
"lease_expires_at=?, result_digest=COALESCE(?, result_digest), last_error=?, updated_at=? "
|
|
"WHERE task_id=? AND step_id=?",
|
|
(
|
|
state.value,
|
|
1 if increment_attempt else 0,
|
|
next_worker,
|
|
next_session,
|
|
lease_expires_at.isoformat() if lease_expires_at else None,
|
|
result_digest,
|
|
last_error,
|
|
self._now(),
|
|
task_id,
|
|
step_id,
|
|
),
|
|
)
|
|
updated = connection.execute(
|
|
"SELECT * FROM steps WHERE task_id=? AND step_id=?", (task_id, step_id)
|
|
).fetchone()
|
|
return self._step_from_row(updated)
|
|
|
|
def set_step_state(
|
|
self,
|
|
task_id: str,
|
|
step_id: str,
|
|
state: StepState,
|
|
**kwargs: Any,
|
|
) -> StoredStep:
|
|
current = self.list_steps(task_id)[step_id]
|
|
return self.transition_step(task_id, step_id, (current.state,), state, **kwargs)
|
|
|
|
def list_steps(self, task_id: str) -> dict[str, StoredStep]:
|
|
with self._connect() as connection:
|
|
rows = connection.execute(
|
|
"SELECT * FROM steps WHERE task_id=? ORDER BY step_id", (task_id,)
|
|
).fetchall()
|
|
return {row["step_id"]: self._step_from_row(row) for row in rows}
|
|
|
|
@staticmethod
|
|
def _step_from_row(row: sqlite3.Row) -> StoredStep:
|
|
return StoredStep(
|
|
task_id=row["task_id"],
|
|
step_id=row["step_id"],
|
|
state=StepState(row["state"]),
|
|
attempt=int(row["attempt"]),
|
|
max_attempts=int(row["max_attempts"]),
|
|
worker_id=row["worker_id"],
|
|
session_id=row["session_id"],
|
|
lease_expires_at=(
|
|
datetime.fromisoformat(row["lease_expires_at"])
|
|
if row["lease_expires_at"]
|
|
else None
|
|
),
|
|
result_digest=row["result_digest"],
|
|
updated_at=datetime.fromisoformat(row["updated_at"]),
|
|
last_error=row["last_error"],
|
|
)
|
|
|
|
def consume_nonce(self, scope: str, nonce: str) -> None:
|
|
try:
|
|
with self.transaction() as connection:
|
|
connection.execute(
|
|
"INSERT INTO consumed_nonces(scope, nonce, consumed_at) VALUES(?, ?, ?)",
|
|
(scope, nonce, self._now()),
|
|
)
|
|
except sqlite3.IntegrityError as error:
|
|
raise StateConflict(
|
|
"NONCE_REPLAY", "nonce was already consumed", {"scope": scope, "nonce": nonce}
|
|
) from error
|
|
|
|
def append_event(
|
|
self,
|
|
task_id: str,
|
|
event_type: str,
|
|
payload: Mapping[str, Any],
|
|
previous_digest: str | None,
|
|
event_digest: str,
|
|
created_at: datetime,
|
|
) -> int:
|
|
with self.transaction() as connection:
|
|
cursor = connection.execute(
|
|
"INSERT INTO events(task_id, event_type, payload_json, previous_digest, digest, created_at) "
|
|
"VALUES(?, ?, ?, ?, ?, ?)",
|
|
(
|
|
task_id,
|
|
event_type,
|
|
canonical_json(payload),
|
|
previous_digest,
|
|
event_digest,
|
|
created_at.astimezone(timezone.utc).isoformat(),
|
|
),
|
|
)
|
|
return int(cursor.lastrowid)
|
|
|
|
def list_events(self, task_id: str) -> list[Mapping[str, Any]]:
|
|
with self._connect() as connection:
|
|
rows = connection.execute(
|
|
"SELECT * FROM events WHERE task_id=? ORDER BY sequence", (task_id,)
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def put_idempotent_result(
|
|
self, namespace: str, operation_id: str, result: Mapping[str, Any]
|
|
) -> Mapping[str, Any]:
|
|
encoded = canonical_json(result)
|
|
with self.transaction() as connection:
|
|
existing = connection.execute(
|
|
"SELECT result_json FROM idempotency WHERE namespace=? AND operation_id=?",
|
|
(namespace, operation_id),
|
|
).fetchone()
|
|
if existing is not None:
|
|
return json.loads(existing["result_json"])
|
|
connection.execute(
|
|
"INSERT INTO idempotency(namespace, operation_id, result_json, created_at) "
|
|
"VALUES(?, ?, ?, ?)",
|
|
(namespace, operation_id, encoded, self._now()),
|
|
)
|
|
return dict(result)
|