from __future__ import annotations import json import sqlite3 import time from pathlib import Path from typing import Iterable from uuid import uuid4 SCHEMA_VERSION = 2 class Store: def __init__(self, path: Path): path.parent.mkdir(parents=True, exist_ok=True) self.db = sqlite3.connect(path, timeout=5) self.db.row_factory = sqlite3.Row self.db.execute("PRAGMA journal_mode=WAL") self.db.execute("PRAGMA foreign_keys=ON") self.db.execute("PRAGMA busy_timeout=5000") self.db.executescript(""" CREATE TABLE IF NOT EXISTS metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL); CREATE TABLE IF NOT EXISTS messages ( message_id TEXT PRIMARY KEY, sender_repo TEXT NOT NULL, target_repo TEXT NOT NULL, body TEXT NOT NULL, created_at REAL NOT NULL, state TEXT NOT NULL, endpoint_id TEXT, provenance TEXT, injected_at REAL, acknowledged_at REAL ); CREATE INDEX IF NOT EXISTS messages_target_state ON messages(target_repo, state); CREATE TABLE IF NOT EXISTS endpoints ( endpoint_id TEXT PRIMARY KEY, pid INTEGER NOT NULL, session TEXT NOT NULL, repos TEXT NOT NULL, delivery_mode TEXT NOT NULL DEFAULT 'manual', connected_at REAL NOT NULL, disconnected_at REAL ); CREATE TABLE IF NOT EXISTS leases ( message_id TEXT PRIMARY KEY REFERENCES messages(message_id) ON DELETE CASCADE, lease_id TEXT NOT NULL, endpoint_id TEXT NOT NULL, acquired_at REAL NOT NULL, expires_at REAL NOT NULL ); """) endpoint_columns = { row["name"] for row in self.db.execute("PRAGMA table_info(endpoints)") } if "delivery_mode" not in endpoint_columns: self.db.execute( "ALTER TABLE endpoints ADD COLUMN delivery_mode TEXT NOT NULL DEFAULT 'manual'" ) self.db.execute( "INSERT INTO metadata(key,value) VALUES('schema_version',?) " "ON CONFLICT(key) DO UPDATE SET value=excluded.value", (str(SCHEMA_VERSION),), ) self.db.commit() def close(self) -> None: self.db.close() def register_endpoint( self, endpoint_id: str, pid: int, session: str, repos: list[str], delivery_mode: str = "manual", ) -> None: import json if delivery_mode not in {"manual", "pane"}: raise ValueError(f"invalid delivery mode: {delivery_mode}") self.db.execute( "UPDATE endpoints SET disconnected_at=strftime('%s','now') " "WHERE pid=? AND session=? AND endpoint_id<>? AND disconnected_at IS NULL", (pid, session, endpoint_id), ) self.db.execute( "INSERT INTO endpoints(endpoint_id,pid,session,repos,delivery_mode,connected_at,disconnected_at) VALUES(?,?,?,?,?,strftime('%s','now'),NULL) " "ON CONFLICT(endpoint_id) DO UPDATE SET pid=excluded.pid, session=excluded.session, repos=excluded.repos, delivery_mode=excluded.delivery_mode, connected_at=excluded.connected_at, disconnected_at=NULL", (endpoint_id, pid, session, json.dumps(repos), delivery_mode), ) self.db.commit() def disconnect_endpoint(self, endpoint_id: str) -> None: self.db.execute("UPDATE endpoints SET disconnected_at=strftime('%s','now') WHERE endpoint_id=?", (endpoint_id,)) self.db.commit() def disconnect_all(self) -> None: self.db.execute("UPDATE endpoints SET disconnected_at=strftime('%s','now') WHERE disconnected_at IS NULL") self.db.commit() def endpoints(self) -> list[sqlite3.Row]: return list( self.db.execute( "SELECT * FROM endpoints WHERE disconnected_at IS NULL ORDER BY endpoint_id" ) ) def endpoint(self, endpoint_id: str) -> sqlite3.Row | None: row = self.db.execute("SELECT * FROM endpoints WHERE endpoint_id=? AND disconnected_at IS NULL", (endpoint_id,)).fetchone() if row is not None: return row if endpoint_id.startswith("tmux-amq-"): visible_pid = endpoint_id.removeprefix("tmux-amq-") if not visible_pid.isdigit(): return None pid = int(visible_pid) return self.db.execute("SELECT * FROM endpoints WHERE pid=? AND disconnected_at IS NULL ORDER BY connected_at DESC LIMIT 1", (pid,)).fetchone() return None def history_stats(self) -> tuple[int, float | None]: size = 0 try: size = self.db.execute("SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()").fetchone()[0] except sqlite3.DatabaseError: pass oldest = self.db.execute("SELECT MIN(created_at) FROM messages").fetchone()[0] return size, oldest def add(self, sender: str, target: str, body: str, *, endpoint: str | None = None, provenance: str | None = None) -> str: if len(body.encode("utf-8")) > 8192: raise ValueError("message body exceeds 8 KiB limit") message_id = f"m-{uuid4()}" self.db.execute( "INSERT INTO messages VALUES(?,?,?,?,?,?,?,?,?,?)", (message_id, sender, target, body, time.time(), "pending", endpoint, provenance, None, None), ) self.db.commit() return message_id def list(self, target: str | None = None, state: str | None = None) -> list[sqlite3.Row]: clauses, values = [], [] if target: clauses.append("target_repo=?"); values.append(target) if state: clauses.append("state=?"); values.append(state) where = f" WHERE {' AND '.join(clauses)}" if clauses else "" return list(self.db.execute(f"SELECT * FROM messages{where} ORDER BY created_at", values)) def set_state(self, message_id: str, state: str) -> None: column = {"injected": "injected_at", "acknowledged": "acknowledged_at"}.get(state) if column: self.db.execute(f"UPDATE messages SET state=?, {column}=? WHERE message_id=?", (state, time.time(), message_id)) else: self.db.execute("UPDATE messages SET state=? WHERE message_id=?", (state, message_id)) self.db.commit() def acknowledge(self, message_id: str) -> bool: row = self.db.execute("SELECT 1 FROM messages WHERE message_id=?", (message_id,)).fetchone() if row is None: return False self.set_state(message_id, "acknowledged") return True def claim(self, message_id: str, endpoint_id: str, ttl: float = 30.0) -> str | None: now = time.time() lease_id = f"lease-{uuid4()}" with self.db: self.db.execute("DELETE FROM leases WHERE expires_at < ?", (now,)) try: self.db.execute("INSERT INTO leases VALUES(?,?,?,?,?)", (message_id, lease_id, endpoint_id, now, now + ttl)) except sqlite3.IntegrityError: return None return lease_id def renew(self, message_id: str, lease_id: str, ttl: float = 30.0) -> bool: with self.db: result = self.db.execute("UPDATE leases SET expires_at=? WHERE message_id=? AND lease_id=?", (time.time() + ttl, message_id, lease_id)) return result.rowcount == 1 def release(self, message_id: str, lease_id: str, state: str = "injected") -> bool: with self.db: result = self.db.execute("DELETE FROM leases WHERE message_id=? AND lease_id=?", (message_id, lease_id)) if result.rowcount: self.set_state(message_id, state) return result.rowcount == 1 def export(self, rows: Iterable[sqlite3.Row], path: Path) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as stream: for row in rows: stream.write(json.dumps(dict(row), sort_keys=True) + "\n") def purge(self, before: float | None = None, max_bytes: int | None = None) -> int: rows = self.list() ids = [r["message_id"] for r in rows if before is not None and r["created_at"] < before] if max_bytes is not None and self.db.execute("SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()").fetchone()[0] > max_bytes: ids += [r["message_id"] for r in rows if r["message_id"] not in ids][:max(0, len(rows) // 2)] if ids: self.db.executemany("DELETE FROM messages WHERE message_id=?", ((i,) for i in set(ids))) self.db.commit() return len(set(ids))