"""Worker-shared Find Work catalog state.""" from __future__ import annotations import json import os import secrets import stat import sqlite3 import time from dataclasses import dataclass from pathlib import Path 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): self.path = Path(path) self.clock = clock or time.time self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) os.chmod(self.path.parent, stat.S_IRWXU) old_umask = os.umask(0o077) try: 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)" ) finally: os.umask(old_umask) os.chmod(self.path, stat.S_IRUSR | stat.S_IWUSR) def _connect(self): connection = sqlite3.connect(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() return AvailableIssueSnapshotState( items=json.loads(row["items_json"]) if row["items_json"] else None, 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", (json.dumps(items, separators=(",", ":")), 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 = json.loads(row["items_json"]) 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", (json.dumps(items, separators=(",", ":")),), ) 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 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()