fix(database): refuse unsafe restore file fallback (#226)

* fix(database): refuse unsafe restore file fallback (#223)

* fix(database): retain safety backup limit after restore refusal (#223)

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:
         ``<db_path>.backup_<YYYYmmdd_HHMMSS>`` (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}")

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")

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(

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"])