diff --git a/better_memory/db/schema.py b/better_memory/db/schema.py index 5332e99..64014b6 100644 --- a/better_memory/db/schema.py +++ b/better_memory/db/schema.py @@ -5,21 +5,46 @@ version (the ``NNNN`` prefix) is recorded in the ``schema_migrations`` table so re-running :func:`apply_migrations` is a no-op. +Concurrent starts (two MCP servers pointing at the same DB) are serialised +via a two-phase claim on the version row: + +* ``applied_at IS NULL`` — a process has claimed the version and is running + the DDL. Peers seeing this row poll until it either transitions or + disappears. +* ``applied_at IS NOT NULL`` — the DDL completed. Peers skip. +* row absent — no one holds a claim. A peer may attempt to claim. + +The claim ``INSERT OR IGNORE`` is atomic under SQLite's writer serialisation, +so exactly one process wins the race and runs the DDL. On DDL failure the +winner deletes its claim row so peers (or the next start) can retry rather +than skipping past a half-applied migration. Peers time out after a generous +per-version budget and raise, so a stuck peer never lets a loser return with +a broken schema (the pre-fix regression this design closes). + Each file is executed via :meth:`sqlite3.Connection.executescript`, which issues an implicit ``COMMIT`` before running the script; migrations are -therefore **not** atomic. On failure the database may be left in a partial -state — for first-time installs the recovery is to discard the DB file and -re-run. Multi-file migrations that require atomicity must use a different -execution path. +therefore **not** atomic within a single file. On failure the database may +be left in a partial state — for first-time installs the recovery is to +discard the DB file and re-run. Multi-file migrations that require +atomicity must use a different execution path. """ from __future__ import annotations import sqlite3 +import time from pathlib import Path _DEFAULT_MIGRATIONS_DIR = Path(__file__).parent / "migrations" +# How long a peer will wait for the current claim-holder to finish DDL on +# one migration before giving up and raising. Sized generously against the +# slowest realistic migration; hitting the deadline means the claim-holder +# died without cleaning up, which we surface as an error rather than +# silently skipping. +_CLAIM_WAIT_SECONDS = 120.0 +_CLAIM_POLL_INTERVAL = 0.1 + def _ensure_schema_migrations_table(conn: sqlite3.Connection) -> None: """Bootstrap the migrations-tracking table if it does not yet exist.""" @@ -35,7 +60,16 @@ def _ensure_schema_migrations_table(conn: sqlite3.Connection) -> None: def _applied_versions(conn: sqlite3.Connection) -> set[str]: - rows = conn.execute("SELECT version FROM schema_migrations").fetchall() + """Return versions that are recorded as *completed*. + + A row with ``applied_at IS NULL`` is an in-progress claim, not a + completed migration; excluding it here means a peer that sees the + row still enters the polling loop rather than mistaking a live claim + for a finished apply. + """ + rows = conn.execute( + "SELECT version FROM schema_migrations WHERE applied_at IS NOT NULL" + ).fetchall() return {row[0] for row in rows} @@ -46,6 +80,118 @@ def _version_from_filename(path: Path) -> str: return version +def _try_claim(conn: sqlite3.Connection, version: str) -> bool: + """Attempt to atomically claim ``version`` as in-progress. + + Explicit ``applied_at = NULL`` overrides the column default so peers + can tell a live claim from a completed apply. + """ + cur = conn.execute( + "INSERT OR IGNORE INTO schema_migrations (version, applied_at) " + "VALUES (?, NULL)", + (version,), + ) + conn.commit() + return cur.rowcount == 1 + + +def _mark_complete(conn: sqlite3.Connection, version: str) -> None: + conn.execute( + "UPDATE schema_migrations SET applied_at = CURRENT_TIMESTAMP " + "WHERE version = ?", + (version,), + ) + conn.commit() + + +def _release_claim(conn: sqlite3.Connection, version: str) -> None: + """Best-effort delete of the claim row on DDL failure.""" + try: + conn.execute( + "DELETE FROM schema_migrations WHERE version = ?", (version,), + ) + conn.commit() + except sqlite3.Error: + pass + + +def _wait_for_peer(conn: sqlite3.Connection, version: str) -> str: + """Poll a peer's claim on ``version``. + + Returns: + * ``"completed"`` — peer set ``applied_at`` non-NULL; caller skips. + * ``"released"`` — peer deleted the claim (DDL failed); caller + should attempt to claim it themselves. + + Raises ``sqlite3.OperationalError`` on timeout. Timing out is + intentional: it turns "peer died mid-migration" into a loud error + at the losing process rather than a silent return with a partial + schema (the regression BugBot flagged on the first cut of #106). + """ + deadline = time.monotonic() + _CLAIM_WAIT_SECONDS + while time.monotonic() < deadline: + row = conn.execute( + "SELECT applied_at FROM schema_migrations WHERE version = ?", + (version,), + ).fetchone() + if row is None: + return "released" + if row[0] is not None: + return "completed" + time.sleep(_CLAIM_POLL_INTERVAL) + raise sqlite3.OperationalError( + f"Timed out waiting for peer to apply migration {version} " + f"after {_CLAIM_WAIT_SECONDS:.0f}s; a previous starter likely " + f"died mid-migration. Inspect schema_migrations and the DB " + f"before retrying." + ) + + +def _apply_one( + conn: sqlite3.Connection, + version: str, + sql_file: Path, +) -> bool: + """Apply a single migration, coordinating with any concurrent peer. + + Returns True iff we actually ran the DDL (i.e. we won the claim); + False if a peer had already completed it. Raises on DDL failure or + peer-wait timeout. + """ + while True: + if _try_claim(conn, version): + # Everything after the claim lands in the same try/except so a + # failure in read_text (OSError / UnicodeDecodeError), the DDL + # itself, OR _mark_complete releases the claim row. Otherwise + # the row would be leaked with applied_at IS NULL and every + # future start would poll _CLAIM_WAIT_SECONDS then raise — + # permanently blocking the version until someone hand-cleans + # schema_migrations. + # + # ``executescript`` issues its own COMMIT before running, so we + # cannot wrap it in an outer BEGIN. On failure, SQLite + # auto-rolls back the individual failing statement; a partial + # init is equivalent to a corrupt fresh DB — discard and retry. + try: + sql = sql_file.read_text(encoding="utf-8") + conn.executescript(sql) + _mark_complete(conn, version) + except Exception: + _release_claim(conn, version) + try: + conn.rollback() + except sqlite3.Error: + pass + raise + return True + + # Lost the claim; wait for the peer to finish or release. + outcome = _wait_for_peer(conn, version) + if outcome == "completed": + return False + # Peer released (DDL failed). Loop to try claiming ourselves. + + def apply_migrations( conn: sqlite3.Connection, migrations_dir: Path | None = None, @@ -72,29 +218,7 @@ def apply_migrations( version = _version_from_filename(sql_file) if version in applied: continue - - sql = sql_file.read_text(encoding="utf-8") - - # ``executescript`` issues its own COMMIT before running, so we cannot - # wrap it in an outer BEGIN. On failure, SQLite auto-rolls back the - # individual failing statement; the caller can inspect partial state or - # recreate the DB. For the init migration this is acceptable because a - # partial init is equivalent to a corrupt fresh DB — discard and retry. - try: - conn.executescript(sql) - except Exception: - # Defensive: if anything is left pending, clean up. - try: - conn.rollback() - except sqlite3.Error: - pass - raise - - conn.execute( - "INSERT INTO schema_migrations (version) VALUES (?)", - (version,), - ) - conn.commit() - applied_now.append(version) + if _apply_one(conn, version, sql_file): + applied_now.append(version) return applied_now diff --git a/tests/db/test_schema.py b/tests/db/test_schema.py index 52e0a35..09082c9 100644 --- a/tests/db/test_schema.py +++ b/tests/db/test_schema.py @@ -238,6 +238,272 @@ def test_apply_migrations_is_idempotent(tmp_memory_db: Path) -> None: conn.close() +def test_apply_migrations_skips_versions_recorded_after_snapshot( + tmp_memory_db: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Concurrent-race regression for #106. + + Reproduces the window between the initial ``_applied_versions`` snapshot + and the pre-claim ``INSERT``: a peer completed the version row (both + ``applied_at IS NOT NULL`` and the DDL) between our snapshot and our + claim. The ``INSERT OR IGNORE`` must find a conflict and the poll must + see ``applied_at IS NOT NULL`` and skip — otherwise executescript + would run ``CREATE TABLE ... already exists`` and crash startup. + """ + from better_memory.db import schema as schema_mod + + migrations_dir = tmp_path / "migs" + migrations_dir.mkdir() + (migrations_dir / "0001_race.sql").write_text( + "CREATE TABLE race_table (id INTEGER PRIMARY KEY);" + ) + + conn = connect(tmp_memory_db) + try: + schema_mod._ensure_schema_migrations_table(conn) + conn.execute("CREATE TABLE race_table (id INTEGER PRIMARY KEY)") + # Peer completed: applied_at explicitly non-NULL. + conn.execute( + "INSERT INTO schema_migrations (version, applied_at) " + "VALUES ('0001', CURRENT_TIMESTAMP)" + ) + conn.commit() + monkeypatch.setattr( + schema_mod, "_applied_versions", lambda _conn: set() + ) + + applied = apply_migrations(conn, migrations_dir=migrations_dir) + assert applied == [] + finally: + conn.close() + + +def test_apply_migrations_releases_claim_on_ddl_failure( + tmp_memory_db: Path, tmp_path: Path, +) -> None: + """A DDL failure must remove the pre-claim row so the next run retries.""" + migrations_dir = tmp_path / "migs" + migrations_dir.mkdir() + (migrations_dir / "0001_bad.sql").write_text( + "CREATE TABLE valid (id INTEGER PRIMARY KEY); " + "THIS IS NOT SQL;" + ) + + conn = connect(tmp_memory_db) + try: + with pytest.raises(sqlite3.OperationalError): + apply_migrations(conn, migrations_dir=migrations_dir) + rows = conn.execute( + "SELECT version FROM schema_migrations WHERE version = '0001'" + ).fetchall() + assert rows == [] + finally: + conn.close() + + +def test_apply_migrations_waits_for_peer_and_retries_after_release( + tmp_memory_db: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Peer's DDL failed → loser must retry the claim, not silently skip. + + Regression against the BugBot finding on the first cut of #106: if a + peer holds an in-progress claim (``applied_at IS NULL``) and then + fails and deletes the row, a loser that had already skipped would + return with a broken schema. The two-phase protocol makes the loser + poll until the row disappears, then re-attempt the claim itself. + """ + from better_memory.db import schema as schema_mod + + migrations_dir = tmp_path / "migs" + migrations_dir.mkdir() + (migrations_dir / "0001_late.sql").write_text( + "CREATE TABLE late_table (id INTEGER PRIMARY KEY);" + ) + + conn = connect(tmp_memory_db) + try: + schema_mod._ensure_schema_migrations_table(conn) + # Peer holds an in-progress claim (applied_at NULL). + conn.execute( + "INSERT INTO schema_migrations (version, applied_at) " + "VALUES ('0001', NULL)" + ) + conn.commit() + monkeypatch.setattr( + schema_mod, "_applied_versions", lambda _conn: set() + ) + + # On the FIRST sleep, simulate the peer failing and releasing. + calls: list[float] = [] + + def fake_sleep(seconds: float) -> None: + calls.append(seconds) + if len(calls) == 1: + conn.execute( + "DELETE FROM schema_migrations WHERE version = '0001'" + ) + conn.commit() + + monkeypatch.setattr(schema_mod.time, "sleep", fake_sleep) + + applied = apply_migrations(conn, migrations_dir=migrations_dir) + # Loser detected the release, claimed the version, and ran the DDL. + assert applied == ["0001"] + assert calls, "loser must have polled while peer held the claim" + row = conn.execute( + "SELECT applied_at FROM schema_migrations WHERE version = '0001'" + ).fetchone() + assert row is not None + assert row[0] is not None # completed, not in-progress + # And the DDL actually ran. + tables = conn.execute( + "SELECT name FROM sqlite_master WHERE name = 'late_table'" + ).fetchall() + assert tables, "loser did not run the DDL after peer released" + finally: + conn.close() + + +def test_apply_migrations_times_out_on_stuck_peer( + tmp_memory_db: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A peer that holds the claim forever must produce a loud error. + + The whole point of the two-phase claim is that a losing process never + silently returns with a broken schema — a stuck peer surfaces as an + ``OperationalError``, not a partial-schema startup. + """ + from better_memory.db import schema as schema_mod + + migrations_dir = tmp_path / "migs" + migrations_dir.mkdir() + (migrations_dir / "0001_stuck.sql").write_text( + "CREATE TABLE stuck_table (id INTEGER PRIMARY KEY);" + ) + + conn = connect(tmp_memory_db) + try: + schema_mod._ensure_schema_migrations_table(conn) + conn.execute( + "INSERT INTO schema_migrations (version, applied_at) " + "VALUES ('0001', NULL)" + ) + conn.commit() + monkeypatch.setattr( + schema_mod, "_applied_versions", lambda _conn: set() + ) + # Shrink the wait so the test is fast; peer never transitions. + monkeypatch.setattr(schema_mod, "_CLAIM_WAIT_SECONDS", 0.05) + monkeypatch.setattr(schema_mod, "_CLAIM_POLL_INTERVAL", 0.01) + + with pytest.raises(sqlite3.OperationalError) as excinfo: + apply_migrations(conn, migrations_dir=migrations_dir) + assert "0001" in str(excinfo.value) + finally: + conn.close() + + +def test_apply_migrations_releases_claim_when_sql_read_fails( + tmp_memory_db: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A read_text failure after claiming must not leak the claim row. + + Otherwise the version would poll for the full timeout on every future + start until someone hand-cleans schema_migrations. + """ + migrations_dir = tmp_path / "migs" + migrations_dir.mkdir() + sql_file = migrations_dir / "0001_unreadable.sql" + sql_file.write_text("CREATE TABLE ok (id INTEGER PRIMARY KEY);") + + original_read_text = Path.read_text + + def failing_read_text(self, *args, **kwargs): + if self == sql_file: + raise OSError("simulated read failure") + return original_read_text(self, *args, **kwargs) + + monkeypatch.setattr(Path, "read_text", failing_read_text) + + conn = connect(tmp_memory_db) + try: + with pytest.raises(OSError, match="simulated read failure"): + apply_migrations(conn, migrations_dir=migrations_dir) + rows = conn.execute( + "SELECT version FROM schema_migrations WHERE version = '0001'" + ).fetchall() + assert rows == [], "claim row must be released after read failure" + finally: + conn.close() + + +def test_apply_migrations_releases_claim_when_mark_complete_fails( + tmp_memory_db: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A _mark_complete failure after successful DDL must not leak the claim. + + If the completion UPDATE raises (transient DB error, thread issue), + the claim row would otherwise sit at applied_at IS NULL forever. + """ + from better_memory.db import schema as schema_mod + + migrations_dir = tmp_path / "migs" + migrations_dir.mkdir() + (migrations_dir / "0001_ok.sql").write_text( + "CREATE TABLE marked_ok (id INTEGER PRIMARY KEY);" + ) + + def failing_mark(_conn, _version): + raise sqlite3.OperationalError("simulated mark failure") + + monkeypatch.setattr(schema_mod, "_mark_complete", failing_mark) + + conn = connect(tmp_memory_db) + try: + with pytest.raises(sqlite3.OperationalError, match="simulated mark"): + apply_migrations(conn, migrations_dir=migrations_dir) + rows = conn.execute( + "SELECT version FROM schema_migrations WHERE version = '0001'" + ).fetchall() + assert rows == [], "claim row must be released after mark failure" + finally: + conn.close() + + +def test_applied_versions_excludes_in_progress_claims( + tmp_memory_db: Path, +) -> None: + """``_applied_versions`` must not count claims where ``applied_at IS NULL``. + + Otherwise a peer's live claim would look like a completed apply, and + the whole polling protocol would collapse to the pre-fix "trust the + bare row" behaviour. + """ + from better_memory.db import schema as schema_mod + + conn = connect(tmp_memory_db) + try: + schema_mod._ensure_schema_migrations_table(conn) + conn.execute( + "INSERT INTO schema_migrations (version, applied_at) " + "VALUES ('0001', CURRENT_TIMESTAMP), ('0002', NULL)" + ) + conn.commit() + assert schema_mod._applied_versions(conn) == {"0001"} + finally: + conn.close() + + def test_episodic_indexes_exist(tmp_memory_db: Path) -> None: """The two episodic indexes are created.""" conn = connect(tmp_memory_db)