"""Worker-shared Find Work catalog state.""" from __future__ import annotations import secrets import sqlite3 import time from dataclasses import dataclass from pathlib import Path from src.private_state import connect_private_sqlite from src.state_encryption import ( PrivateStateCipher, PrivateStateEncryptionError, private_state_encryption_key, ) class RefreshLeaseLost(RuntimeError): pass @dataclass(frozen=True) class AvailableIssueSnapshotState: items: list[dict] | None created_at: float | None retry_at: float | None refreshing: bool lease_expires_at: float | None class AvailableIssueSnapshotStore: def __init__(self, path, *, clock=None, encryption_key=None): self.path = Path(path) self.clock = clock or time.time self._cipher = PrivateStateCipher( encryption_key if encryption_key is not None else private_state_encryption_key(), store="available-issue-snapshot", ) with self._connect() as connection: connection.executescript( """ CREATE TABLE IF NOT EXISTS available_issue_snapshot ( singleton INTEGER PRIMARY KEY CHECK (singleton = 1), items_json TEXT, created_at REAL, retry_at REAL ); CREATE TABLE IF NOT EXISTS available_issue_refresh_lease ( singleton INTEGER PRIMARY KEY CHECK (singleton = 1), owner TEXT NOT NULL, expires_at REAL NOT NULL ); CREATE TABLE IF NOT EXISTS available_issue_claim ( repository TEXT NOT NULL, number INTEGER NOT NULL, PRIMARY KEY (repository, number) ); """ ) connection.execute( "INSERT OR IGNORE INTO available_issue_snapshot VALUES (1, NULL, NULL, NULL)" ) def _connect(self): connection = connect_private_sqlite(self.path, timeout=1.0, isolation_level=None) connection.row_factory = sqlite3.Row connection.execute("PRAGMA busy_timeout = 1000") return connection def try_acquire_refresh(self, *, lease_seconds: float) -> str | None: if lease_seconds <= 0: raise ValueError("refresh lease duration must be positive") now = self.clock() owner = secrets.token_hex(16) with self._connect() as connection: connection.execute("BEGIN IMMEDIATE") active = connection.execute( "SELECT expires_at FROM available_issue_refresh_lease WHERE singleton = 1" ).fetchone() if active is not None and active["expires_at"] > now: connection.rollback() return None connection.execute("DELETE FROM available_issue_refresh_lease WHERE singleton = 1") connection.execute( "INSERT INTO available_issue_refresh_lease VALUES (1, ?, ?)", (owner, now + lease_seconds), ) connection.commit() return owner def load(self) -> AvailableIssueSnapshotState: now = self.clock() with self._connect() as connection: row = connection.execute( "SELECT * FROM available_issue_snapshot WHERE singleton = 1" ).fetchone() lease = connection.execute( "SELECT expires_at FROM available_issue_refresh_lease " "WHERE singleton = 1 AND expires_at > ?", (now,) ).fetchone() items = None if row["items_json"]: items, legacy = self._cipher.open(row["items_json"]) if not isinstance(items, list): raise PrivateStateEncryptionError("private state could not be decrypted") if legacy: migrated = self._cipher.seal(items) with self._connect() as connection: connection.execute( "UPDATE available_issue_snapshot SET items_json = ? " "WHERE singleton = 1 AND items_json = ?", (migrated, row["items_json"]), ) return AvailableIssueSnapshotState( items=items, created_at=row["created_at"], retry_at=row["retry_at"], refreshing=lease is not None, lease_expires_at=lease["expires_at"] if lease else None, ) def publish(self, owner: str | None, *, items: list[dict]) -> AvailableIssueSnapshotState: now = self.clock() with self._connect() as connection: connection.execute("BEGIN IMMEDIATE") lease = connection.execute( "SELECT owner, expires_at FROM available_issue_refresh_lease WHERE singleton = 1" ).fetchone() if ( owner is None or lease is None or lease["owner"] != owner or lease["expires_at"] <= now ): connection.rollback() raise RefreshLeaseLost("available issue refresh lease expired or changed owner") claimed = { (row["repository"], row["number"]) for row in connection.execute( "SELECT repository, number FROM available_issue_claim" ) } items = [ item for item in items if (item.get("repository"), item.get("number")) not in claimed ] connection.execute( "UPDATE available_issue_snapshot SET items_json = ?, created_at = ?, " "retry_at = NULL WHERE singleton = 1", (self._cipher.seal(items), now), ) connection.execute("DELETE FROM available_issue_refresh_lease WHERE singleton = 1") connection.commit() return self.load() def remove_claimed(self, repository: str, number: int) -> AvailableIssueSnapshotState: with self._connect() as connection: connection.execute("BEGIN IMMEDIATE") connection.execute( "INSERT OR IGNORE INTO available_issue_claim VALUES (?, ?)", (repository, number), ) row = connection.execute( "SELECT items_json FROM available_issue_snapshot WHERE singleton = 1" ).fetchone() items = self._cipher.open(row["items_json"])[0] if row["items_json"] else None if items is not None: items = [ item for item in items if item.get("repository") != repository or item.get("number") != number ] connection.execute( "UPDATE available_issue_snapshot SET items_json = ? WHERE singleton = 1", (self._cipher.seal(items),), ) connection.commit() return self.load() def invalidate(self) -> AvailableIssueSnapshotState: with self._connect() as connection: connection.execute("BEGIN IMMEDIATE") connection.execute( "UPDATE available_issue_snapshot SET items_json = NULL, " "created_at = NULL, retry_at = NULL WHERE singleton = 1" ) connection.execute("DELETE FROM available_issue_claim") connection.commit() return self.load() def release_refresh(self, owner: str | None) -> None: if owner is None: return with self._connect() as connection: connection.execute( "DELETE FROM available_issue_refresh_lease " "WHERE singleton = 1 AND owner = ?", (owner,), ) def fail_refresh(self, owner: str | None, *, retry_at: float) -> AvailableIssueSnapshotState: now = self.clock() with self._connect() as connection: connection.execute("BEGIN IMMEDIATE") lease = connection.execute( "SELECT owner, expires_at FROM available_issue_refresh_lease WHERE singleton = 1" ).fetchone() if ( owner is None or lease is None or lease["owner"] != owner or lease["expires_at"] <= now ): connection.rollback() raise RefreshLeaseLost("available issue refresh lease expired or changed owner") connection.execute( "UPDATE available_issue_snapshot SET retry_at = ? WHERE singleton = 1", (retry_at,), ) connection.execute("DELETE FROM available_issue_refresh_lease WHERE singleton = 1") connection.commit() return self.load()