Vito Sansevero
1f422d2c43
restore_from_file copied the uploaded file over the live database and deleted its WAL while other PromptModel instances (nodes, image monitor, queue hook) still held connections, which corrupts on POSIX and fails on Windows. It now copies pages through the online backup API into the live connection so every open handle sees the new content. PromptDatabase.close_all() closes the connections of every thread, which tests need before deleting a database. Output scans skip symlinked files, filesystem entries get a stable digest id instead of a per-process hash(), and thumbnail streaming stops when the client disconnects.
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()
|
|
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()
|
|
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()
|
|
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"])
|