diff --git a/database/models.py b/database/models.py index 6383949..02ffbc1 100644 --- a/database/models.py +++ b/database/models.py @@ -124,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() @@ -550,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 @@ -1123,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 b060416..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() @@ -1506,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 54e4885..e82011f 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -318,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"])