From 6e2fdfacaa346e095681e11093ffed3e86673e96 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sun, 18 Oct 2026 09:34:07 -0700 Subject: [PATCH] fix(database): refuse unsafe restore file fallback (#223) --- database/models.py | 25 ++---- tests/test_api_admin.py | 72 +++++++++++++++++ tests/test_database_backup.py | 148 ++++++++++++++++++++++++++++++++++ 3 files changed, 227 insertions(+), 18 deletions(-) diff --git a/database/models.py b/database/models.py index ff0bb9a..6383949 100644 --- a/database/models.py +++ b/database/models.py @@ -6,7 +6,6 @@ import contextlib import datetime import json import math -import shutil import sqlite3 import os import threading @@ -1061,29 +1060,17 @@ class PromptModel: PromptModel instances that share the file (nodes, image monitor, queue hook, API) simply see the new pages on their next statement. Nothing is swapped underneath open handles and no -wal/-shm file has - to be deleted, which also keeps Windows happy. When SQLite refuses - (page sizes differ while the target is in WAL mode) fall back to a - file copy with every connection of this instance closed first. + to be deleted, which also keeps Windows happy. If SQLite refuses + replacement (for example, incompatible page sizes in WAL mode), + propagate the error rather than copy files under other open owners. """ source = sqlite3.connect( Path(src_path).resolve().as_uri() + "?mode=ro", uri=True ) try: - try: - source.backup(self.get_connection()) - return - except sqlite3.OperationalError as e: - self.logger.warning( - f"Online restore failed ({e}); falling back to a file copy" - ) + source.backup(self.get_connection()) finally: source.close() - self.close_all() - shutil.copy2(src_path, self.db_path) - for suffix in ("-wal", "-shm"): - stale = self.db_path + suffix - if os.path.exists(stale): - os.remove(stale) def restore_from_file(self, src_path: str) -> str: """ @@ -1094,6 +1081,8 @@ class PromptModel: ``.backup_`` (WAL content included), and the staged pages replace the live pages through SQLite's backup API. Migration failures never replace live content or modify the source. + An online replacement failure leaves the live database intact and + propagates; the completed safety backup remains available. Args: src_path: SQLite file to restore @@ -1104,7 +1093,7 @@ class PromptModel: Raises: RuntimeError: If the previous database could not be backed up - OSError: If the restored file could not be copied into place + sqlite3.OperationalError: If SQLite refuses online replacement """ ok, reason = self.verify_database_file(src_path) if not ok: diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index 1f91093..b060416 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -1425,6 +1425,78 @@ class TestBackupRestoreContract(AdminAPITestCase): conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0], 1 ) + async def test_incompatible_wal_restore_returns_failure_without_publication( + self, + ) -> None: + import sqlite3 + from contextlib import closing + + from database.execution_history import subscribe_history_restore + + self.api.db.save_prompt(text="keep live", prompt_hash="live") + previous = self.api.db + live = previous.model.get_connection() + epoch = previous.get_execution_epoch() + self.assertEqual(live.execute("PRAGMA page_size").fetchone()[0], 4096) + self.assertEqual(live.execute("PRAGMA journal_mode").fetchone()[0], "wal") + compatible = self._snapshot_upload() + source_path = Path(self.tmpdir) / "large-pages.db" + with closing(sqlite3.connect(source_path)) as conn: + conn.execute("PRAGMA page_size=8192") + conn.execute("VACUUM") + source = PromptDatabase(str(source_path)) + try: + source.save_prompt(text="do not publish", prompt_hash="source") + self.assertEqual( + source.model.get_connection().execute("PRAGMA page_size").fetchone()[0], + 8192, + ) + finally: + source.close_all() + upload = source_path.read_bytes() + notices = [] + self.addCleanup( + subscribe_history_restore( + previous.model.db_path, lambda: notices.append(True) + ) + ) + + resp = await self._restore("database_file", upload) + + body = await resp.text() + self.assertEqual(resp.status, 500, body) + data = json.loads(body) + self.assertFalse(data["success"]) + self.assertIn("Failed to restore database", data["error"]) + self.assertLess(len(data["error"]), 300) + self.assertNotIn(self.tmpdir, body) + self.assertNotIn("/tmp/", body) + self.assertNotIn("backup_created", data) + self.assertIs(self.api.db, previous) + self.assertEqual(previous.get_execution_epoch(), epoch) + self.assertEqual(notices, []) + self.assertEqual( + live.execute("SELECT text FROM prompts").fetchone()[0], "keep live" + ) + previous.save_prompt(text="still writable", prompt_hash="after") + self.assertEqual(source_path.read_bytes(), upload) + backups = list(Path(self.tmpdir).glob("prompts.db.backup_*")) + self.assertEqual(len(backups), 1) + with closing(sqlite3.connect(backups[0])) as conn: + self.assertEqual( + conn.execute("SELECT text FROM prompts").fetchone()[0], "keep live" + ) + self.assertEqual(conn.execute("PRAGMA integrity_check").fetchone()[0], "ok") + + retry = await self._restore("database_file", compatible) + self.assertEqual(retry.status, 200, await retry.text()) + self.assertTrue((await retry.json())["success"]) + self.assertEqual(notices, [True]) + self.assertNotEqual(self.api.db.get_execution_epoch(), epoch) + self.assertEqual( + [row["text"] for row in self.api.db.search_prompts()], ["keep live"] + ) + async def test_restore_backup_created_is_a_public_path(self): self.api.db.save_prompt(text="restored", prompt_hash="h1") resp = await self._restore("database_file", self._snapshot_upload()) diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 79f31b4..54e4885 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -12,6 +12,9 @@ import sys import tempfile import threading import unittest +from contextlib import closing +from pathlib import Path +from unittest.mock import patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -156,6 +159,151 @@ class TestVerifyDatabaseFile(BackupTestCase): class TestRestoreFromFile(BackupTestCase): + def test_incompatible_wal_restore_preserves_other_owner_and_allows_retry( + self, + ) -> None: + from database.execution_history import subscribe_history_restore + + self._save(self.db, "keep live") + live = self.db.model.get_connection() + self.assertEqual(live.execute("PRAGMA journal_mode").fetchone()[0], "wal") + self.assertEqual(live.execute("PRAGMA page_size").fetchone()[0], 4096) + epoch = self.db.get_execution_epoch() + compatible = self._make_source_db(["compatible retry"]) + source_path = str(Path(self.tmpdir) / "large-pages.db") + with closing(sqlite3.connect(source_path)) as conn: + conn.execute("PRAGMA page_size=8192") + conn.execute("VACUUM") + source = PromptDatabase(source_path) + try: + self._save(source, "incompatible source") + self.assertEqual( + source.model.get_connection().execute("PRAGMA page_size").fetchone()[0], + 8192, + ) + finally: + source.close_all() + original_source = Path(source_path).read_bytes() + notices = [] + self.addCleanup( + subscribe_history_restore(self.path, lambda: notices.append(True)) + ) + unrelated = PromptDatabase(str(Path(self.tmpdir) / "unrelated.db")) + self.addCleanup(unrelated.close_all) + self._save(unrelated, "unrelated") + unrelated_connection = unrelated.model.get_connection() + bystander = PromptDatabase(self.path) + self.addCleanup(bystander.close_all) + ready, release, refused = (threading.Event() for _ in range(3)) + results = {} + + def hold_reader_then_write() -> None: + try: + conn = bystander.model.get_connection() + conn.execute("BEGIN") + results["before"] = conn.execute("SELECT text FROM prompts").fetchone()[ + 0 + ] + ready.set() + if not release.wait(10): + raise TimeoutError("Restore did not release held reader") + results["held"] = conn.execute("SELECT text FROM prompts").fetchone()[0] + conn.rollback() + # Never write through the original unsafe fallback during RED. + if refused.is_set(): + self._save(bystander, "other owner still writes") + results["after"] = conn.execute( + "SELECT text FROM prompts ORDER BY id" + ).fetchall() + except Exception as exc: + results["error"] = exc + finally: + ready.set() + bystander.close() + + worker = threading.Thread(target=hold_reader_then_write) + worker.start() + try: + self.assertTrue(ready.wait(10), "Held reader did not start") + self.assertNotIn("error", results) + with self.assertRaises(sqlite3.OperationalError): + try: + self.db.model.restore_from_file(source_path) + except sqlite3.OperationalError: + refused.set() + raise + finally: + release.set() + worker.join(10) + + self.assertFalse(worker.is_alive(), "Held reader did not finish") + self.assertNotIn("error", results) + self.assertEqual(results["before"], "keep live") + self.assertEqual(results["held"], "keep live") + self.assertEqual( + [row[0] for row in results["after"]], + ["keep live", "other owner still writes"], + ) + self.assertEqual(self.db.get_execution_epoch(), epoch) + self.assertEqual(notices, []) + self.assertEqual(_texts(self.path), ["keep live", "other owner still writes"]) + self._save(self.db, "restoring owner still writes") + self.assertEqual(live.execute("PRAGMA integrity_check").fetchone()[0], "ok") + self.assertEqual( + unrelated_connection.execute("SELECT text FROM prompts").fetchone()[0], + "unrelated", + ) + self.assertEqual(Path(source_path).read_bytes(), original_source) + backups = list(Path(self.tmpdir).glob("prompts.db.backup_*")) + self.assertEqual(len(backups), 1) + self.assertEqual(_texts(str(backups[0])), ["keep live"]) + for path in (source_path, str(backups[0])): + with closing(sqlite3.connect(path)) as conn: + self.assertEqual( + conn.execute("PRAGMA integrity_check").fetchone()[0], "ok" + ) + + self.db.model.restore_from_file(compatible) + self.assertEqual(_texts(self.path), ["compatible retry"]) + self.assertNotEqual(self.db.get_execution_epoch(), epoch) + self.assertEqual(notices, [True]) + + def test_online_operational_error_preserves_live_data_and_closes_source( + self, + ) -> None: + self._save(self.db, "keep live") + source = self._make_source_db(["do not publish"]) + epoch = self.db.get_execution_epoch() + live = self.db.model.get_connection() + connect = sqlite3.connect + failed_sources = [] + live_path = os.path.realpath(self.path) + + class FailingLiveBackup(sqlite3.Connection): + def backup(self, target: sqlite3.Connection, *args, **kwargs) -> None: + target_path = target.execute("PRAGMA database_list").fetchone()[2] + if os.path.realpath(target_path) == live_path: + failed_sources.append(self) + raise sqlite3.OperationalError("injected live backup I/O failure") + super().backup(target, *args, **kwargs) + + def connect_with_failure(*args, **kwargs) -> sqlite3.Connection: + return connect(*args, **kwargs, factory=FailingLiveBackup) + + with patch("database.models.sqlite3.connect", side_effect=connect_with_failure): + with self.assertRaises(sqlite3.OperationalError): + self.db.model.restore_from_file(source) + + self.assertEqual(_texts(self.path), ["keep live"]) + self.assertEqual(self.db.get_execution_epoch(), epoch) + self.assertEqual( + live.execute("SELECT text FROM prompts").fetchone()[0], "keep live" + ) + self.assertEqual(_texts(source), ["do not publish"]) + self.assertEqual(len(failed_sources), 1) + with self.assertRaises(sqlite3.ProgrammingError): + failed_sources[0].execute("SELECT 1") + def test_restore_replaces_content_and_backs_up_the_old_database(self): self._save(self.db, "old prompt") src = self._make_source_db(["restored one", "restored two"])