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

766 lines
28 KiB
Python

"""Durable command inbox and client recovery cursors."""
from __future__ import annotations
import hashlib
import json
import os
import sqlite3
import stat
import threading
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
SCHEMA_VERSION = 3
# A client daemon executes unrelated jobs concurrently. SQLite still permits
# only one writer, so serialize this process's short state transactions rather
# than allowing an otherwise healthy job to fail after its busy timeout.
_DATABASE_LOCK = threading.RLock()
class CommandConflict(RuntimeError):
pass
class JobConflict(RuntimeError):
pass
class FileOperationConflict(RuntimeError):
pass
@dataclass(frozen=True)
class CommandAcceptance:
duplicate: bool
acknowledgement_json: str
state: 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")
version = 2
if version == 2:
connection.executescript(_SCHEMA_V3)
connection.execute("PRAGMA user_version = 3")
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, state
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"], existing["state"]
)
status = json.loads(acknowledgement).get("status")
state = (
"rejected"
if status is not None
and status != "COMMAND_ACK_STATUS_ACCEPTED"
else "accepted"
)
connection.execute(
"""
INSERT INTO commands (
command_id, payload_sha256, command_json,
acknowledgement_json, state
) VALUES (?, ?, ?, ?, ?)
""",
(command_id, digest, payload, acknowledgement, state),
)
return CommandAcceptance(False, acknowledgement, state)
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 WHERE state = 'accepted' 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 ensure_job_definition(self, job_id: str, definition_json: str) -> None:
"""Durably record an assignment without inventing a global event."""
definition = _canonical(json.loads(definition_json))
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing = connection.execute(
"SELECT definition_json FROM jobs WHERE job_id = ?", (job_id,)
).fetchone()
if existing is not None:
if existing["definition_json"] != definition:
raise JobConflict("job definition is immutable")
return
connection.execute(
"""
INSERT INTO jobs (job_id, definition_json, state, revision,
last_event_sequence, committed)
VALUES (?, ?, 'JOB_STATE_QUEUED', 0, 0, 0)
""",
(job_id, definition),
)
def is_command_active(self, command_id: str) -> bool:
with self._connect() as connection:
row = connection.execute(
"SELECT state FROM commands WHERE command_id = ?", (command_id,)
).fetchone()
return row is not None and row["state"] == "accepted"
def reconcile_job(
self,
*,
job_id: str,
definition_json: str,
state: str,
revision: int,
last_event_sequence: int,
committed: bool,
superseded_command_ids: list[str],
) -> None:
"""Apply control's cursor, retiring only stale command leases.
This drops unacknowledged local journal rows beyond control's cursor;
resource files and operation artifacts are intentionally retained.
"""
definition = _canonical(json.loads(definition_json))
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing = connection.execute(
"SELECT definition_json FROM jobs WHERE job_id = ?", (job_id,)
).fetchone()
if existing is not None and existing["definition_json"] != definition:
raise JobConflict("reconciliation has a different job definition")
connection.execute(
"DELETE FROM events WHERE job_id = ? AND sequence > ?",
(job_id, last_event_sequence),
)
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)),
)
if superseded_command_ids:
placeholders = ",".join("?" for _ in superseded_command_ids)
rows = connection.execute(
f"SELECT command_id, command_json FROM commands WHERE command_id IN ({placeholders})",
superseded_command_ids,
).fetchall()
owned_ids = [
row["command_id"] for row in rows
if _command_job_id(row["command_json"]) == job_id
]
if owned_ids:
owned_placeholders = ",".join("?" for _ in owned_ids)
connection.execute(
f"UPDATE commands SET state = 'superseded' WHERE command_id IN ({owned_placeholders})",
owned_ids,
)
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 put_job_artifact(
self, job_id: str, kind: str, value: dict[str, object]
) -> dict[str, object]:
encoded = _canonical(value)
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing = connection.execute(
"""
SELECT value_json FROM job_artifacts
WHERE job_id = ? AND kind = ?
""",
(job_id, kind),
).fetchone()
if existing is not None and existing["value_json"] != encoded:
raise JobConflict(f"job artifact {kind} is immutable")
connection.execute(
"""
INSERT INTO job_artifacts (job_id, kind, value_json)
VALUES (?, ?, ?)
ON CONFLICT(job_id, kind) DO NOTHING
""",
(job_id, kind, encoded),
)
row = connection.execute(
"""
SELECT job_id, kind, value_json FROM job_artifacts
WHERE job_id = ? AND kind = ?
""",
(job_id, kind),
).fetchone()
result = dict(row)
result["value"] = json.loads(str(result.pop("value_json")))
return result
def get_job_artifact(
self, job_id: str, kind: str
) -> dict[str, object] | None:
with self._connect() as connection:
row = connection.execute(
"""
SELECT job_id, kind, value_json FROM job_artifacts
WHERE job_id = ? AND kind = ?
""",
(job_id, kind),
).fetchone()
if row is None:
return None
result = dict(row)
result["value"] = json.loads(str(result.pop("value_json")))
return result
def record_job_event(
self,
*,
job_id: str,
definition_json: str,
event_id: str,
event_json: str,
state: str,
revision: int,
sequence: int,
committed: bool,
) -> dict[str, object]:
definition = _canonical(json.loads(definition_json))
event = _canonical(json.loads(event_json))
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
existing_event = connection.execute(
"""
SELECT job_id, sequence, event_json
FROM events WHERE event_id = ?
""",
(event_id,),
).fetchone()
if existing_event:
if (
existing_event["job_id"] != job_id
or existing_event["sequence"] != sequence
or existing_event["event_json"] != event
):
raise JobConflict(
"job event ID was reused with different content"
)
return dict(existing_event)
existing_sequence = connection.execute(
"""
SELECT event_id, event_json FROM events
WHERE job_id = ? AND sequence = ?
""",
(job_id, sequence),
).fetchone()
if existing_sequence:
if existing_sequence["event_json"] != event:
raise JobConflict(
"job event sequence has conflicting content"
)
return dict(existing_sequence)
job = connection.execute(
"""
SELECT definition_json, revision, last_event_sequence,
committed
FROM jobs WHERE job_id = ?
""",
(job_id,),
).fetchone()
if job and job["definition_json"] != definition:
raise JobConflict("job definition is immutable")
if job and (
revision < job["revision"]
or sequence <= job["last_event_sequence"]
or (job["committed"] and not committed)
):
raise JobConflict("job event cursor cannot move backwards")
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, sequence,
int(committed),
),
)
connection.execute(
"""
INSERT INTO events (
event_id, job_id, sequence, event_json
) VALUES (?, ?, ?, ?)
""",
(event_id, job_id, sequence, event),
)
row = connection.execute(
"SELECT * FROM events WHERE event_id = ?", (event_id,)
).fetchone()
return dict(row)
def job_event_rows(
self, job_id: str, after_sequence: int = 0
) -> list[dict[str, object]]:
with self._connect() as connection:
rows = connection.execute(
"""
SELECT event_id, job_id, sequence, event_json
FROM events
WHERE job_id = ? AND sequence > ?
ORDER BY sequence
""",
(job_id, after_sequence),
).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]
@contextmanager
def _connect(self):
"""Yield one connection while serializing local SQLite writers.
The longer SQLite timeout also covers a short lock held by a separate
maintenance process such as a backup. The lock is deliberately held
for the full transaction, including ``BEGIN IMMEDIATE``.
"""
with _DATABASE_LOCK:
connection = sqlite3.connect(
self.database, isolation_level=None, timeout=30
)
try:
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")
# Preserve the original ``with connection`` commit/rollback
# behavior used by every store operation.
with connection:
yield connection
finally:
connection.close()
def _command_job_id(command_json: str) -> str:
"""Return the job target from canonical protobuf JSON, if it has one."""
command = json.loads(command_json)
if "assignJob" in command:
return str(command["assignJob"].get("job", {}).get("jobId", ""))
if "executeStep" in command:
return str(command["executeStep"].get("jobId", ""))
if "cancelJob" in command:
return str(command["cancelJob"].get("jobId", ""))
return ""
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)
);
"""
_SCHEMA_V3 = """
CREATE TABLE job_artifacts (
job_id TEXT NOT NULL REFERENCES jobs(job_id),
kind TEXT NOT NULL,
value_json TEXT NOT NULL,
PRIMARY KEY(job_id, kind)
);
"""