diff --git a/tests/browser/fast_reset.py b/tests/browser/fast_reset.py index 2e4b786b..49f6ac4b 100644 --- a/tests/browser/fast_reset.py +++ b/tests/browser/fast_reset.py @@ -130,7 +130,7 @@ class ScratchSeeder: raise RuntimeError("the seed changed the database schema") if live.execute("PRAGMA data_version").fetchone()[0] != version: raise RuntimeError("the live database changed while the seed ran") - self._publish(live, scratch.model.db_path, changed) + self._publish(live, scratch.model.db_path, changed, version) return result def _scratch_database(self): @@ -150,33 +150,48 @@ class ScratchSeeder: return conn @staticmethod - def _publish(live: sqlite3.Connection, path: str, tables: list[str]) -> None: + def _publish( + live: sqlite3.Connection, path: str, tables: list[str], version: int + ) -> None: + """Replace ``tables`` with the scratch rows in one live transaction. + + ``version`` is the live ``data_version`` read before the seed. It is + checked again under the write lock: a writer whose transaction was + open during the seed commits before ``BEGIN IMMEDIATE`` returns, and + its rows must fail the reset instead of being overwritten. + """ if not tables: return foreign_keys = live.execute("PRAGMA foreign_keys").fetchone()[0] # Rows are copied as they are; no cascade may reach a table left alone. live.execute("PRAGMA foreign_keys = OFF") - live.execute(f"ATTACH DATABASE ? AS {SCRATCH_SCHEMA}", (path,)) try: - live.execute("BEGIN IMMEDIATE") + live.execute(f"ATTACH DATABASE ? AS {SCRATCH_SCHEMA}", (path,)) try: - for table in tables: - columns = ", ".join( - _quote(c) if c != "rowid" else "rowid" - for c in _columns(live, table) - ) - name = _quote(table) - live.execute(f"DELETE FROM main.{name}") - live.execute( - f"INSERT INTO main.{name} ({columns}) " - f"SELECT {columns} FROM {SCRATCH_SCHEMA}.{name}" - ) - live.commit() - except BaseException: - live.rollback() - raise + live.execute("BEGIN IMMEDIATE") + try: + if live.execute("PRAGMA data_version").fetchone()[0] != version: + raise RuntimeError( + "the live database changed while the seed ran" + ) + for table in tables: + columns = ", ".join( + _quote(c) if c != "rowid" else "rowid" + for c in _columns(live, table) + ) + name = _quote(table) + live.execute(f"DELETE FROM main.{name}") + live.execute( + f"INSERT INTO main.{name} ({columns}) " + f"SELECT {columns} FROM {SCRATCH_SCHEMA}.{name}" + ) + live.commit() + except BaseException: + live.rollback() + raise + finally: + live.execute(f"DETACH DATABASE {SCRATCH_SCHEMA}") finally: - live.execute(f"DETACH DATABASE {SCRATCH_SCHEMA}") live.execute(f"PRAGMA foreign_keys = {int(foreign_keys)}") def close(self) -> None: diff --git a/tests/test_browser_fast_reset.py b/tests/test_browser_fast_reset.py index 27cc316a..61834d8d 100644 --- a/tests/test_browser_fast_reset.py +++ b/tests/test_browser_fast_reset.py @@ -3,9 +3,12 @@ import hashlib import sqlite3 import tempfile +import threading +import time import types import unittest from pathlib import Path +from unittest import mock try: from tests.open_handles import assert_closed @@ -13,6 +16,7 @@ except ImportError: # discovered with tests/ as the top-level directory from open_handles import assert_closed from database.operations import PromptDatabase +from tests.browser import fast_reset from tests.browser.fast_reset import ScratchSeeder, table_digests @@ -162,6 +166,59 @@ class TestScratchSeeder(unittest.TestCase): self.assertIn("written to live meanwhile", texts) self.assertNotIn("seeded 0", texts) + def test_a_write_open_during_the_seed_and_committed_later_fails_it(self): + other = self.database("live.db") + began, finished = threading.Event(), [] + + def write(): + with other.model.get_connection() as conn: + conn.execute("BEGIN IMMEDIATE") + conn.execute( + "INSERT INTO prompts (text, hash) VALUES ('concurrent', 'c')" + ) + began.set() + time.sleep(0.5) # the publish waits on this lock meanwhile + other.model.close() + finished.append(True) + + writer = threading.Thread(target=write) + + def racing(api): + reseed(api) + writer.start() + self.assertTrue(began.wait(10)) + + with self.assertRaisesRegex(RuntimeError, "changed while"): + self.seeder.run(racing) + writer.join() + self.assertEqual(finished, [True]) + with self.live.model.get_connection() as conn: + texts = [row[0] for row in conn.execute("SELECT text FROM prompts")] + self.assertIn("concurrent", texts) + self.assertNotIn("seeded 0", texts) + self.assert_live_connection_restored() + + def test_a_failed_publish_leaves_live_and_its_connection_as_they_were(self): + before = digests(self.live) + columns = fast_reset._columns + + def broken_columns(conn, table, *schema): + # The publish asks for the live table's columns without a schema. + return columns(conn, table, *schema) if schema else ["no_such_column"] + + with mock.patch.object(fast_reset, "_columns", broken_columns): + with self.assertRaises(sqlite3.OperationalError): + self.seeder.run(reseed) + self.assertEqual(digests(self.live), before) + self.assert_live_connection_restored() + + def assert_live_connection_restored(self): + conn = self.live.model.get_connection() + self.assertFalse(conn.in_transaction) + self.assertEqual(conn.execute("PRAGMA foreign_keys").fetchone()[0], 1) + names = [row[1] for row in conn.execute("PRAGMA database_list")] + self.assertEqual(names, ["main"]) + def test_the_live_connection_is_left_as_it_was(self): self.seeder.run(reseed) conn = self.live.model.get_connection()