An unbounded functools.cache style decorator that persists results to a SQLite file

disk_cache.py

"""An unbounded functools.cache style decorator that persists results to a SQLite file.

Usage (note that the decorator takes a required .sqlite path):

    from disk_cache import disk_cache

    @disk_cache("/var/cache/myapp/expensive.sqlite")
    def expensive(x, y, *, scale=1):
        ...

    expensive.cache_clear()   # delete every cached row
    expensive.cache_vacuum()  # repack the file now (also runs at exit if dirty)
    expensive.cache_path      # the Path to the .sqlite file

Storage:  a single SQLite database with one row per distinct call, mapping the
          pickle of (args, kwargs) to the pickled result. The key is not a hash,
          so it stays human-readable for reverse engineering -- pickle.loads(key)
          gives back the call that produced the row:
              pickle.loads(key) -> (args, kwargs)
          e.g. the call expensive(3, 7, scale=2) reads back as
          ((3, 7), {'scale': 2}). Nothing is compressed; the result is a plain
          pickle BLOB.
          The database uses the default rollback journal, so at rest there is
          exactly one file: the -journal appears only during an active write
          and SQLite removes it on commit. No WAL, no lingering sidecars.

Concurrency:
  * SQLite's own file locking does all the coordinating -- across threads and
    processes alike. Each thread gets its own connection (a sqlite3 connection
    may not be shared between threads), and SQLite locks between connections
    just as it does between processes: readers share, a writer escalates to an
    exclusive lock on the whole database, and the PENDING lock keeps a waiting
    writer from being starved by a stream of new readers. A busy timeout makes
    a blocked caller wait rather than raise SQLITE_BUSY. No application-level
    lock is needed on top: our reads are short autocommit SELECTs and our
    writes are a single INSERT, so we never hold a read lock while upgrading to
    a write -- the one situation the busy timeout cannot resolve.
  * The function body runs BEFORE the write, never inside a transaction, so a
    slow miss never holds a lock.
  * Concurrent callers of the same missing key still run the body once each
    (no dedup). Stores use INSERT OR REPLACE, so writes never corrupt; the
    last writer wins.
  * Under heavy write concurrency from many threads in one process, writers
    lean on the busy timeout's retry and may raise SQLITE_BUSY if it is
    exhausted; raise busy_timeout (or add a write mutex) if that ever happens.

Maintenance:
  * Writing one row per miss leaves the file fragmented. A process that wrote
    at least once runs VACUUM on a clean exit to repack it; read-only runs skip
    it. It is multi-process safe: only writers vacuum, a short busy timeout
    means a contended vacuum is skipped rather than stalling exit, and errors
    are swallowed. Call <func>.cache_vacuum() to repack on demand. cache_clear()
    vacuums immediately, since a DELETE only frees pages without shrinking the
    file. VACUUM cost scales with file size, so for a very large cache you may
    prefer to vacuum manually rather than on every dirty exit.

Notes:
  * Arguments and results must be picklable, and the key pickle must be stable
    across process restarts -- effectively the same restriction functools.cache
    enforces by requiring hashable args. In particular do NOT key on set or
    frozenset: their iteration order is hash-randomized, so the same value can
    pickle to different bytes across restarts. Never a wrong hit, but since
    nothing here evicts, each restart's fresh key leaves a permanent dead row --
    an unbounded leak, not a one-off miss. (Pin PYTHONHASHSEED to avoid it.)
  * Unlike functools.cache (typed=False), equal but differently typed keys are
    not collapsed: f(1), f(1.0) and f(True) map to distinct rows.
  * Use one .sqlite file per function. The key encodes only the arguments, so
    two functions sharing a file would collide on equal arguments.
"""

from __future__ import annotations

import atexit
import functools
import pickle
import sqlite3
import threading
from pathlib import Path
from typing import Any, Callable, ParamSpec, TypeVar

P = ParamSpec("P")
R = TypeVar("R")

# Pinned so identical arguments produce the same key across Python versions.
_PICKLE_PROTOCOL = 5

# Distinguishes a genuine cache miss from a stored result that happens to be
# None, without a second lookup.
_MISS = object()


