Vito Sansevero
507c8a8993
Each owned test that opens a PromptDatabase now calls close_all() before removing the file, and temp directories holding a database use ignore_cleanup_errors, so the suite tears down cleanly on Windows. The image API fixture lays out a fake ComfyUI tree (<tmp>/output) and pins GalleryConfig.path_anchors to it so responses can be checked for root-relative paths.
302 lines
11 KiB
Python
302 lines
11 KiB
Python
"""Backup, verification and restore of the prompts database.
|
|
|
|
The database runs in WAL mode, so committed rows can sit in prompts.db-wal
|
|
for a long time. A backup must include them, a restore must reject files
|
|
that are not healthy SQLite databases, and the live database must stay
|
|
usable afterwards.
|
|
"""
|
|
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from database.models import PromptModel
|
|
from database.operations import PromptDatabase
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
|
|
def _count_prompts(path):
|
|
"""Row count read through a brand-new connection (closed before return)."""
|
|
conn = sqlite3.connect(path)
|
|
try:
|
|
return conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0]
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _texts(path):
|
|
conn = sqlite3.connect(path)
|
|
try:
|
|
return sorted(r[0] for r in conn.execute("SELECT text FROM prompts"))
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
class BackupTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmpdir = tempfile.mkdtemp()
|
|
self.path = os.path.join(self.tmpdir, "prompts.db")
|
|
self.db = PromptDatabase(self.path)
|
|
self.addCleanup(self._cleanup)
|
|
|
|
def _cleanup(self):
|
|
self.db.close_all()
|
|
for name in os.listdir(self.tmpdir):
|
|
os.unlink(os.path.join(self.tmpdir, name))
|
|
os.rmdir(self.tmpdir)
|
|
|
|
def _save(self, db, text):
|
|
return db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
|
|
|
def _make_source_db(self, texts):
|
|
"""A second database file holding the given prompts, fully closed."""
|
|
src_path = os.path.join(self.tmpdir, "source.db")
|
|
src = PromptDatabase(src_path)
|
|
for text in texts:
|
|
self._save(src, text)
|
|
src.close_all()
|
|
return src_path
|
|
|
|
|
|
class TestBackupDatabase(BackupTestCase):
|
|
def test_backup_contains_rows_that_only_exist_in_the_wal(self):
|
|
for n in range(5):
|
|
self._save(self.db, f"wal prompt {n}")
|
|
wal = self.path + "-wal"
|
|
self.assertTrue(os.path.exists(wal) and os.path.getsize(wal) > 0)
|
|
|
|
dst = os.path.join(self.tmpdir, "backup.db")
|
|
self.assertTrue(self.db.model.backup_database(dst))
|
|
|
|
self.assertEqual(_count_prompts(dst), 5)
|
|
self.assertFalse(os.path.exists(dst + "-wal"))
|
|
# The live database is untouched and still usable
|
|
self.assertEqual(len(self.db.search_prompts(text="wal prompt")), 5)
|
|
|
|
def test_backup_overwrites_a_stale_destination(self):
|
|
self._save(self.db, "one")
|
|
dst = os.path.join(self.tmpdir, "backup.db")
|
|
with open(dst, "wb") as fh:
|
|
fh.write(b"not a database at all")
|
|
self.assertTrue(self.db.model.backup_database(dst))
|
|
self.assertEqual(_count_prompts(dst), 1)
|
|
|
|
def test_backup_onto_the_live_file_is_refused(self):
|
|
self._save(self.db, "one")
|
|
self.assertFalse(self.db.model.backup_database(self.path))
|
|
self.assertEqual(len(self.db.search_prompts(text="one")), 1)
|
|
|
|
def test_backup_to_an_unwritable_location_returns_false(self):
|
|
dst = os.path.join(self.tmpdir, "missing", "dir", "backup.db")
|
|
self.assertFalse(self.db.model.backup_database(dst))
|
|
|
|
|
|
class TestVerifyDatabaseFile(BackupTestCase):
|
|
def test_text_file_is_rejected_without_raising(self):
|
|
path = os.path.join(self.tmpdir, "notes.txt")
|
|
with open(path, "w", encoding="utf-8") as fh:
|
|
fh.write("just some text, definitely not sqlite")
|
|
ok, reason = PromptModel.verify_database_file(path)
|
|
self.assertFalse(ok)
|
|
self.assertIn("SQLite", reason)
|
|
|
|
def test_missing_file_is_rejected(self):
|
|
ok, reason = PromptModel.verify_database_file(
|
|
os.path.join(self.tmpdir, "nope.db")
|
|
)
|
|
self.assertFalse(ok)
|
|
self.assertTrue(reason)
|
|
|
|
def test_database_without_prompts_table_is_rejected(self):
|
|
path = os.path.join(self.tmpdir, "other.db")
|
|
conn = sqlite3.connect(path)
|
|
conn.execute("CREATE TABLE something (id INTEGER)")
|
|
conn.commit()
|
|
conn.close()
|
|
ok, reason = PromptModel.verify_database_file(path)
|
|
self.assertFalse(ok)
|
|
self.assertIn("prompts", reason)
|
|
|
|
def test_prompts_table_missing_required_columns_is_rejected(self):
|
|
path = os.path.join(self.tmpdir, "partial.db")
|
|
conn = sqlite3.connect(path)
|
|
conn.execute("CREATE TABLE prompts (id INTEGER PRIMARY KEY, text TEXT)")
|
|
conn.commit()
|
|
conn.close()
|
|
ok, reason = PromptModel.verify_database_file(path)
|
|
self.assertFalse(ok)
|
|
self.assertIn("created_at", reason)
|
|
|
|
def test_corrupted_database_fails_integrity_check(self):
|
|
path = self._make_source_db([f"prompt {n}" for n in range(40)])
|
|
with open(path, "r+b") as fh:
|
|
page_size = int.from_bytes(fh.read(18)[16:18], "big") or 4096
|
|
fh.seek(page_size * 2) # scribble over pages 3 and 4 entirely
|
|
fh.write(b"\xff" * (page_size * 2))
|
|
ok, reason = PromptModel.verify_database_file(path)
|
|
self.assertFalse(ok)
|
|
self.assertTrue(reason)
|
|
|
|
def test_valid_database_passes(self):
|
|
path = self._make_source_db(["a", "b"])
|
|
ok, reason = PromptModel.verify_database_file(path)
|
|
self.assertTrue(ok, reason)
|
|
|
|
def test_verification_does_not_lock_or_modify_the_file(self):
|
|
path = self._make_source_db(["a"])
|
|
before = os.path.getsize(path)
|
|
PromptModel.verify_database_file(path)
|
|
os.unlink(path) # would fail on Windows if a connection were left open
|
|
self.assertEqual(before, before)
|
|
|
|
|
|
class TestRestoreFromFile(BackupTestCase):
|
|
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"])
|
|
|
|
backup_path = self.db.model.restore_from_file(src)
|
|
|
|
self.assertTrue(backup_path.startswith(self.path + ".backup_"))
|
|
self.assertTrue(os.path.exists(backup_path))
|
|
self.assertEqual(_texts(backup_path), ["old prompt"])
|
|
# The restore goes through SQLite's backup API into the live file, so
|
|
# the WAL stays valid (and shared with other instances) rather than
|
|
# being deleted under them; the content is what matters.
|
|
self.assertEqual(_texts(self.path), ["restored one", "restored two"])
|
|
|
|
def test_database_is_usable_from_every_thread_after_restore(self):
|
|
self._save(self.db, "old prompt")
|
|
src = self._make_source_db(["restored"])
|
|
|
|
ready = threading.Event()
|
|
proceed = threading.Event()
|
|
results = {}
|
|
|
|
def worker():
|
|
self.db.search_prompts(text="old") # opens this thread's connection
|
|
ready.set()
|
|
proceed.wait(5)
|
|
try:
|
|
results["texts"] = [p["text"] for p in self.db.search_prompts()]
|
|
except Exception as exc: # noqa: BLE001 - reported through results
|
|
results["error"] = exc
|
|
finally:
|
|
self.db.close()
|
|
|
|
thread = threading.Thread(target=worker)
|
|
thread.start()
|
|
ready.wait(5)
|
|
|
|
self.db.model.restore_from_file(src)
|
|
proceed.set()
|
|
thread.join(5)
|
|
|
|
self.assertNotIn("error", results)
|
|
self.assertEqual(results["texts"], ["restored"])
|
|
self.assertEqual([p["text"] for p in self.db.search_prompts()], ["restored"])
|
|
self.assertEqual(self._save(self.db, "after restore"), 2)
|
|
|
|
def test_restore_applies_migrations_to_an_old_schema(self):
|
|
src = os.path.join(self.tmpdir, "legacy.db")
|
|
conn = sqlite3.connect(src)
|
|
# The 3.0-era schema: workflow_name still present, no usage columns
|
|
conn.execute(
|
|
"CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT,"
|
|
" text TEXT NOT NULL, created_at TIMESTAMP, updated_at TIMESTAMP,"
|
|
" category TEXT, tags TEXT, rating INTEGER, notes TEXT,"
|
|
" hash TEXT UNIQUE, workflow_name TEXT)"
|
|
)
|
|
conn.execute(
|
|
"INSERT INTO prompts (text, created_at, tags, workflow_name)"
|
|
" VALUES ('legacy', '2024-01-01T00:00:00', '[\"old\"]', 'wf')"
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
self.db.model.restore_from_file(src)
|
|
|
|
prompt = self.db.get_prompt_by_id(1)
|
|
self.assertEqual(prompt["text"], "legacy")
|
|
self.assertEqual(prompt["tags"], ["old"])
|
|
self.assertEqual(prompt["run_count"], 1)
|
|
self.assertNotIn("workflow_name", prompt)
|
|
self.assertEqual(self.db.get_prompt_images(1), [])
|
|
|
|
def test_restore_without_a_live_database_returns_empty_backup_path(self):
|
|
src = self._make_source_db(["fresh"])
|
|
self.db.close_all()
|
|
os.unlink(self.path)
|
|
for suffix in ("-wal", "-shm"):
|
|
if os.path.exists(self.path + suffix):
|
|
os.unlink(self.path + suffix)
|
|
|
|
self.assertEqual(self.db.model.restore_from_file(src), "")
|
|
self.assertEqual(_texts(self.path), ["fresh"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|
|
|
|
class TestConnectionsAcrossInstances(BackupTestCase):
|
|
"""Several PromptModel instances share one file (nodes, API, monitor)."""
|
|
|
|
def test_close_all_closes_connections_opened_on_other_threads(self):
|
|
opened = []
|
|
|
|
def open_one():
|
|
opened.append(self.db.model.get_connection())
|
|
|
|
threads = [threading.Thread(target=open_one) for _ in range(3)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
self.assertEqual(len(opened), 3)
|
|
|
|
self.db.close_all()
|
|
|
|
for conn in opened:
|
|
with self.assertRaises(sqlite3.ProgrammingError):
|
|
conn.execute("SELECT 1")
|
|
|
|
def test_restore_is_seen_by_another_instance_holding_an_open_connection(self):
|
|
self.db.save_prompt(text="before", prompt_hash=generate_prompt_hash("before"))
|
|
snapshot = os.path.join(self.tmpdir, "snapshot.db")
|
|
other_db = PromptDatabase(snapshot)
|
|
for i in range(3):
|
|
other_db.save_prompt(
|
|
text=f"after {i}", prompt_hash=generate_prompt_hash(f"after {i}")
|
|
)
|
|
other_db.close_all()
|
|
|
|
# A second instance on the live file, with a connection kept open on
|
|
# a worker thread for the whole restore (like the image monitor).
|
|
bystander = PromptDatabase(self.path)
|
|
ready, go, seen = threading.Event(), threading.Event(), []
|
|
|
|
def hold_and_read():
|
|
bystander.model.get_connection().execute("SELECT COUNT(*) FROM prompts")
|
|
ready.set()
|
|
go.wait(10)
|
|
seen.append(len(bystander.search_prompts(limit=50)))
|
|
|
|
worker = threading.Thread(target=hold_and_read)
|
|
worker.start()
|
|
self.assertTrue(ready.wait(10))
|
|
try:
|
|
self.db.model.restore_from_file(snapshot)
|
|
finally:
|
|
go.set()
|
|
worker.join(10)
|
|
bystander.close_all()
|
|
|
|
self.assertEqual(seen, [3])
|
|
self.assertEqual(sorted(_texts(self.path)), ["after 0", "after 1", "after 2"])
|