conjurer/durable_queue.py

"""Dependency-free, disk-backed queue for durable message delivery.

One JSON file per key under a directory. Used on both sides of the
librarian <-> bot result path so an expensive (hours-long) search result is
never lost to a transient network failure or a restart:

* the librarian keeps a result in its OUTBOX until the bot acks it,
* the bot keeps a result in its INBOX until it is actually rendered, and
  remembers delivered uuids so duplicate resends are idempotent.

Only ``json`` + ``os`` are imported, so the logic is unit-testable without
flask, discord, or the network. Writes are atomic (temp file + ``os.replace``)
so a crash mid-write can never leave a half-written record that poisons replay.
"""
import json
import os
import tempfile
import time


def _safe_name(key: str) -> str:
    """Filesystem-safe file stem for a key (uuids are safe; be defensive)."""
    stem = "".join(c for c in str(key) if c.isalnum() or c in "-_.")
    return stem or "_"


class DiskQueue:
    """A directory of ``<key>.json`` records, each ``{key, payload, ts}``."""

    def __init__(self, directory: str):
        # No disk touch here on purpose: constructing a DiskQueue at import time
        # must not create directories (tests, read-only default paths). The
        # directory is created lazily on the first put().
        self.directory = directory

    def _path(self, key) -> str:
        return os.path.join(self.directory, _safe_name(key) + ".json")

    def put(self, key, payload) -> None:
        """Atomically write (overwrite) the record for ``key``."""
        os.makedirs(self.directory, exist_ok=True)
        record = {"key": str(key), "payload": payload, "ts": time.time()}
        fd, tmp = tempfile.mkstemp(dir=self.directory, suffix=".tmp")
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as handle:
                json.dump(record, handle)
            os.replace(tmp, self._path(key))  # atomic on POSIX
        finally:
            if os.path.exists(tmp):
                os.remove(tmp)

    def contains(self, key) -> bool:
        return os.path.exists(self._path(key))

    def get(self, key):
        """Return the payload stored for ``key``, or None if absent/unreadable."""
        try:
            with open(self._path(key), encoding="utf-8") as handle:
                return json.load(handle)["payload"]
        except (OSError, ValueError, KeyError, TypeError):
            return None

    def remove(self, key) -> None:
        try:
            os.remove(self._path(key))
        except FileNotFoundError:
            pass

    def items(self):
        """Return ``[(key, payload, ts), ...]`` oldest-first.

        Unreadable / half-written / corrupt files are skipped (never raise),
        so one bad file can't stall replay of the rest.
        """
        out = []
        try:
            names = os.listdir(self.directory)
        except FileNotFoundError:
            return out
        for name in names:
            if not name.endswith(".json"):
                continue
            try:
                with open(os.path.join(self.directory, name), encoding="utf-8") as handle:
                    record = json.load(handle)
                out.append((record["key"], record["payload"], record.get("ts", 0)))
            except (OSError, ValueError, KeyError, TypeError):
                continue
        out.sort(key=lambda triple: triple[2])
        return out

    def keys(self):
        return [key for key, _payload, _ts in self.items()]

    def __len__(self) -> int:
        return len(self.items())

    def prune(self, max_entries: int) -> int:
        """Keep only the newest ``max_entries`` (by ts); drop the rest.

        Used for the delivered-uuid set so it cannot grow without bound.
        Returns how many were dropped.
        """
        if max_entries < 0:
            return 0
        entries = self.items()  # oldest first
        excess = len(entries) - max_entries
        dropped = 0
        for key, _payload, _ts in entries[: max(0, excess)]:
            self.remove(key)
            dropped += 1
        return dropped