Files
archive-clients/src/archive_clients/state.py
T

463 lines
16 KiB
Python

"""Durable command inbox and client recovery cursors."""
from __future__ import annotations
import hashlib
import json
import os
import sqlite3
import stat
from dataclasses import dataclass
from pathlib import Path
SCHEMA_VERSION = 2
class CommandConflict(RuntimeError):
pass
class JobConflict(RuntimeError):
pass
class FileOperationConflict(RuntimeError):
pass
@dataclass(frozen=True)
class CommandAcceptance:
duplicate: bool
acknowledgement_json: str
class ClientStore:
def __init__(self, database: Path):
self.database = database
def initialize(self) -> None:
self.database.parent.mkdir(parents=True, exist_ok=True)
existed = self.database.exists()
if existed:
metadata = self.database.stat()
if not stat.S_ISREG(metadata.st_mode):
raise RuntimeError("client database is not a regular file")
if metadata.st_mode & (stat.S_IRWXG | stat.S_IRWXO):
raise RuntimeError("client database file permissions are unsafe")
with self._connect() as connection:
version = connection.execute("PRAGMA user_version").fetchone()[0]
if version > SCHEMA_VERSION:
raise RuntimeError(
f"client database schema {version} is newer than supported"
)
if version == 0:
connection.executescript(_SCHEMA_V1)
connection.execute("PRAGMA user_version = 1")
version = 1
if version == 1:
connection.executescript(_SCHEMA_V2)
connection.execute("PRAGMA user_version = 2")
if connection.execute("PRAGMA foreign_key_check").fetchall():
raise RuntimeError("client database foreign-key check failed")
if not existed:
os.chmod(self.database, 0o600)
def accept_command(
self, command_id: str, command_json: str, acknowledgement_json: str
) -> CommandAcceptance:
payload = _canonical(json.loads(command_json))
acknowledgement = _canonical(json.loads(acknowledgement_json))
digest = hashlib.sha256(payload.encode()).hexdigest()
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing = connection.execute(
"""
SELECT payload_sha256, command_json, acknowledgement_json
FROM commands WHERE command_id = ?
""",
(command_id,),
).fetchone()
if existing:
if existing["payload_sha256"] != digest or existing["command_json"] != payload:
raise CommandConflict("command ID was reused with different content")
return CommandAcceptance(True, existing["acknowledgement_json"])
connection.execute(
"""
INSERT INTO commands (
command_id, payload_sha256, command_json,
acknowledgement_json, state
) VALUES (?, ?, ?, ?, 'received')
""",
(command_id, digest, payload, acknowledgement),
)
return CommandAcceptance(False, acknowledgement)
def list_active_job_cursors(self) -> list[dict[str, object]]:
with self._connect() as connection:
rows = connection.execute(
"""
SELECT job_id, revision, last_event_sequence, state, committed
FROM jobs WHERE state NOT IN ('JOB_STATE_SUCCEEDED',
'JOB_STATE_FAILED', 'JOB_STATE_CANCELLED')
ORDER BY job_id
"""
).fetchall()
return [dict(row) for row in rows]
def list_accepted_commands(self) -> list[dict[str, object]]:
with self._connect() as connection:
rows = connection.execute(
"""
SELECT command_id, command_json, acknowledgement_json
FROM commands ORDER BY rowid
"""
).fetchall()
return [dict(row) for row in rows]
def job_snapshot_rows(
self, job_ids: list[str]
) -> list[dict[str, object]]:
if not job_ids:
return []
placeholders = ",".join("?" for _ in job_ids)
with self._connect() as connection:
rows = connection.execute(
f"""
SELECT job_id, definition_json, state, revision,
last_event_sequence, committed
FROM jobs WHERE job_id IN ({placeholders}) ORDER BY job_id
""",
job_ids,
).fetchall()
return [dict(row) for row in rows]
def save_job(
self,
job_id: str,
definition_json: str,
state: str,
revision: int,
last_event_sequence: int,
committed: bool,
) -> None:
definition = _canonical(json.loads(definition_json))
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing = connection.execute(
"""
SELECT definition_json, state, revision,
last_event_sequence, committed
FROM jobs WHERE job_id = ?
""",
(job_id,),
).fetchone()
if existing and existing["definition_json"] != definition:
raise JobConflict("job definition is immutable")
if existing and (
revision < existing["revision"]
or last_event_sequence < existing["last_event_sequence"]
or (existing["committed"] and not committed)
):
raise JobConflict("job cursor cannot move backwards")
if existing and (
revision == existing["revision"]
and last_event_sequence == existing["last_event_sequence"]
and (
state != existing["state"]
or int(committed) != existing["committed"]
)
):
raise JobConflict("equal job cursor has conflicting state")
connection.execute(
"""
INSERT INTO jobs (
job_id, definition_json, state, revision,
last_event_sequence, committed
) VALUES (?, ?, ?, ?, ?, ?)
ON CONFLICT(job_id) DO UPDATE SET
state = excluded.state,
revision = excluded.revision,
last_event_sequence = excluded.last_event_sequence,
committed = excluded.committed
""",
(
job_id, definition, state, revision,
last_event_sequence, int(committed),
),
)
def begin_file_operation(
self,
operation_id: str,
job_id: str,
intent_json: str,
) -> dict[str, object]:
intent = _canonical(json.loads(intent_json))
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing = connection.execute(
"SELECT * FROM file_journal WHERE operation_id = ?",
(operation_id,),
).fetchone()
if existing:
if (
existing["job_id"] != job_id
or existing["intent_json"] != intent
):
raise FileOperationConflict(
"file operation ID was reused with different intent"
)
return dict(existing)
connection.execute(
"""
INSERT INTO file_journal (
operation_id, job_id, intent_json, state
) VALUES (?, ?, ?, 'intent')
""",
(operation_id, job_id, intent),
)
row = connection.execute(
"SELECT * FROM file_journal WHERE operation_id = ?",
(operation_id,),
).fetchone()
return dict(row)
def complete_file_operation(
self,
operation_id: str,
result_json: str,
) -> dict[str, object]:
result = _canonical(json.loads(result_json))
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing = connection.execute(
"SELECT * FROM file_journal WHERE operation_id = ?",
(operation_id,),
).fetchone()
if existing is None:
raise FileOperationConflict("file operation intent does not exist")
if existing["state"] == "completed":
if existing["result_json"] != result:
raise FileOperationConflict(
"completed file operation has a different result"
)
return dict(existing)
connection.execute(
"""
UPDATE file_journal
SET result_json = ?, state = 'completed'
WHERE operation_id = ?
""",
(result, operation_id),
)
row = connection.execute(
"SELECT * FROM file_journal WHERE operation_id = ?",
(operation_id,),
).fetchone()
return dict(row)
def file_operation_rows(self, job_id: str) -> list[dict[str, object]]:
with self._connect() as connection:
rows = connection.execute(
"""
SELECT operation_id, job_id, intent_json, result_json, state
FROM file_journal WHERE job_id = ? ORDER BY operation_id
""",
(job_id,),
).fetchall()
return [dict(row) for row in rows]
def begin_route_attempt(
self,
command_id: str,
route_id: str,
spec_json: str,
nonce: str,
) -> dict[str, object]:
spec = _canonical(json.loads(spec_json))
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing = connection.execute(
"SELECT * FROM route_attempts WHERE command_id = ?",
(command_id,),
).fetchone()
if existing:
if (
existing["route_id"] != route_id
or existing["spec_json"] != spec
):
raise CommandConflict(
"route command ID was reused with different content"
)
return dict(existing)
connection.execute(
"""
INSERT INTO route_attempts (
command_id, route_id, spec_json, nonce, state,
last_sequence
) VALUES (?, ?, ?, ?, 'accepted', 0)
""",
(command_id, route_id, spec, nonce),
)
row = connection.execute(
"SELECT * FROM route_attempts WHERE command_id = ?",
(command_id,),
).fetchone()
return dict(row)
def record_route_update(
self,
command_id: str,
sequence: int,
state: str,
update_json: str,
) -> None:
update = _canonical(json.loads(update_json))
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing = connection.execute(
"""
SELECT update_json FROM route_attempt_updates
WHERE command_id = ? AND sequence = ?
""",
(command_id, sequence),
).fetchone()
if existing:
if existing["update_json"] != update:
raise CommandConflict(
"route update sequence has different content"
)
return
attempt = connection.execute(
"""
SELECT last_sequence FROM route_attempts
WHERE command_id = ?
""",
(command_id,),
).fetchone()
if attempt is None:
raise CommandConflict("route attempt does not exist")
expected = attempt["last_sequence"] + 1
if sequence != expected:
raise CommandConflict(
f"route update sequence must be {expected}, received {sequence}"
)
connection.execute(
"""
INSERT INTO route_attempt_updates (
command_id, sequence, update_json
) VALUES (?, ?, ?)
""",
(command_id, sequence, update),
)
connection.execute(
"""
UPDATE route_attempts
SET state = ?, last_sequence = ?
WHERE command_id = ?
""",
(state, sequence, command_id),
)
def record_route_ownership(
self, command_id: str, created_device: bool, created_folder: bool
) -> None:
with self._connect() as connection:
cursor = connection.execute(
"""
UPDATE route_attempts
SET created_device = MAX(created_device, ?),
created_folder = MAX(created_folder, ?)
WHERE command_id = ?
""",
(int(created_device), int(created_folder), command_id),
)
if cursor.rowcount != 1:
raise CommandConflict("route attempt does not exist")
def get_route_attempt(self, command_id: str) -> dict[str, object] | None:
with self._connect() as connection:
row = connection.execute(
"SELECT * FROM route_attempts WHERE command_id = ?",
(command_id,),
).fetchone()
return dict(row) if row else None
def route_update_rows(self, command_id: str) -> list[dict[str, object]]:
with self._connect() as connection:
rows = connection.execute(
"""
SELECT sequence, update_json FROM route_attempt_updates
WHERE command_id = ? ORDER BY sequence
""",
(command_id,),
).fetchall()
return [dict(row) for row in rows]
def _connect(self) -> sqlite3.Connection:
connection = sqlite3.connect(self.database, isolation_level=None, timeout=5)
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 = 5000")
return connection
def _canonical(value: object) -> str:
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
_SCHEMA_V1 = """
CREATE TABLE IF NOT EXISTS commands (
command_id TEXT PRIMARY KEY,
payload_sha256 TEXT NOT NULL,
command_json TEXT NOT NULL,
acknowledgement_json TEXT NOT NULL,
state TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS jobs (
job_id TEXT PRIMARY KEY,
definition_json TEXT NOT NULL,
state TEXT NOT NULL,
revision INTEGER NOT NULL,
last_event_sequence INTEGER NOT NULL,
committed INTEGER NOT NULL CHECK(committed IN (0, 1))
);
CREATE TABLE IF NOT EXISTS events (
event_id TEXT PRIMARY KEY,
job_id TEXT NOT NULL REFERENCES jobs(job_id),
sequence INTEGER NOT NULL,
event_json TEXT NOT NULL,
UNIQUE(job_id, sequence)
);
CREATE TABLE IF NOT EXISTS file_journal (
operation_id TEXT PRIMARY KEY,
job_id TEXT NOT NULL,
intent_json TEXT NOT NULL,
result_json TEXT,
state TEXT NOT NULL
);
"""
_SCHEMA_V2 = """
CREATE TABLE route_attempts (
command_id TEXT PRIMARY KEY REFERENCES commands(command_id),
route_id TEXT NOT NULL,
spec_json TEXT NOT NULL,
nonce TEXT NOT NULL,
state TEXT NOT NULL,
last_sequence INTEGER NOT NULL DEFAULT 0,
created_device INTEGER NOT NULL DEFAULT 0 CHECK(created_device IN (0, 1)),
created_folder INTEGER NOT NULL DEFAULT 0 CHECK(created_folder IN (0, 1))
);
CREATE TABLE route_attempt_updates (
command_id TEXT NOT NULL REFERENCES route_attempts(command_id),
sequence INTEGER NOT NULL,
update_json TEXT NOT NULL,
PRIMARY KEY(command_id, sequence)
);
"""