nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Vito Sansevero
2026-10-18 10:01:35 -07:00
parent ac23f22a37
commit 6e2fdfacaa
3 changed files with 227 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
database/models.py
+7 -18
View File
@@ -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:
``<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:
tests/test_api_admin.py
+72
View File
@@ -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())
tests/test_database_backup.py
+148
View File
@@ -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"])