SQLite Job Queue with Leases

import sqlite3
import time
import uuid
from contextlib import contextmanager
from collections.abc import Iterator

SCHEMA = """
CREATE TABLE IF NOT EXISTS jobs (
    id TEXT PRIMARY KEY,
    payload TEXT NOT NULL,
    run_at REAL NOT NULL,
    locked_until REAL,
    attempts INTEGER NOT NULL DEFAULT 0
);
CREATE INDEX IF NOT EXISTS jobs_ready ON jobs (run_at);
"""


@contextmanager
def connect(path: str) -> Iterator[sqlite3.Connection]:
    conn = sqlite3.connect(path, isolation_level=None, timeout=5)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=5000")
    try:
        yield conn
    finally:
        conn.close()


def enqueue(conn: sqlite3.Connection, payload: str, delay: float = 0.0) -> str:
    job_id = uuid.uuid4().hex
    conn.execute(
        "INSERT INTO jobs (id, payload, run_at) VALUES (?, ?, ?)",
        (job_id, payload, time.time() + delay),
    )
    return job_id


def claim(conn: sqlite3.Connection, lease: float = 30.0) -> dict[str, object] | None:
    now = time.time()
    conn.execute("BEGIN IMMEDIATE")
    try:
        row = conn.execute(
            """
            UPDATE jobs SET locked_until = ?, attempts = attempts + 1
            WHERE id = (
                SELECT id FROM jobs
                WHERE run_at <= ? AND (locked_until IS NULL OR locked_until < ?)
                ORDER BY run_at LIMIT 1
            )
            RETURNING id, payload, attempts
            """,
            (now + lease, now, now),
        ).fetchone()
        conn.execute("COMMIT")
    except Exception:
        conn.execute("ROLLBACK")
        raise
    return dict(row) if row else None


def finish(conn: sqlite3.Connection, job_id: str) -> None:
    conn.execute("DELETE FROM jobs WHERE id = ?", (job_id,))


def retry_later(conn: sqlite3.Connection, job_id: str, delay: float = 0.0) -> None:
    conn.execute(
        "UPDATE jobs SET locked_until = NULL, run_at = ? WHERE id = ?",
        (time.time() + delay, job_id),
    )


if __name__ == "__main__":
    with connect(":memory:") as conn:
        conn.executescript(SCHEMA)
        enqueue(conn, "send-email")
        enqueue(conn, "resize-image")
        first = claim(conn)
        print("claimed:", first["payload"], "attempts:", first["attempts"])
        finish(conn, str(first["id"]))
        second = claim(conn)
        print("next:", second["payload"])
        retry_later(conn, str(second["id"]), delay=60)
        print("waiting jobs:", claim(conn))