tmux-amq/src/tamq/store.py

from __future__ import annotations

import json
import sqlite3
import time
from pathlib import Path
from typing import Iterable
from uuid import uuid4

SCHEMA_VERSION = 1


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,
          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
        );
        """)
        self.db.execute("INSERT OR IGNORE INTO metadata(key,value) VALUES('schema_version',?)", (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]) -> None:
        import json
        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,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, connected_at=excluded.connected_at, disconnected_at=NULL",
            (endpoint_id, pid, session, json.dumps(repos)),
        )
        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))