pr-reviewer/app/store.py

182 lines
7.2 KiB
Python
Raw Permalink Normal View History

"""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