diff --git a/database/models.py b/database/models.py index ff0bb9a..02ffbc1 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 @@ -125,7 +124,7 @@ class PromptModel: self.logger.debug(f"Initializing database model with path: {db_path}") self.db_path = db_path # One connection per thread (see get_connection); every open connection - # is also tracked here so restore_from_file can close them all. + # is also tracked here so close_all can clean up across threads. self._local = threading.local() self._connections: Set[sqlite3.Connection] = set() self._conn_lock = threading.Lock() @@ -551,7 +550,7 @@ class PromptModel: if conn is not None and self._local.generation == self._generation: return conn - # check_same_thread=False only so restore_from_file can close other + # check_same_thread=False only so close_all can clean up other # threads' connections; each connection is still used by one thread. conn = sqlite3.connect( self.db_path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS @@ -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: @@ -1134,7 +1123,11 @@ class PromptModel: if not self.backup_database(backup_path): raise RuntimeError("Could not back up the current database") self.logger.info(f"Current database backed up to: {backup_path}") - self._replace_contents_from(staged_path) + try: + self._replace_contents_from(staged_path) + except sqlite3.OperationalError as exc: + exc._prompt_manager_restore_backup_created = bool(backup_path) + raise with PromptModel._init_lock: PromptModel._initialized_paths.add(self._schema_key()) self.logger.info(f"Database restored from {src_path}") diff --git a/py/api/admin.py b/py/api/admin.py index 206792e..0f3eb88 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -1248,9 +1248,24 @@ class AdminRoutesMixin: {"success": False, "error": reason}, status=400 ) - backup_path = await self._run_in_executor( - model.restore_from_file, temp_path - ) + try: + backup_path = await self._run_in_executor( + model.restore_from_file, temp_path + ) + except sqlite3.OperationalError as exc: + if ( + getattr(exc, "_prompt_manager_restore_backup_created", False) + is True + ): + try: + await self._run_in_executor( + _prune_safety_backups, model.db_path + ) + except Exception: + self.logger.warning( + "Could not prune safety backups after failed restore" + ) + raise removed = await self._run_in_executor(_prune_safety_backups, model.db_path) if removed: self.logger.info(f"Removed {len(removed)} old safety backups") diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index 1f91093..96329a1 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -13,7 +13,7 @@ import tempfile import types import unittest from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch _mock_server = MagicMock() _mock_server.PromptServer.instance.routes = MagicMock() @@ -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()) @@ -1434,6 +1506,159 @@ class TestBackupRestoreContract(AdminAPITestCase): data = json.loads(body) self.assertTrue(data["backup_created"].startswith("prompts.db.backup_")) + def _retention_upload(self) -> bytes: + import sqlite3 + from contextlib import closing + + path = Path(self.tmpdir) / "retention-source.db" + with closing(sqlite3.connect(path)) as conn: + conn.execute("PRAGMA page_size=8192") + conn.execute("VACUUM") + source = PromptDatabase(str(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() + return path.read_bytes() + + def _seed_retention_backups(self, count: int) -> list[Path]: + snapshot = self._snapshot_upload() + paths = [] + for index in range(count): + path = Path( + f"{self.api.db.model.db_path}.backup_202001{index + 1:02d}_000000" + ) + path.write_bytes(snapshot) + os.utime(path, (1_000_000 + index, 1_000_000 + index)) + paths.append(path) + return paths + + async def test_refused_restore_retains_ten_backups_and_new_recovery_copy( + self, + ) -> None: + import sqlite3 + from contextlib import closing + + from database.execution_history import subscribe_history_restore + + previous = self.api.db + previous.save_prompt(text="keep live", prompt_hash="live") + live = previous.model.get_connection() + self.assertEqual(live.execute("PRAGMA page_size").fetchone()[0], 4096) + self.assertEqual(live.execute("PRAGMA journal_mode").fetchone()[0], "wal") + epoch = previous.get_execution_epoch() + old_backups = self._seed_retention_backups(10) + upload = self._retention_upload() + notices = [] + self.addCleanup( + subscribe_history_restore( + previous.model.db_path, lambda: notices.append(True) + ) + ) + + response = await self._restore("database_file", upload) + + data = await response.json() + self.assertEqual(response.status, 500, data) + self.assertFalse(data["success"]) + self.assertNotIn(self.tmpdir, data["error"]) + remaining = set(Path(self.tmpdir).glob("prompts.db.backup_*")) + self.assertEqual(len(remaining), 10) + self.assertEqual(remaining.intersection(old_backups), set(old_backups[1:])) + new_backups = remaining.difference(old_backups) + self.assertEqual(len(new_backups), 1) + with closing(sqlite3.connect(new_backups.pop())) as conn: + self.assertEqual( + conn.execute("SELECT text FROM prompts").fetchone()[0], "keep live" + ) + self.assertEqual(conn.execute("PRAGMA integrity_check").fetchone()[0], "ok") + 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 writes", prompt_hash="after") + self.assertEqual( + (Path(self.tmpdir) / "retention-source.db").read_bytes(), upload + ) + + async def test_prebackup_restore_failures_leave_existing_backups_untouched( + self, + ) -> None: + import sqlite3 + + self.api.db.save_prompt(text="keep live", prompt_hash="live") + old_backups = self._seed_retention_backups(11) + original = {path: path.read_bytes() for path in old_backups} + epoch = self.api.db.get_execution_epoch() + invalid = await self._restore("database_file", b"not a database") + self.assertEqual(invalid.status, 400) + self.assertFalse((await invalid.json())["success"]) + compatible = self._snapshot_upload() + with patch( + "database.models.rotate_history_epoch", + side_effect=sqlite3.OperationalError("staging failed before safety backup"), + ): + staged = await self._restore("database_file", compatible) + self.assertEqual(staged.status, 500) + self.assertIn("staging failed", (await staged.json())["error"]) + self.assertEqual( + set(Path(self.tmpdir).glob("prompts.db.backup_*")), set(old_backups) + ) + self.assertEqual({path: path.read_bytes() for path in old_backups}, original) + self.assertEqual(self.api.db.get_execution_epoch(), epoch) + self.assertIsNotNone(self.api.db.get_prompt_by_hash("live")) + + async def test_retention_failure_keeps_original_restore_error_and_backup( + self, + ) -> None: + import sqlite3 + from contextlib import closing + + self.api.db.save_prompt(text="keep live", prompt_hash="live") + previous = self.api.db + epoch = previous.get_execution_epoch() + old_backups = self._seed_retention_backups(10) + upload = self._retention_upload() + with patch( + "py.api.admin._prune_safety_backups", + side_effect=OSError("retention cleanup failure /tmp/private"), + ): + with self.assertLogs(self.api.logger, level="WARNING") as logs: + response = await self._restore("database_file", upload) + data = await response.json() + self.assertEqual(response.status, 500, data) + self.assertFalse(data["success"]) + self.assertIn("readonly database", data["error"]) + self.assertNotIn("retention cleanup failure", data["error"]) + self.assertNotIn("/tmp/", data["error"]) + warnings = [ + record.getMessage() + for record in logs.records + if record.levelname == "WARNING" + ] + self.assertTrue(any("backup" in message.lower() for message in warnings)) + self.assertTrue( + all("/tmp/" not in message and len(message) < 300 for message in warnings) + ) + self.assertIs(self.api.db, previous) + self.assertEqual(previous.get_execution_epoch(), epoch) + self.assertIsNotNone(previous.get_prompt_by_hash("live")) + remaining = set(Path(self.tmpdir).glob("prompts.db.backup_*")) + self.assertEqual(len(remaining), 11) + new_backups = remaining.difference(old_backups) + self.assertEqual(len(new_backups), 1) + with closing(sqlite3.connect(new_backups.pop())) as conn: + self.assertEqual( + conn.execute("SELECT text FROM prompts").fetchone()[0], "keep live" + ) + self.assertEqual(conn.execute("PRAGMA integrity_check").fetchone()[0], "ok") + async def test_restore_rejects_database_with_a_trigger(self): self.api.db.save_prompt(text="keep me", prompt_hash="h1") upload = self._snapshot_upload( diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 79f31b4..e82011f 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"]) @@ -170,6 +318,25 @@ class TestRestoreFromFile(BackupTestCase): # being deleted under them; the content is what matters. self.assertEqual(_texts(self.path), ["restored one", "restored two"]) + def test_refusal_preserves_exception_identity_and_marks_completed_backup( + self, + ) -> None: + self._save(self.db, "keep live") + source = self._make_source_db(["replacement"]) + failure = sqlite3.OperationalError("online replacement refused") + with patch.object(self.db.model, "_replace_contents_from", side_effect=failure): + with self.assertRaises(sqlite3.OperationalError) as raised: + self.db.model.restore_from_file(source) + self.assertIs(raised.exception, failure) + self.assertIs( + getattr(raised.exception, "_prompt_manager_restore_backup_created", None), + True, + ) + self.assertEqual(_texts(self.path), ["keep live"]) + backups = list(Path(self.tmpdir).glob("prompts.db.backup_*")) + self.assertEqual(len(backups), 1) + self.assertEqual(_texts(str(backups[0])), ["keep live"]) + def test_database_is_usable_from_every_thread_after_restore(self): self._save(self.db, "old prompt") src = self._make_source_db(["restored"])