drew
0bc734c020
Applies `ruff format` to the accumulated formatting debt on this branch. Formatting-only — no behavioral changes. Required for CI/lint's format gate (`nox -s format -- --check`), which the branch was failing on 288 tracked files that drifted from ruff's canonical style. In-progress WIP files are intentionally excluded so this commit stays a clean formatting-only diff. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1084 lines
41 KiB
Python
1084 lines
41 KiB
Python
"""SQLite-backed store for the PR State Warmer's snapshot of open PRs.
|
||
|
||
The :mod:`pr_state_warmer` sidecar polls Forgejo's ``/pulls?state=open``
|
||
endpoint every ``PR_STATE_WARMER_INTERVAL_S`` seconds (default 30s),
|
||
walks every page, and writes the full PR object (number, head.sha,
|
||
labels, body, mergeable, updated_at, etc.) here. Dispatchers READ from
|
||
this cache instead of calling Forgejo per cycle.
|
||
|
||
v3 (2026-05-17): adds ``comments_refreshed_updated_at`` to drive the
|
||
warmer's comments-cache refresh queue from persistent state, not
|
||
from the per-cycle ``upsert_prs`` delta. Stores the PR's
|
||
``updated_at`` value at the time of the refresh — pending refreshes
|
||
are then ``WHERE updated_at != COALESCE(comments_refreshed_updated_at, '')``.
|
||
Using the PR's own ``updated_at`` (string equality) rather than a
|
||
wall-clock stamp avoids two failure modes: (1) every cycle's
|
||
unchanged-upsert advances ``last_seen_at`` which would trigger
|
||
spurious refreshes, and (2) Forgejo's ``updated_at`` format and
|
||
our wall-clock format have different string-sort semantics. The
|
||
change is ADDITIVE — the migration runs ``ALTER TABLE ADD COLUMN``
|
||
rather than DROP/CREATE so a v2 → v3 upgrade preserves data and a
|
||
deploy with version skew can co-exist (a v2 reader sees the table
|
||
without the new column, which is harmless because v2 readers never
|
||
SELECT it).
|
||
|
||
Why this design
|
||
---------------
|
||
|
||
The dispatcher's per-cycle ``/pulls`` call was both:
|
||
|
||
1. **Truncated at 50 PRs** — single page, no pagination. Past 50 open
|
||
PRs the 51st+ were silently invisible.
|
||
2. **Flaky** — Forgejo's response builds 16 head+base repo+owner
|
||
profiles inline (127 KB for 8 PRs). First request after the
|
||
server's hot-cache window expires takes 24-30s; dispatcher's
|
||
short timeout fires and serves stale.
|
||
|
||
The warmer eliminates both problems: pagination is a warmer-side
|
||
detail; Forgejo's hot cache stays permanently primed by the 30s
|
||
poll cadence (shorter than the ~60-90s cache TTL), so cold rebuilds
|
||
collapse to ~1/day. Dispatchers do a local SQLite read (~0ms).
|
||
|
||
Storage
|
||
-------
|
||
|
||
SQLite at ``/tmp/cleveragents-pr-state/state.sqlite3`` (overridable
|
||
via ``PR_STATE_CACHE_DIR``). WAL mode so the warmer-writer and
|
||
dispatcher-readers can coexist without lock contention.
|
||
|
||
Schema::
|
||
|
||
CREATE TABLE pr_state (
|
||
owner TEXT NOT NULL, -- repo owner (e.g. 'drew')
|
||
repo TEXT NOT NULL, -- repo name (e.g. 'cleveragents-core')
|
||
number INTEGER NOT NULL, -- PR number within (owner, repo)
|
||
body_json TEXT NOT NULL, -- full PR object as JSON
|
||
head_sha TEXT NOT NULL,
|
||
state TEXT NOT NULL, -- 'open' / 'closed' / 'merged'
|
||
labels_json TEXT NOT NULL, -- JSON list of label-name strings
|
||
updated_at TEXT NOT NULL, -- last-modified at Forgejo
|
||
first_seen_at TEXT NOT NULL, -- when the warmer first added the row
|
||
last_seen_at TEXT NOT NULL, -- last warmer cycle that found it in /pulls
|
||
vanished_at TEXT, -- NULL while open; set when removed from /pulls
|
||
PRIMARY KEY (owner, repo, number)
|
||
);
|
||
|
||
The ``(owner, repo)`` discriminator prevents a warmer/dispatcher
|
||
cfg mismatch (e.g. running tests against the canonical repo while
|
||
the prod warmer wrote rows for the fork) from silently serving the
|
||
wrong-repo PRs. Every read takes ``(owner, repo)`` and queries
|
||
within that scope only.
|
||
|
||
Vanish semantics: when a PR disappears from the warmer's enumeration
|
||
(closed/merged externally, or the operator deleted it), the row is
|
||
NOT immediately deleted — instead ``vanished_at`` is stamped on the
|
||
first cycle that misses it AND the row keeps for a grace period
|
||
(default 7 days). This lets late-arriving dispatcher cycles still
|
||
see the PR's last known state for telemetry/audit purposes.
|
||
|
||
Errors
|
||
------
|
||
|
||
Operational errors raise :class:`PRStateCacheError`. SQLite errors
|
||
propagate as ``sqlite3.DatabaseError`` so a corrupt DB is
|
||
distinguishable from a missing row.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import datetime as _dt
|
||
import fcntl
|
||
import functools
|
||
import json
|
||
import logging
|
||
import os
|
||
import sqlite3
|
||
import threading
|
||
import time as _time
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
_logger = logging.getLogger("pr_state_cache")
|
||
|
||
# v2 added the (owner, repo) primary-key prefix so a warmer/dispatcher
|
||
# cfg mismatch can't silently serve wrong-repo PRs.
|
||
# v3 added comments_refreshed_at to drive deferral state from
|
||
# persistent storage instead of in-memory.
|
||
# Bump when changing schema. Migration policy below for upgrade paths.
|
||
SCHEMA_VERSION = 3
|
||
|
||
_DEFAULT_CACHE_DIR = Path("/tmp/cleveragents-pr-state")
|
||
_CACHE_DIR_ENV = "PR_STATE_CACHE_DIR"
|
||
_DISABLE_ENV = "PR_STATE_CACHE_DISABLE"
|
||
|
||
# How long a vanished row sticks around before janitor sweeps it.
|
||
# A week is generous — most operators close PRs and don't care, but
|
||
# audit/telemetry walks may want to see the last-known state of a
|
||
# recently-closed PR.
|
||
_VANISHED_GRACE_S_DEFAULT = 7 * 24 * 3600
|
||
_VANISHED_GRACE_S_ENV = "PR_STATE_CACHE_VANISHED_GRACE_S"
|
||
|
||
|
||
class PRStateCacheError(RuntimeError):
|
||
"""Operational error from the PR-state cache."""
|
||
|
||
|
||
def cache_dir() -> Path:
|
||
return Path(os.environ.get(_CACHE_DIR_ENV) or str(_DEFAULT_CACHE_DIR))
|
||
|
||
|
||
def cache_path() -> Path:
|
||
return cache_dir() / "state.sqlite3"
|
||
|
||
|
||
def is_disabled() -> bool:
|
||
raw = os.environ.get(_DISABLE_ENV, "").strip().lower()
|
||
return raw in {"1", "true", "yes", "on"}
|
||
|
||
|
||
def _vanished_grace_s() -> int:
|
||
raw = os.environ.get(_VANISHED_GRACE_S_ENV)
|
||
if raw:
|
||
try:
|
||
return max(1, int(raw))
|
||
except ValueError:
|
||
pass
|
||
return _VANISHED_GRACE_S_DEFAULT
|
||
|
||
|
||
def _now() -> str:
|
||
return _dt.datetime.now(_dt.timezone.utc).isoformat()
|
||
|
||
|
||
def _normalize_updated_at(value: str) -> str:
|
||
"""Strip sub-second precision and canonicalize the timezone marker
|
||
so two values representing the same instant compare equal.
|
||
|
||
Why: ``comments_refreshed_updated_at`` is matched against
|
||
``pr_state.updated_at`` via string equality in the pending-refresh
|
||
query. Forgejo's ``updated_at`` can emit different representations
|
||
of the same instant across versions / proxies (``Z`` vs ``+00:00``,
|
||
microseconds vs whole-second). Without normalization a single
|
||
representation flap triggers infinite re-refresh for every PR.
|
||
|
||
Conservative: input that doesn't look like an ISO-8601 timestamp
|
||
is returned untouched."""
|
||
if not value:
|
||
return value
|
||
t_pos = value.find("T")
|
||
if t_pos < 0:
|
||
return value
|
||
# Find the timezone suffix and split.
|
||
tz_pos = -1
|
||
for sep in ("+", "-"):
|
||
candidate = value.find(sep, t_pos)
|
||
if candidate > t_pos:
|
||
tz_pos = candidate
|
||
break
|
||
if tz_pos > 0:
|
||
dt_part = value[:tz_pos]
|
||
tz_part = "+00:00" if value[tz_pos:] in ("+00:00", "+0000") else value[tz_pos:]
|
||
elif value and value[-1] in ("Z", "z"):
|
||
dt_part = value[:-1]
|
||
tz_part = "+00:00" # canonicalize Z → +00:00
|
||
else:
|
||
dt_part = value
|
||
tz_part = ""
|
||
dot = dt_part.find(".")
|
||
if dot >= 0:
|
||
dt_part = dt_part[:dot]
|
||
return dt_part + tz_part
|
||
|
||
|
||
# SQLite caps placeholders per statement at SQLITE_MAX_VARIABLE_NUMBER
|
||
# (32766 on modern builds). ``mark_vanished`` builds an IN-clause from
|
||
# the seen-numbers set; if the warmer's pagination cap is raised past
|
||
# ~650 pages × 50/page = 32 500 PRs, a single IN-clause would overflow.
|
||
# Chunk well under the limit (3 reserved slots for owner/repo/now).
|
||
_MARK_VANISHED_CHUNK = 32_000
|
||
|
||
_MIGRATION_LOCK_FILENAME = "migration.lock"
|
||
# Bounded retry on the cross-process migration lock. A stale lock
|
||
# (previous holder SIGKILL'd) would block every subsequent `_connect()`
|
||
# forever under unbounded `LOCK_EX`. With LOCK_NB + retry we surface
|
||
# a clear timeout error after N seconds instead.
|
||
_MIGRATION_LOCK_TIMEOUT_S = 30
|
||
_MIGRATION_LOCK_POLL_S = 0.5
|
||
|
||
# Per-process migration state. Schema migration runs at most ONCE per
|
||
# process — the first connect performs the migration if needed and
|
||
# stamps ``_initialized = True``. Subsequent connects skip the
|
||
# migration block entirely.
|
||
#
|
||
# Thread safety: a ``threading.Lock`` guards the check/set of
|
||
# ``_initialized`` so two threads in the same process can't both
|
||
# observe "not initialised", both re-migrate, and stamp the flag
|
||
# while a third thread sees a partially-migrated state. Used in
|
||
# concert with the cross-process file lock.
|
||
#
|
||
# Cross-process safety: an OS-level ``fcntl.flock`` on
|
||
# ``<cache_dir>/migration.lock`` serializes the migration window so N
|
||
# dispatcher processes + the warmer can't race destructive DDL. The
|
||
# migration policy is also restricted (see ``_run_migration_locked``)
|
||
# to ADDITIVE changes whenever possible — a DROP only happens on the
|
||
# pre-v2 cases where the column shape is fundamentally incompatible.
|
||
#
|
||
# Re-heal: if a public-API call gets ``sqlite3.OperationalError`` for
|
||
# ``no such table`` or ``no such column``, we clear ``_initialized``
|
||
# and retry once. Catches the externally-replaced/corrupted-DB case
|
||
# without requiring a process restart. A ``sqlite3.DatabaseError``
|
||
# (file corruption) also triggers re-heal but additionally quarantines
|
||
# the broken file.
|
||
_initialized = False
|
||
_init_lock = threading.Lock()
|
||
|
||
|
||
def _connect() -> sqlite3.Connection:
|
||
"""Open a SQLite connection, running schema migration ONCE per
|
||
process. Subsequent calls in the same process skip the migration
|
||
block.
|
||
|
||
The migration block is guarded by:
|
||
- ``_init_lock`` (in-process): ensures only one thread per
|
||
process runs the migration body; other threads block until
|
||
it returns.
|
||
- ``fcntl.flock`` (cross-process, inside ``_run_migration_locked``):
|
||
ensures only one process across the host runs destructive DDL.
|
||
"""
|
||
global _initialized
|
||
target = cache_path()
|
||
target.parent.mkdir(parents=True, exist_ok=True)
|
||
conn = sqlite3.connect(
|
||
str(target),
|
||
timeout=10.0,
|
||
isolation_level=None, # autocommit; explicit BEGIN/COMMIT below
|
||
)
|
||
conn.row_factory = sqlite3.Row
|
||
conn.execute("PRAGMA journal_mode = WAL")
|
||
conn.execute("PRAGMA synchronous = NORMAL")
|
||
if not _initialized:
|
||
with _init_lock:
|
||
# Re-check under the lock — another thread might have
|
||
# initialised while we waited.
|
||
if not _initialized:
|
||
_run_migration_locked(conn, target.parent)
|
||
_initialized = True
|
||
return conn
|
||
|
||
|
||
def _run_migration_locked(conn: sqlite3.Connection, lock_dir: Path) -> None:
|
||
"""Acquire an exclusive file lock and run migrations.
|
||
|
||
Policy:
|
||
- ``on_disk == SCHEMA_VERSION``: no-op (another process raced
|
||
ahead while we waited for the lock).
|
||
- ``on_disk == 0`` (pre-versioning): DROP + CREATE. The pre-v2
|
||
column shape is fundamentally incompatible (no ``owner`` /
|
||
``repo`` columns) so an ALTER cannot preserve the rows
|
||
meaningfully.
|
||
- ``0 < on_disk < SCHEMA_VERSION``: additive ALTER TABLE ADD
|
||
COLUMN per known step. Preserves data across upgrade.
|
||
- ``on_disk > SCHEMA_VERSION``: REFUSE to mutate. A newer
|
||
version has already touched the DB; downgrading would lose
|
||
data the newer process expects. Log + leave alone — caller
|
||
will likely see a schema-shape error and we let it surface
|
||
rather than papering over with a DROP.
|
||
"""
|
||
lock_path = lock_dir / _MIGRATION_LOCK_FILENAME
|
||
with open(lock_path, "w") as lock_handle:
|
||
# Bounded retry instead of unbounded LOCK_EX: a stale lock
|
||
# (previous holder SIGKILL'd, NFS hang, etc.) would otherwise
|
||
# block every subsequent _connect() forever with no signal.
|
||
# LOCK_NB + poll gives us a clear timeout error and lets the
|
||
# operator see WHICH process is stuck.
|
||
deadline = _time.monotonic() + _MIGRATION_LOCK_TIMEOUT_S
|
||
while True:
|
||
try:
|
||
fcntl.flock(
|
||
lock_handle.fileno(),
|
||
fcntl.LOCK_EX | fcntl.LOCK_NB,
|
||
)
|
||
break
|
||
except BlockingIOError:
|
||
if _time.monotonic() >= deadline:
|
||
# Read holder PID from the lock file if any.
|
||
try:
|
||
holder = lock_path.read_text().strip() or "(unknown)"
|
||
except OSError:
|
||
holder = "(read failed)"
|
||
raise PRStateCacheError(
|
||
f"pr_state cache migration lock at {lock_path} "
|
||
f"held by another process (pid={holder}) for "
|
||
f">{_MIGRATION_LOCK_TIMEOUT_S}s; refusing to "
|
||
f"wait further. Inspect with: lsof {lock_path}"
|
||
) from None
|
||
_time.sleep(_MIGRATION_LOCK_POLL_S)
|
||
# Stamp our PID so a future timeout can identify the holder.
|
||
try:
|
||
lock_handle.seek(0)
|
||
lock_handle.truncate()
|
||
lock_handle.write(f"{os.getpid()}\n")
|
||
lock_handle.flush()
|
||
except OSError:
|
||
pass # cosmetic
|
||
try:
|
||
on_disk = int(conn.execute("PRAGMA user_version").fetchone()[0])
|
||
# Detect the externally-dropped case: user_version stamp
|
||
# is still correct (it lives in the SQLite header, not
|
||
# the table) but the table itself is gone. Treat the same
|
||
# as a fresh DB — CREATE TABLE puts us back in shape and
|
||
# the user_version stamp is already correct.
|
||
table_exists = (
|
||
conn.execute(
|
||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='pr_state'"
|
||
).fetchone()
|
||
is not None
|
||
)
|
||
if on_disk == SCHEMA_VERSION and table_exists:
|
||
# Another process won the race (or migration is
|
||
# already current); nothing to do.
|
||
_ensure_indices(conn)
|
||
return
|
||
if on_disk == SCHEMA_VERSION and not table_exists:
|
||
_logger.warning(
|
||
"pr_state cache: user_version matches but table "
|
||
"is missing — recreating (external drop / "
|
||
"corruption recovery)",
|
||
)
|
||
_create_table_at_latest(conn)
|
||
_ensure_indices(conn)
|
||
return
|
||
if on_disk > SCHEMA_VERSION:
|
||
_logger.error(
|
||
"pr_state cache user_version=%s is NEWER than this "
|
||
"process knows about (expected %s). Refusing to "
|
||
"mutate — a downgrade would corrupt data the newer "
|
||
"process expects. Upgrade this process or wipe the "
|
||
"cache dir (%s) manually if you really want to "
|
||
"downgrade.",
|
||
on_disk,
|
||
SCHEMA_VERSION,
|
||
lock_dir,
|
||
)
|
||
return
|
||
# Versions < 2 used a fundamentally different schema
|
||
# (no owner/repo columns) — ALTER cannot preserve the
|
||
# rows meaningfully. From v2 onward, every step is
|
||
# additive, so the upgrade preserves data.
|
||
if on_disk < 2:
|
||
_logger.info(
|
||
"pr_state cache: legacy DB found (user_version=%s); "
|
||
"destructive rebuild required — pre-v2 schemas "
|
||
"lack the owner/repo columns",
|
||
on_disk,
|
||
)
|
||
conn.execute("DROP TABLE IF EXISTS pr_state")
|
||
_create_table_at_latest(conn)
|
||
else:
|
||
# 2 <= on_disk < SCHEMA_VERSION → additive upgrade.
|
||
_logger.info(
|
||
"pr_state cache: additive upgrade from v%s to v%s",
|
||
on_disk,
|
||
SCHEMA_VERSION,
|
||
)
|
||
_apply_additive_upgrades(conn, on_disk)
|
||
conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
|
||
_ensure_indices(conn)
|
||
finally:
|
||
# flock auto-releases on fd close (the `with open(...)`
|
||
# block); explicit release is redundant but documents
|
||
# intent for the reader.
|
||
fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN)
|
||
|
||
|
||
def _create_table_at_latest(conn: sqlite3.Connection) -> None:
|
||
"""Build the table in its current shape. Called from a fresh-DB
|
||
path or after a pre-versioning DROP."""
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS pr_state (
|
||
owner TEXT NOT NULL,
|
||
repo TEXT NOT NULL,
|
||
number INTEGER NOT NULL,
|
||
body_json TEXT NOT NULL,
|
||
head_sha TEXT NOT NULL,
|
||
state TEXT NOT NULL,
|
||
labels_json TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL,
|
||
first_seen_at TEXT NOT NULL,
|
||
last_seen_at TEXT NOT NULL,
|
||
vanished_at TEXT,
|
||
comments_refreshed_updated_at TEXT,
|
||
PRIMARY KEY (owner, repo, number)
|
||
)
|
||
"""
|
||
)
|
||
|
||
|
||
def _apply_additive_upgrades(conn: sqlite3.Connection, from_version: int) -> None:
|
||
"""Apply ALTER TABLE steps from ``from_version`` up to
|
||
``SCHEMA_VERSION``. Each step is non-destructive (ADD COLUMN
|
||
only), so a deploy with version skew can co-exist: a v2 reader
|
||
sees the v3 table without the new column and ignores it; a v3
|
||
reader sees its own column populated by the v3 writer."""
|
||
# v2 → v3: add comments_refreshed_updated_at.
|
||
if from_version < 3:
|
||
try:
|
||
conn.execute(
|
||
"ALTER TABLE pr_state ADD COLUMN comments_refreshed_updated_at TEXT"
|
||
)
|
||
except sqlite3.OperationalError as exc:
|
||
# ``duplicate column name`` is benign — concurrent
|
||
# migration must have raced (shouldn't with the lock,
|
||
# but be defensive).
|
||
if "duplicate column" not in str(exc).lower():
|
||
raise
|
||
|
||
|
||
def _ensure_indices(conn: sqlite3.Connection) -> None:
|
||
"""Create the indices we depend on. ``IF NOT EXISTS`` makes this
|
||
idempotent so callers can run it on every migration outcome."""
|
||
conn.execute("CREATE INDEX IF NOT EXISTS idx_pr_state_state ON pr_state(state)")
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_pr_state_updated_at ON pr_state(updated_at)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_pr_state_vanished_at ON pr_state(vanished_at)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_pr_state_comments_refreshed "
|
||
"ON pr_state(comments_refreshed_updated_at)"
|
||
)
|
||
|
||
|
||
_CORRUPT_DB_MARKERS = (
|
||
"malformed",
|
||
"database disk image is malformed",
|
||
"file is not a database",
|
||
"not a database",
|
||
"database is corrupt",
|
||
)
|
||
|
||
|
||
def _quarantine_corrupt_db() -> None:
|
||
"""Rename the SQLite file out of the way so the next ``_connect()``
|
||
starts fresh. Called when ``sqlite3.DatabaseError`` indicates
|
||
structural corruption (header damage, truncated WAL).
|
||
|
||
Best-effort: a failure to rename logs + propagates so the caller
|
||
can decide. We rename rather than delete so the file is available
|
||
for forensics; the warmer's next cycle re-seeds within seconds."""
|
||
target = cache_path()
|
||
if not target.exists():
|
||
return
|
||
ts = _dt.datetime.now(_dt.timezone.utc).strftime("%Y%m%dT%H%M%S")
|
||
quarantine = target.with_suffix(f".sqlite3.corrupt.{ts}")
|
||
try:
|
||
target.rename(quarantine)
|
||
_logger.error(
|
||
"pr_state cache: quarantined corrupt DB to %s; next cycle "
|
||
"will re-seed from a fresh file",
|
||
quarantine,
|
||
)
|
||
except OSError as exc:
|
||
_logger.error(
|
||
"pr_state cache: FAILED to quarantine corrupt DB at %s: %s "
|
||
"— manual intervention required",
|
||
target,
|
||
exc,
|
||
)
|
||
raise
|
||
|
||
|
||
def _reheal_and_retry(exc: sqlite3.Error) -> bool:
|
||
"""Decide whether a SQLite error warrants clearing the
|
||
per-process ``_initialized`` flag and letting the caller retry.
|
||
|
||
Two paths:
|
||
1. ``OperationalError: no such table`` / ``no such column`` —
|
||
the DB file was replaced or its schema mutated out from under
|
||
a long-running warmer. Re-arm; retry against the re-migrated
|
||
schema.
|
||
2. ``DatabaseError`` with corruption marker — the file itself
|
||
is broken. Quarantine it and re-arm; the retry will start
|
||
from a fresh file.
|
||
|
||
Returns True iff caller should retry. False if the exception
|
||
isn't one we know how to recover from.
|
||
"""
|
||
global _initialized
|
||
msg = str(exc).lower()
|
||
if isinstance(exc, sqlite3.OperationalError) and (
|
||
"no such table" in msg or "no such column" in msg
|
||
):
|
||
with _init_lock:
|
||
_initialized = False
|
||
_logger.warning(
|
||
"pr_state cache: re-arming migration after operational "
|
||
"error (%s) — likely external DB replacement / schema drop",
|
||
exc,
|
||
)
|
||
return True
|
||
if isinstance(exc, sqlite3.DatabaseError) and any(
|
||
marker in msg for marker in _CORRUPT_DB_MARKERS
|
||
):
|
||
try:
|
||
_quarantine_corrupt_db()
|
||
except OSError:
|
||
return False
|
||
with _init_lock:
|
||
_initialized = False
|
||
return True
|
||
return False
|
||
|
||
|
||
# ─── Public API ─────────────────────────────────────────────────────
|
||
|
||
|
||
def _with_reheal(fn):
|
||
"""Wrap a public-API call so a recoverable SQLite error triggers
|
||
a one-shot re-arm + retry. Targets two failure classes:
|
||
|
||
- ``OperationalError`` "no such table"/"no such column" —
|
||
externally-replaced DB; re-migrate against the file as-is.
|
||
- ``DatabaseError`` corruption markers — quarantine the
|
||
broken file, re-migrate from scratch.
|
||
|
||
Single retry — if the second attempt also fails, the error
|
||
propagates so the caller sees a real signal instead of an
|
||
infinite loop. ``DatabaseError`` is a SUPERCLASS of
|
||
``OperationalError`` so the catch order matters: subclass first.
|
||
"""
|
||
|
||
@functools.wraps(fn)
|
||
def wrapper(*args, **kwargs):
|
||
try:
|
||
return fn(*args, **kwargs)
|
||
except sqlite3.OperationalError as exc:
|
||
if _reheal_and_retry(exc):
|
||
return fn(*args, **kwargs)
|
||
raise
|
||
except sqlite3.DatabaseError as exc:
|
||
if _reheal_and_retry(exc):
|
||
return fn(*args, **kwargs)
|
||
raise
|
||
|
||
return wrapper
|
||
|
||
|
||
@_with_reheal
|
||
def upsert_prs(
|
||
prs: list[dict[str, Any]],
|
||
*,
|
||
owner: str,
|
||
repo: str,
|
||
) -> dict[str, Any]:
|
||
"""Write each PR's state within ``(owner, repo)``. Returns counts
|
||
AND the changed-number list (union of inserted + updated). See
|
||
module docstring for the schema details.
|
||
|
||
Caller passes the full PR object dicts returned by Forgejo's
|
||
``/pulls`` endpoint."""
|
||
if is_disabled():
|
||
raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
|
||
if not prs:
|
||
return {"inserted": 0, "updated": 0, "unchanged": 0, "changed_numbers": []}
|
||
now = _now()
|
||
counts = {"inserted": 0, "updated": 0, "unchanged": 0}
|
||
changed: list[int] = []
|
||
conn = _connect()
|
||
try:
|
||
with conn:
|
||
for pr in prs:
|
||
if not isinstance(pr, dict):
|
||
continue
|
||
try:
|
||
number = int(pr.get("number"))
|
||
except (TypeError, ValueError):
|
||
continue
|
||
if number <= 0:
|
||
continue
|
||
head_sha = ""
|
||
head = pr.get("head")
|
||
if isinstance(head, dict):
|
||
head_sha = str(head.get("sha") or "")
|
||
state = str(pr.get("state") or "open")
|
||
labels = []
|
||
raw_labels = pr.get("labels") or []
|
||
if isinstance(raw_labels, list):
|
||
labels = [
|
||
str(l.get("name") or "")
|
||
for l in raw_labels
|
||
if isinstance(l, dict) and l.get("name")
|
||
]
|
||
# Canonicalize on write so the pending-comments-refresh
|
||
# query (string equality vs comments_refreshed_updated_at)
|
||
# is robust to Forgejo's timezone-marker drift.
|
||
updated_at = _normalize_updated_at(str(pr.get("updated_at") or ""))
|
||
body_json = json.dumps(pr, default=str, sort_keys=True)
|
||
labels_json = json.dumps(labels)
|
||
cached = conn.execute(
|
||
"SELECT updated_at FROM pr_state "
|
||
"WHERE owner = ? AND repo = ? AND number = ?",
|
||
(owner, repo, number),
|
||
).fetchone()
|
||
if cached is None:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO pr_state (
|
||
owner, repo, number, body_json, head_sha, state,
|
||
labels_json, updated_at, first_seen_at,
|
||
last_seen_at, vanished_at
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL)
|
||
""",
|
||
(
|
||
owner,
|
||
repo,
|
||
number,
|
||
body_json,
|
||
head_sha,
|
||
state,
|
||
labels_json,
|
||
updated_at,
|
||
now,
|
||
now,
|
||
),
|
||
)
|
||
counts["inserted"] += 1
|
||
changed.append(number)
|
||
elif cached["updated_at"] != updated_at:
|
||
conn.execute(
|
||
"""
|
||
UPDATE pr_state SET
|
||
body_json = ?,
|
||
head_sha = ?,
|
||
state = ?,
|
||
labels_json = ?,
|
||
updated_at = ?,
|
||
last_seen_at = ?,
|
||
vanished_at = NULL
|
||
WHERE owner = ? AND repo = ? AND number = ?
|
||
""",
|
||
(
|
||
body_json,
|
||
head_sha,
|
||
state,
|
||
labels_json,
|
||
updated_at,
|
||
now,
|
||
owner,
|
||
repo,
|
||
number,
|
||
),
|
||
)
|
||
counts["updated"] += 1
|
||
changed.append(number)
|
||
else:
|
||
# Refresh body_json + labels_json even on the
|
||
# ``updated_at`` unchanged branch — Forgejo
|
||
# mutations like label add/remove do NOT bump
|
||
# ``updated_at``, so without this the warmer's
|
||
# cached body would carry stale labels for as
|
||
# long as the PR sat idle, and downstream
|
||
# consumers (cycle-cap, claim sweeps, filter
|
||
# exclusions) would never see them. Cheap: same
|
||
# JSON we already computed above.
|
||
conn.execute(
|
||
"""
|
||
UPDATE pr_state SET
|
||
body_json = ?,
|
||
labels_json = ?,
|
||
last_seen_at = ?,
|
||
vanished_at = NULL
|
||
WHERE owner = ? AND repo = ? AND number = ?
|
||
""",
|
||
(body_json, labels_json, now, owner, repo, number),
|
||
)
|
||
counts["unchanged"] += 1
|
||
finally:
|
||
conn.close()
|
||
return {**counts, "changed_numbers": changed}
|
||
|
||
|
||
@_with_reheal
|
||
def mark_vanished(
|
||
seen_numbers: set[int],
|
||
*,
|
||
owner: str,
|
||
repo: str,
|
||
) -> int:
|
||
"""Mark any open-state row within ``(owner, repo)`` whose
|
||
``number`` is NOT in ``seen_numbers`` as vanished. Returns the
|
||
number of rows newly marked vanished.
|
||
|
||
Scoped to ``(owner, repo)`` so a multi-repo warmer doesn't
|
||
accidentally vanish PRs from a sibling repo on a single poll.
|
||
|
||
Implementation note: ``seen_numbers`` larger than
|
||
``_MARK_VANISHED_CHUNK`` would overflow SQLite's per-statement
|
||
placeholder cap. We invert the semantics for large sets — write
|
||
a temp table of seen IDs, then ``NOT IN (SELECT ...)`` — same
|
||
result, no placeholder limit. The large-set branch runs both the
|
||
temp-table populate AND the UPDATE inside an explicit
|
||
BEGIN/COMMIT pair so a signal mid-batch can't leave the
|
||
intermediate temp data without its accompanying mutation.
|
||
"""
|
||
if is_disabled():
|
||
raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
|
||
now = _now()
|
||
conn = _connect()
|
||
try:
|
||
if not seen_numbers:
|
||
with conn:
|
||
cur = conn.execute(
|
||
"UPDATE pr_state SET vanished_at = ? "
|
||
"WHERE owner = ? AND repo = ? AND vanished_at IS NULL",
|
||
(now, owner, repo),
|
||
)
|
||
return int(cur.rowcount or 0)
|
||
if len(seen_numbers) <= _MARK_VANISHED_CHUNK:
|
||
with conn:
|
||
placeholders = ",".join("?" * len(seen_numbers))
|
||
cur = conn.execute(
|
||
f"""
|
||
UPDATE pr_state SET vanished_at = ?
|
||
WHERE owner = ? AND repo = ?
|
||
AND vanished_at IS NULL
|
||
AND number NOT IN ({placeholders})
|
||
""",
|
||
(now, owner, repo, *seen_numbers),
|
||
)
|
||
return int(cur.rowcount or 0)
|
||
# Large set: stage in a TEMP table to dodge the placeholder
|
||
# cap. Autocommit + ``with conn:`` is a no-op (no implicit
|
||
# BEGIN), so we drive the transaction explicitly to make the
|
||
# whole temp-populate + UPDATE atomic. Failure mid-batch
|
||
# rolls back the temp inserts AND any partial UPDATE; success
|
||
# commits both as one unit. TEMP tables are connection-scoped
|
||
# so cleanup happens on close.
|
||
conn.execute("BEGIN")
|
||
try:
|
||
conn.execute("DROP TABLE IF EXISTS _seen_pr_numbers")
|
||
conn.execute("CREATE TEMP TABLE _seen_pr_numbers (n INTEGER PRIMARY KEY)")
|
||
conn.executemany(
|
||
"INSERT OR IGNORE INTO _seen_pr_numbers (n) VALUES (?)",
|
||
((int(n),) for n in seen_numbers),
|
||
)
|
||
cur = conn.execute(
|
||
"""
|
||
UPDATE pr_state SET vanished_at = ?
|
||
WHERE owner = ? AND repo = ?
|
||
AND vanished_at IS NULL
|
||
AND number NOT IN (SELECT n FROM _seen_pr_numbers)
|
||
""",
|
||
(now, owner, repo),
|
||
)
|
||
rowcount = int(cur.rowcount or 0)
|
||
conn.execute("COMMIT")
|
||
return rowcount
|
||
except BaseException:
|
||
try:
|
||
conn.execute("ROLLBACK")
|
||
except sqlite3.DatabaseError:
|
||
pass
|
||
raise
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
@_with_reheal
|
||
def list_open_prs(*, owner: str, repo: str) -> list[dict[str, Any]]:
|
||
"""Return the full PR objects for every cached row within
|
||
``(owner, repo)`` that is currently open, sorted by
|
||
``updated_at`` desc (matches Forgejo's ``?sort=newest``)."""
|
||
if is_disabled():
|
||
raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
|
||
conn = _connect()
|
||
try:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT body_json FROM pr_state
|
||
WHERE owner = ? AND repo = ? AND vanished_at IS NULL
|
||
ORDER BY updated_at DESC
|
||
""",
|
||
(owner, repo),
|
||
).fetchall()
|
||
finally:
|
||
conn.close()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
out.append(json.loads(r["body_json"]))
|
||
except (ValueError, TypeError) as exc:
|
||
_logger.warning("malformed body_json in pr_state cache: %s", exc)
|
||
return out
|
||
|
||
|
||
@_with_reheal
|
||
def get_pr(number: int, *, owner: str, repo: str) -> dict[str, Any] | None:
|
||
"""Return one PR's full object within ``(owner, repo)`` by number,
|
||
or None if not cached (or if the row is marked vanished)."""
|
||
if is_disabled():
|
||
raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
|
||
conn = _connect()
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT body_json FROM pr_state "
|
||
"WHERE owner = ? AND repo = ? AND number = ? "
|
||
"AND vanished_at IS NULL",
|
||
(owner, repo, int(number)),
|
||
).fetchone()
|
||
finally:
|
||
conn.close()
|
||
if row is None:
|
||
return None
|
||
try:
|
||
return json.loads(row["body_json"])
|
||
except (ValueError, TypeError):
|
||
return None
|
||
|
||
|
||
@_with_reheal
|
||
def count_rows(
|
||
*,
|
||
owner: str | None = None,
|
||
repo: str | None = None,
|
||
) -> dict[str, int]:
|
||
"""Operator-facing diagnostic: ``{"total", "open", "vanished"}``.
|
||
|
||
When ``owner`` + ``repo`` are both given, counts are scoped to
|
||
that repo. Both omitted: global counts across every cached repo."""
|
||
if is_disabled():
|
||
raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
|
||
conn = _connect()
|
||
try:
|
||
where = ""
|
||
args: tuple[Any, ...] = ()
|
||
if owner is not None and repo is not None:
|
||
where = " WHERE owner = ? AND repo = ?"
|
||
args = (owner, repo)
|
||
total = conn.execute(
|
||
f"SELECT COUNT(*) AS c FROM pr_state{where}",
|
||
args,
|
||
).fetchone()["c"]
|
||
open_where = where + (" AND " if where else " WHERE ") + "vanished_at IS NULL"
|
||
open_n = conn.execute(
|
||
f"SELECT COUNT(*) AS c FROM pr_state{open_where}",
|
||
args,
|
||
).fetchone()["c"]
|
||
finally:
|
||
conn.close()
|
||
return {"total": int(total), "open": int(open_n), "vanished": int(total - open_n)}
|
||
|
||
|
||
@_with_reheal
|
||
def latest_write_at(*, owner: str, repo: str) -> str | None:
|
||
"""Most-recent ``last_seen_at`` across all rows within
|
||
``(owner, repo)`` — i.e. when the warmer last successfully
|
||
polled. Consumers use this as a staleness signal: if the
|
||
timestamp is older than ~2× warmer interval, the warmer is
|
||
likely dead and the cache is stale.
|
||
|
||
Returns ``None`` when no rows exist (cold cache)."""
|
||
if is_disabled():
|
||
raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
|
||
conn = _connect()
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT MAX(last_seen_at) AS latest FROM pr_state "
|
||
"WHERE owner = ? AND repo = ?",
|
||
(owner, repo),
|
||
).fetchone()
|
||
finally:
|
||
conn.close()
|
||
if row is None or row["latest"] is None:
|
||
return None
|
||
return str(row["latest"])
|
||
|
||
|
||
@_with_reheal
|
||
def mark_comments_refreshed(
|
||
refreshed: list[tuple[int, str]],
|
||
*,
|
||
owner: str,
|
||
repo: str,
|
||
) -> int:
|
||
"""Stamp each row's ``comments_refreshed_updated_at`` to the PR's
|
||
``updated_at`` value at the moment of refresh. Returns rowcount.
|
||
|
||
``refreshed`` is a list of ``(pr_number, pr_updated_at)`` tuples
|
||
— the caller passes the ``updated_at`` value the comments cache
|
||
was just refreshed against, so a subsequent Forgejo-side change
|
||
to the PR bumps ``pr_state.updated_at`` past the stamped value
|
||
and the pending-refresh query re-surfaces the PR.
|
||
"""
|
||
if is_disabled():
|
||
raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
|
||
if not refreshed:
|
||
return 0
|
||
# Normalize the stamped values so a Forgejo timezone-format flap
|
||
# (``Z`` vs ``+00:00``, microseconds vs none) on the same instant
|
||
# doesn't make the next ``list_pending`` query treat the row as
|
||
# "changed since refresh" forever.
|
||
rows_to_stamp = [
|
||
(_normalize_updated_at(str(updated_at)), owner, repo, int(number))
|
||
for number, updated_at in refreshed
|
||
]
|
||
conn = _connect()
|
||
try:
|
||
# With ``isolation_level=None`` (autocommit) the ``with conn:``
|
||
# context is a no-op — it does NOT wrap the loop in a single
|
||
# transaction. Drive BEGIN/COMMIT explicitly so the N updates
|
||
# are atomic: a failure on row K rolls back rows 0..K-1 too,
|
||
# and the warmer's retry won't see partial state.
|
||
conn.execute("BEGIN")
|
||
try:
|
||
conn.executemany(
|
||
"""
|
||
UPDATE pr_state SET comments_refreshed_updated_at = ?
|
||
WHERE owner = ? AND repo = ? AND number = ?
|
||
""",
|
||
rows_to_stamp,
|
||
)
|
||
# executemany on a CONNECTION-level cursor doesn't expose
|
||
# a stable per-statement rowcount across drivers; run a
|
||
# confirmation SELECT for the diag count.
|
||
placeholders = ",".join("?" * len(refreshed))
|
||
row = conn.execute(
|
||
f"""
|
||
SELECT COUNT(*) AS c FROM pr_state
|
||
WHERE owner = ? AND repo = ?
|
||
AND number IN ({placeholders})
|
||
""",
|
||
(owner, repo, *(int(n) for n, _ in refreshed)),
|
||
).fetchone()
|
||
conn.execute("COMMIT")
|
||
return int(row["c"] or 0)
|
||
except BaseException:
|
||
try:
|
||
conn.execute("ROLLBACK")
|
||
except sqlite3.DatabaseError:
|
||
pass
|
||
raise
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
@_with_reheal
|
||
def count_pending_comments_refresh(*, owner: str, repo: str) -> int:
|
||
"""Count rows whose comments-cache refresh is stale or unwritten.
|
||
Lets the warmer report accurate ``comments_deferred`` in its diag
|
||
without materialising every pending row into memory.
|
||
|
||
Raises ``PRStateCacheError`` when the cache is disabled — matches
|
||
every other public-API call (was silently returning 0 in an
|
||
earlier revision, which hid the disabled state in operator diags)."""
|
||
if is_disabled():
|
||
raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
|
||
conn = _connect()
|
||
try:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT COUNT(*) AS c FROM pr_state
|
||
WHERE owner = ? AND repo = ?
|
||
AND vanished_at IS NULL
|
||
AND updated_at != COALESCE(comments_refreshed_updated_at, '')
|
||
""",
|
||
(owner, repo),
|
||
).fetchone()
|
||
finally:
|
||
conn.close()
|
||
return int(row["c"] or 0)
|
||
|
||
|
||
@_with_reheal
|
||
def list_pending_comments_refresh(
|
||
*,
|
||
owner: str,
|
||
repo: str,
|
||
limit: int,
|
||
) -> list[tuple[int, str]]:
|
||
"""Return ``(pr_number, updated_at)`` tuples for rows within
|
||
``(owner, repo)`` whose comments-cache refresh has not seen the
|
||
current ``updated_at`` value — i.e. either the row has never been
|
||
refreshed (NULL stamp) or the PR has been touched on Forgejo
|
||
since the last refresh.
|
||
|
||
Ordered by ``updated_at`` desc (most-actionable first), capped
|
||
at ``limit``. Drives the warmer's refresh queue from PERSISTENT
|
||
state rather than the per-cycle ``upsert_prs`` delta: a burst
|
||
that overflows the per-cycle cap surfaces deferred rows next
|
||
cycle even if nothing else changes (the old in-memory deferral
|
||
lost them on warmer restart AND on the next cycle, because
|
||
``upsert_prs`` had already updated the cached ``updated_at`` so
|
||
the per-cycle delta no longer flagged them).
|
||
"""
|
||
if is_disabled():
|
||
raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
|
||
if limit <= 0:
|
||
return []
|
||
conn = _connect()
|
||
try:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT number, updated_at FROM pr_state
|
||
WHERE owner = ? AND repo = ?
|
||
AND vanished_at IS NULL
|
||
AND updated_at != COALESCE(comments_refreshed_updated_at, '')
|
||
ORDER BY updated_at DESC
|
||
LIMIT ?
|
||
""",
|
||
(owner, repo, int(limit)),
|
||
).fetchall()
|
||
finally:
|
||
conn.close()
|
||
return [(int(r["number"]), str(r["updated_at"])) for r in rows]
|
||
|
||
|
||
def janitor() -> int:
|
||
"""Delete vanished rows older than ``_vanished_grace_s()``.
|
||
Returns the number of rows removed. Best-effort: never raises."""
|
||
if is_disabled():
|
||
return 0
|
||
cutoff = (
|
||
_dt.datetime.now(_dt.timezone.utc) - _dt.timedelta(seconds=_vanished_grace_s())
|
||
).isoformat()
|
||
conn: sqlite3.Connection | None = None
|
||
try:
|
||
conn = _connect()
|
||
with conn:
|
||
cur = conn.execute(
|
||
"DELETE FROM pr_state WHERE vanished_at IS NOT NULL "
|
||
"AND vanished_at <= ?",
|
||
(cutoff,),
|
||
)
|
||
return int(cur.rowcount or 0)
|
||
except (OSError, sqlite3.DatabaseError) as exc:
|
||
_logger.warning("pr-state janitor swallowed error: %s", exc)
|
||
return 0
|
||
finally:
|
||
if conn is not None:
|
||
try:
|
||
conn.close()
|
||
except sqlite3.DatabaseError:
|
||
pass
|
||
|
||
|
||
__all__ = (
|
||
"PRStateCacheError",
|
||
"SCHEMA_VERSION",
|
||
"cache_dir",
|
||
"cache_path",
|
||
"count_rows",
|
||
"get_pr",
|
||
"is_disabled",
|
||
"janitor",
|
||
"latest_write_at",
|
||
"count_pending_comments_refresh",
|
||
"list_open_prs",
|
||
"list_pending_comments_refresh",
|
||
"mark_comments_refreshed",
|
||
"mark_vanished",
|
||
"upsert_prs",
|
||
)
|