pr-reviewer/app/store.py

"""SQLite cache of pre-computed PR reviews.

One row per PR, keyed by ``(owner, repo, number)``. The background worker writes;
the API reads. SQLite allows only a single writer at a time, so every write goes
through one module-level lock and connections use a busy timeout; readers open
their own short-lived connection. The DB file is local and gitignored.

Status lifecycle:
  reviewing -> ready | error         (worker reviews a new/changed PR)
  ready     -> ready                 (unchanged; refreshed metadata only)
  *         -> closed                (PR left the queue; kept briefly, then purged)
"""

from __future__ import annotations

import sqlite3
import threading
import time
from typing import Any

# One writer at a time. SQLite serializes writes anyway; this keeps the Python
# side from piling up concurrent write transactions from the worker pool.
_WRITE_LOCK = threading.Lock()

_SCHEMA = """
CREATE TABLE IF NOT EXISTS reviews (
    owner         TEXT    NOT NULL,
    repo          TEXT    NOT NULL,
    number        INTEGER NOT NULL,
    status        TEXT    NOT NULL,          -- reviewing | ready | error | closed
    diff_hash     TEXT,
    review_json   TEXT,
    error         TEXT,
    attempts      INTEGER NOT NULL DEFAULT 0,
    pr_title      TEXT,
    pr_author     TEXT,
    pr_url        TEXT,
    pr_updated_at TEXT,                       -- PR's updatedAt on GitHub
    pr_created_at TEXT,                        -- PR's createdAt (for oldest->newest sort)
    pr_node_id    TEXT,                        -- PR GraphQL global id (for auto-merge)
    cached_at     REAL    NOT NULL,           -- when this row was last written
    closed_at     REAL,                       -- when the PR left the queue
    PRIMARY KEY (owner, repo, number)
);
"""

# Nullable columns added after the initial release. Migrated in via ALTER TABLE
# on an existing db. Only nullable, no-default columns are safe to ALTER-ADD in
# SQLite, so every entry here must be a bare nullable type.
_MIGRATIONS = (
    ("pr_created_at", "TEXT"),
    ("pr_node_id", "TEXT"),
)

# Columns a caller may set via upsert(); owner/repo/number are the key and
# cached_at is stamped automatically.
_UPSERTABLE = (
    "status",
    "diff_hash",
    "review_json",
    "error",
    "attempts",
    "pr_title",
    "pr_author",
    "pr_url",
    "pr_updated_at",
    "pr_created_at",
    "pr_node_id",
    "closed_at",
)


class ReviewStore:
    def __init__(self, db_path: str) -> None:
        self.db_path = db_path
        with _WRITE_LOCK, self._connect() as conn:
            conn.executescript(_SCHEMA)
            self._migrate(conn)

    def _migrate(self, conn: sqlite3.Connection) -> None:
        """Add columns introduced after a db was first created. Each ALTER is
        idempotent (guarded by a table_info check) and adds only a nullable
        column, so it is safe on an existing db and safe to re-run. Errors
        propagate so a broken db fails loudly rather than running half-migrated."""
        existing = {row[1] for row in conn.execute("PRAGMA table_info(reviews)")}
        for name, coltype in _MIGRATIONS:
            if name not in existing:
                conn.execute(f"ALTER TABLE reviews ADD COLUMN {name} {coltype}")
        conn.commit()

    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.db_path, timeout=30)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA busy_timeout=30000")
        return conn

    # --- reads --------------------------------------------------------------
    def get(self, owner: str, repo: str, number: int) -> dict[str, Any] | None:
        with self._connect() as conn:
            row = conn.execute(
                "SELECT * FROM reviews WHERE owner=? AND repo=? AND number=?",
                (owner, repo, number),
            ).fetchone()
        return dict(row) if row else None

    def list(self, *, include_closed: bool = True) -> list[dict[str, Any]]:
        sql = "SELECT * FROM reviews"
        if not include_closed:
            sql += " WHERE status != 'closed'"
        sql += " ORDER BY pr_updated_at DESC"
        with self._connect() as conn:
            rows = conn.execute(sql).fetchall()
        return [dict(r) for r in rows]

    # --- writes (serialized) ------------------------------------------------
    def upsert(self, owner: str, repo: str, number: int, **fields: Any) -> None:
        """Insert or update one PR row. Only recognized columns are written;
        ``cached_at`` is stamped automatically. Unspecified columns keep their
        previous value on update (and their default on insert)."""
        bad = set(fields) - set(_UPSERTABLE)
        if bad:
            raise ValueError(f"unknown review columns: {sorted(bad)}")

        cols = ["owner", "repo", "number", "cached_at", *fields]
        vals = [owner, repo, number, time.time(), *fields.values()]
        placeholders = ", ".join("?" for _ in cols)
        # On conflict, update every provided column plus cached_at.
        updates = ", ".join(f"{c}=excluded.{c}" for c in ("cached_at", *fields))
        sql = (
            f"INSERT INTO reviews ({', '.join(cols)}) VALUES ({placeholders}) "
            f"ON CONFLICT(owner, repo, number) DO UPDATE SET {updates}"
        )
        with _WRITE_LOCK, self._connect() as conn:
            conn.execute(sql, vals)
            conn.commit()

    def reset_stale_reviewing(self) -> int:
        """On startup, any row still marked ``reviewing`` was interrupted by a
        crash/restart. Flip it to ``error`` so the next cycle retries it."""
        with _WRITE_LOCK, self._connect() as conn:
            cur = conn.execute(
                "UPDATE reviews SET status='error', error='interrupted' "
                "WHERE status='reviewing'"
            )
            conn.commit()
            return cur.rowcount

    def mark_missing_closed(self, live_keys: set[tuple[str, str, int]]) -> None:
        """Mark rows whose PR is no longer in the queue as ``closed`` (keeping
        them briefly so an open dashboard doesn't lose a review mid-read)."""
        now = time.time()
        with _WRITE_LOCK, self._connect() as conn:
            rows = conn.execute(
                "SELECT owner, repo, number FROM reviews WHERE status != 'closed'"
            ).fetchall()
            gone = [
                (r["owner"], r["repo"], r["number"])
                for r in rows
                if (r["owner"], r["repo"], r["number"]) not in live_keys
            ]
            for owner, repo, number in gone:
                conn.execute(
                    "UPDATE reviews SET status='closed', closed_at=? "
                    "WHERE owner=? AND repo=? AND number=?",
                    (now, owner, repo, number),
                )
            conn.commit()

    def purge_closed(self, older_than_seconds: float) -> int:
        """Hard-delete rows that have been ``closed`` longer than the grace
        window."""
        cutoff = time.time() - older_than_seconds
        with _WRITE_LOCK, self._connect() as conn:
            cur = conn.execute(
                "DELETE FROM reviews WHERE status='closed' AND closed_at IS NOT NULL "
                "AND closed_at < ?",
                (cutoff,),
            )
            conn.commit()
            return cur.rowcount