class _SqliteCache:
    """A tiny key/value store backed by one rollback-journal SQLite file.

    SQLite's file locking plus a busy timeout coordinate all access -- across
    threads and processes. Connections are thread-local because a sqlite3
    connection may not be shared across threads, and SQLite locks between those
    per-thread connections exactly as it does between processes.
    """

    def __init__(self, path: Path) -> None:
        self.path = path
        self._local = threading.local()
        self._dirty = False  # did this process write since the last vacuum?
        path.parent.mkdir(parents=True, exist_ok=True)
        with self._connect() as conn:
            # Rollback journal (the default): one file at rest, no WAL sidecars.
            # Set explicitly so a file left in WAL mode by an older version is
            # converted back on first open.
            conn.execute("PRAGMA journal_mode=DELETE")
            conn.execute(
                "CREATE TABLE IF NOT EXISTS cache ("
                "  key    BLOB PRIMARY KEY,"
                "  result BLOB NOT NULL"
                ")"
            )
        # Repack the file on a clean exit -- but only if we wrote to it, so
        # warm read-only runs skip straight past.
        atexit.register(self._vacuum_on_exit)

    def _connect(self) -> sqlite3.Connection:
        conn = getattr(self._local, "conn", None)
        if conn is None:
            conn = sqlite3.connect(str(self.path))
            # Wait rather than raise if another process holds the write lock.
            conn.execute("PRAGMA busy_timeout=30000")
            self._local.conn = conn
        return conn

    def get(self, key: bytes) -> Any:
        row = self._connect().execute(
            "SELECT result FROM cache WHERE key = ?", (key,)
        ).fetchone()
        if row is None:
            return _MISS
        return pickle.loads(row[0])

    def set(self, key: bytes, result: bytes) -> None:
        conn = self._connect()
        with conn:  # commits (or rolls back) the INSERT
            conn.execute(
                "INSERT OR REPLACE INTO cache (key, result) VALUES (?, ?)",
                (key, result),
            )
        self._dirty = True

    def clear(self) -> None:
        conn = self._connect()
        with conn:
            conn.execute("DELETE FROM cache")
        # DELETE only frees pages; the file keeps its size until a VACUUM. The
        # whole table is now slack, so reclaim it right away rather than waiting
        # for exit. _dirty stays set until vacuum() confirms, so a vacuum lost to
        # a cross-process race is still retried on exit.
        self._dirty = True
        self.vacuum()

    def vacuum(self) -> None:
        """Rewrite the file to reclaim slack (free and fragmented pages).

        One-row-at-a-time writes leave the file fragmented, which VACUUM repacks
        (freelist_count stays ~0 yet the file still shrinks). VACUUM needs brief
        exclusive access and cannot run inside a transaction, so it uses its own
        autocommit connection with a short busy timeout: if another process
        holds the file it is skipped rather than stalling exit. Any error is
        swallowed so a vacuum never breaks the caller.
        """
        try:
            conn = sqlite3.connect(str(self.path), isolation_level=None)
            try:
                conn.execute("PRAGMA busy_timeout=1000")
                conn.execute("VACUUM")
                self._dirty = False
            finally:
                conn.close()
        except sqlite3.Error:
            pass

    def _vacuum_on_exit(self) -> None:
        if self._dirty:
            self.vacuum()


def disk_cache(path: Path | str) -> Callable[[Callable[P, R]], Callable[P, R]]:
    """Persist a function's results to the SQLite file at ``path``."""
    sqlite_path = Path(path)
    assert "sqlite" in sqlite_path.suffix, "For the avoidance of doubt a sqlite file extension is required"

    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        store = _SqliteCache(sqlite_path)

        @functools.wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            # The pickle of the call is the primary key. Not a hash, so it
            # decodes straight back to (args, kwargs) via pickle.loads.
            key = pickle.dumps((args, kwargs), protocol=_PICKLE_PROTOCOL)

            cached = store.get(key)
            if cached is not _MISS:
                return cached

            # Compute before the write, never inside a transaction, so a slow
            # body never holds a database lock.
            result = func(*args, **kwargs)

            store.set(key, pickle.dumps(result, protocol=_PICKLE_PROTOCOL))
            return result

        wrapper.cache_clear = store.clear    # type: ignore[attr-defined]
        wrapper.cache_vacuum = store.vacuum  # type: ignore[attr-defined]
        wrapper.cache_path = sqlite_path     # type: ignore[attr-defined]
        return wrapper

    return decorator