diff --git a/database/models.py b/database/models.py index 14a491b..30499cf 100644 --- a/database/models.py +++ b/database/models.py @@ -3,10 +3,12 @@ Database schema and models for KikoTextEncode prompt storage. """ import datetime +import shutil import sqlite3 import os import threading -from typing import Set +from pathlib import Path +from typing import Set, Tuple # Import logging system try: @@ -31,6 +33,10 @@ def utc_now_iso() -> str: ) +SQLITE_HEADER = b"SQLite format 3\x00" +REQUIRED_PROMPT_COLUMNS = ("id", "text", "created_at") + + class PromptModel: """Database model for prompt storage and schema management.""" @@ -49,6 +55,9 @@ class PromptModel: self._local = threading.local() self._connections: Set[sqlite3.Connection] = set() self._conn_lock = threading.Lock() + # Bumped by _close_all_connections so threads holding a closed + # connection reopen on their next get_connection(). + self._generation = 0 self._ensure_database_exists() def _ensure_database_exists(self) -> None: @@ -62,13 +71,16 @@ class PromptModel: Exception: If database creation fails """ try: - with sqlite3.connect(self.db_path) as conn: + conn = sqlite3.connect(self.db_path) + try: conn.execute("PRAGMA journal_mode = WAL") conn.execute("PRAGMA busy_timeout = 5000") conn.execute("PRAGMA foreign_keys = ON") self._create_tables(conn) self._create_indexes(conn) conn.commit() + finally: + conn.close() except Exception as e: self.logger.error(f"Error creating database: {e}") raise @@ -200,7 +212,7 @@ class PromptModel: sqlite3.Connection: The current thread's configured connection """ conn = getattr(self._local, "conn", None) - if conn is not None: + 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 @@ -211,10 +223,24 @@ class PromptModel: conn.execute("PRAGMA foreign_keys = ON") conn.execute("PRAGMA busy_timeout = 5000") self._local.conn = conn + self._local.generation = self._generation with self._conn_lock: self._connections.add(conn) return conn + def _close_all_connections(self) -> None: + """Close every thread's connection; they reopen lazily on next use.""" + with self._conn_lock: + connections = list(self._connections) + self._connections.clear() + self._generation += 1 + self._local.conn = None + for conn in connections: + try: + conn.close() + except sqlite3.Error as e: + self.logger.warning(f"Error closing database connection: {e}") + def close(self) -> None: """Close the calling thread's connection, if it has one. @@ -599,19 +625,127 @@ class PromptModel: def backup_database(self, backup_path: str) -> bool: """ - Create a backup of the database. + Write a consistent copy of the database to backup_path. + + Uses the SQLite online backup API through this thread's connection, + so rows that are committed but still sit in the -wal file are + included and other threads may keep reading and writing meanwhile. Args: - backup_path: Path where the backup should be saved + backup_path: Path where the backup should be saved (overwritten) Returns: bool: True if backup was successful, False otherwise """ try: - import shutil - - shutil.copy2(self.db_path, backup_path) + if os.path.exists(backup_path) and os.path.samefile( + backup_path, self.db_path + ): + self.logger.error("Refusing to back up the database onto itself") + return False + if os.path.exists(backup_path): + os.remove(backup_path) + source = self.get_connection() + destination = sqlite3.connect(backup_path) + try: + source.backup(destination) + finally: + destination.close() return True - except Exception as e: + except (OSError, sqlite3.Error) as e: self.logger.error(f"Error creating database backup: {e}") return False + + @staticmethod + def verify_database_file(path: str) -> Tuple[bool, str]: + """ + Check that a file is a healthy prompts database before restoring it. + + Opens the file read-only, runs PRAGMA integrity_check and confirms the + prompts table has the columns every version of this extension needs. + + Args: + path: File to inspect + + Returns: + (ok, reason): reason is "ok" when the file is acceptable, otherwise + a human-readable explanation. Never raises for bad input. + """ + if not os.path.isfile(path): + return False, "File not found" + try: + with open(path, "rb") as fh: + header = fh.read(len(SQLITE_HEADER)) + except OSError as e: + return False, f"Cannot read file: {e}" + if header != SQLITE_HEADER: + return False, "Not a SQLite database file" + + uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro" + try: + conn = sqlite3.connect(uri, uri=True) + except sqlite3.Error as e: + return False, f"Invalid SQLite database: {e}" + try: + row = conn.execute("PRAGMA integrity_check").fetchone() + verdict = row[0] if row else "no result" + if verdict != "ok": + return False, f"Integrity check failed: {verdict}" + table = conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'" + ).fetchone() + if not table: + return False, "Database does not contain a 'prompts' table" + columns = {r[1] for r in conn.execute("PRAGMA table_info(prompts)")} + missing = [c for c in REQUIRED_PROMPT_COLUMNS if c not in columns] + if missing: + return ( + False, + "prompts table is missing required columns: " + ", ".join(missing), + ) + return True, "ok" + except sqlite3.Error as e: + return False, f"Invalid SQLite database: {e}" + finally: + conn.close() + + def restore_from_file(self, src_path: str) -> str: + """ + Replace the live database with the file at src_path. + + The current database is first backed up next to itself as + ``.backup_`` (WAL content included), every + thread's connection is closed, the file is copied over, stale -wal and + -shm files are removed and the schema migrations run on the restored + file. Connections reopen lazily on the next get_connection(). + + Callers should run verify_database_file(src_path) first. + + Args: + src_path: SQLite file to restore + + Returns: + Path of the backup taken of the previous database, or "" when + there was no database to back up. + + Raises: + RuntimeError: If the previous database could not be backed up + OSError: If the restored file could not be copied into place + """ + backup_path = "" + if os.path.exists(self.db_path): + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + backup_path = f"{self.db_path}.backup_{timestamp}" + 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._close_all_connections() + 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) + self._ensure_database_exists() + self.logger.info(f"Database restored from {src_path}") + return backup_path diff --git a/py/api/admin.py b/py/api/admin.py index ae9b6a7..7c2fb8e 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -18,6 +18,11 @@ try: except ImportError: from utils.validators import validate_result_timeout +try: + from ...database.operations import PromptDatabase +except ImportError: + from database.operations import PromptDatabase + class AdminRoutesMixin: """Mixin providing admin, diagnostics, and maintenance API endpoints.""" @@ -883,24 +888,33 @@ class AdminRoutesMixin: ) async def backup_database(self, request): - """Backup the entire prompts.db database file.""" + """Download a consistent copy of the prompts database (WAL included).""" try: - db_path = self.db.model.db_path + model = self.db.model + db_path = model.db_path if not os.path.exists(db_path): return web.json_response( {"success": False, "error": "Database file not found"}, status=404 ) - with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as temp_file: - temp_path = temp_file.name - - shutil.copy2(db_path, temp_path) - - with open(temp_path, "rb") as f: - file_data = f.read() - - os.unlink(temp_path) + fd, temp_path = tempfile.mkstemp(suffix=".db") + os.close(fd) + try: + loop = asyncio.get_running_loop() + ok = await loop.run_in_executor(None, model.backup_database, temp_path) + if not ok: + return web.json_response( + {"success": False, "error": "Failed to create database backup"}, + status=500, + ) + # Read once (the file is deleted right after) rather than + # streaming so the temp file never outlives the request. + with open(temp_path, "rb") as fh: + file_data = fh.read() + finally: + if os.path.exists(temp_path): + os.unlink(temp_path) timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"prompts_backup_{timestamp}.db" @@ -922,7 +936,14 @@ class AdminRoutesMixin: ) async def restore_database(self, request): - """Restore the prompts.db database from uploaded file.""" + """Restore the prompts database from an uploaded SQLite file. + + The upload is streamed to a temp file (capped at restore_max_bytes), + verified with PromptModel.verify_database_file and only then copied + over the live database, which is backed up first. + """ + max_bytes = getattr(self, "restore_max_bytes", 100 * 1024 * 1024) + temp_path = None try: reader = await request.multipart() field = await reader.next() @@ -936,95 +957,52 @@ class AdminRoutesMixin: status=400, ) - MAX_RESTORE_SIZE = 100 * 1024 * 1024 # 100MB - file_data = await field.read() + fd, temp_path = tempfile.mkstemp(suffix=".db") + os.close(fd) + size = 0 + with open(temp_path, "wb") as out: + while True: + chunk = await field.read_chunk() + if not chunk: + break + size += len(chunk) + if size > max_bytes: + return web.json_response( + { + "success": False, + "error": f"File too large. Maximum size is {max_bytes // (1024 * 1024)}MB", + }, + status=400, + ) + out.write(chunk) - if not file_data: + if size == 0: return web.json_response( {"success": False, "error": "Uploaded file is empty"}, status=400 ) - if len(file_data) > MAX_RESTORE_SIZE: + model = self.db.model + ok, reason = model.verify_database_file(temp_path) + if not ok: return web.json_response( - { - "success": False, - "error": f"File too large. Maximum size is {MAX_RESTORE_SIZE // (1024*1024)}MB", - }, - status=400, + {"success": False, "error": reason}, status=400 ) - with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as temp_file: - temp_path = temp_file.name - temp_file.write(file_data) + loop = asyncio.get_running_loop() + backup_path = await loop.run_in_executor( + None, model.restore_from_file, temp_path + ) + self.db = PromptDatabase(model.db_path) + prompt_count = self.db.model.get_database_info().get("total_prompts", 0) - try: - with sqlite3.connect(temp_path) as conn: - conn.row_factory = sqlite3.Row - - cursor = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'" - ) - if not cursor.fetchone(): - raise ValueError("Database does not contain a 'prompts' table") - - cursor = conn.execute("PRAGMA table_info(prompts)") - columns = [row["name"] for row in cursor.fetchall()] - required_columns = ["id", "text", "created_at"] - - for col in required_columns: - if col not in columns: - raise ValueError(f"Database missing required column: {col}") - - cursor = conn.execute("SELECT COUNT(*) as count FROM prompts") - prompt_count = cursor.fetchone()["count"] - - db_path = self.db.model.db_path - backup_path = f"{db_path}.backup_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" - - if os.path.exists(db_path): - shutil.copy2(db_path, backup_path) - self.logger.info(f"Current database backed up to: {backup_path}") - - shutil.copy2(temp_path, db_path) - - # Reinitialize the database connection - try: - from ...database.operations import PromptDatabase - except ImportError: - import sys - - sys.path.insert( - 0, - os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ), - ) - from database.operations import PromptDatabase - self.db = PromptDatabase() - - return web.json_response( - { - "success": True, - "message": f"Database restored successfully. Found {prompt_count} prompts.", - "prompt_count": prompt_count, - "backup_created": ( - backup_path if os.path.exists(db_path) else None - ), - } - ) - - except sqlite3.Error as e: - return web.json_response( - {"success": False, "error": f"Invalid SQLite database: {str(e)}"}, - status=400, - ) - except ValueError as e: - return web.json_response( - {"success": False, "error": str(e)}, status=400 - ) - finally: - if os.path.exists(temp_path): - os.unlink(temp_path) + return web.json_response( + { + "success": True, + "message": f"Database restored successfully. Found {prompt_count} prompts.", + "prompt_count": prompt_count, + "backup_created": backup_path or None, + } + ) except Exception as e: self.logger.error(f"Restore error: {e}", exc_info=True) @@ -1032,6 +1010,11 @@ class AdminRoutesMixin: {"success": False, "error": f"Failed to restore database: {str(e)}"}, status=500, ) + finally: + if temp_path: + for suffix in ("", "-wal", "-shm", "-journal"): + if os.path.exists(temp_path + suffix): + os.unlink(temp_path + suffix) async def scan_images(self, request): """Scan ComfyUI output images for prompt metadata and add them to the database.""" diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py new file mode 100644 index 0000000..06c100f --- /dev/null +++ b/tests/test_api_backup_restore.py @@ -0,0 +1,239 @@ +"""Route tests for GET /prompt_manager/backup and POST /prompt_manager/restore. + +The backup must include rows still sitting in the WAL, the restore must +stream the upload to disk, reject anything that is not a healthy prompts +database and leave the live database untouched when it does. +""" + +import os +import sqlite3 +import sys +import tempfile +import unittest +from unittest.mock import patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import aiohttp +from aiohttp import web +from aiohttp.test_utils import AioHTTPTestCase + +from database.operations import PromptDatabase +from py.api import PromptManagerAPI +from utils.hashing import generate_prompt_hash + + +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 BackupRestoreTestCase(AioHTTPTestCase): + async def get_application(self): + self.tmpdir = tempfile.mkdtemp() + self.db_path = os.path.join(self.tmpdir, "prompts.db") + + app = web.Application() + routes = web.RouteTableDef() + self.api = PromptManagerAPI() + self.api.db = PromptDatabase(self.db_path) + self.api.add_routes(routes) + app.router.add_routes(routes) + return app + + async def tearDownAsync(self): + self.api.db.close() + for name in os.listdir(self.tmpdir): + os.unlink(os.path.join(self.tmpdir, name)) + os.rmdir(self.tmpdir) + + def _save(self, text): + return self.api.db.save_prompt( + text=text, prompt_hash=generate_prompt_hash(text) + ) + + def _write(self, name, data): + path = os.path.join(self.tmpdir, name) + with open(path, "wb") as fh: + fh.write(data) + return path + + def _make_db_file(self, texts, name="upload.db"): + path = os.path.join(self.tmpdir, name) + db = PromptDatabase(path) + for text in texts: + db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text)) + db.close() + return path + + async def _restore(self, path, field="database_file"): + with open(path, "rb") as fh: + form = aiohttp.FormData() + form.add_field( + field, + fh.read(), + filename=os.path.basename(path), + content_type="application/octet-stream", + ) + return await self.client.post("/prompt_manager/restore", data=form) + + +class TestBackupRoute(BackupRestoreTestCase): + async def test_backup_download_contains_rows_still_in_the_wal(self): + for n in range(3): + self._save(f"backup me {n}") + self.assertGreater(os.path.getsize(self.db_path + "-wal"), 0) + created = [] + real_mkstemp = tempfile.mkstemp + + def spy_mkstemp(*args, **kwargs): + result = real_mkstemp(*args, **kwargs) + created.append(result[1]) + return result + + with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp): + resp = await self.client.get("/prompt_manager/backup") + + self.assertEqual(resp.status, 200) + self.assertEqual(resp.content_type, "application/octet-stream") + disposition = resp.headers["Content-Disposition"] + self.assertIn("attachment", disposition) + self.assertIn("prompts_backup_", disposition) + self.assertTrue(disposition.rstrip('"').endswith(".db")) + + body = await resp.read() + downloaded = self._write("downloaded.db", body) + self.assertEqual( + _texts(downloaded), ["backup me 0", "backup me 1", "backup me 2"] + ) + # The server-side temp file is gone once the response is complete + self.assertEqual(len(created), 1) + self.assertFalse(os.path.exists(created[0])) + + async def test_backup_of_missing_database_is_404(self): + self.api.db.model.db_path = os.path.join(self.tmpdir, "nope.db") + resp = await self.client.get("/prompt_manager/backup") + self.assertEqual(resp.status, 404) + data = await resp.json() + self.assertFalse(data["success"]) + + async def test_backup_failure_is_500(self): + self._save("x") + with patch.object(self.api.db.model, "backup_database", return_value=False): + resp = await self.client.get("/prompt_manager/backup") + self.assertEqual(resp.status, 500) + data = await resp.json() + self.assertFalse(data["success"]) + self.assertIn("backup", data["error"].lower()) + + +class TestRestoreRoute(BackupRestoreTestCase): + async def test_restore_replaces_database_and_reports_backup(self): + self._save("old prompt") + old_db = self.api.db + upload = self._make_db_file(["restored a", "restored b"]) + + resp = await self._restore(upload) + + self.assertEqual(resp.status, 200) + data = await resp.json() + self.assertTrue(data["success"]) + self.assertEqual(data["prompt_count"], 2) + self.assertTrue(os.path.exists(data["backup_created"])) + self.assertEqual(_texts(data["backup_created"]), ["old prompt"]) + self.assertIsNot(self.api.db, old_db) + self.assertEqual(self.api.db.model.db_path, self.db_path) + + texts = sorted(p["text"] for p in self.api.db.search_prompts()) + self.assertEqual(texts, ["restored a", "restored b"]) + + async def test_restore_leaves_no_temp_files_behind(self): + upload = self._make_db_file(["restored"]) + created = [] + real_mkstemp = tempfile.mkstemp + + def spy_mkstemp(*args, **kwargs): + result = real_mkstemp(*args, **kwargs) + created.append(result[1]) + return result + + with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp): + resp = await self._restore(upload) + + self.assertEqual(resp.status, 200) + self.assertEqual(len(created), 1) + for suffix in ("", "-wal", "-shm", "-journal"): + self.assertFalse(os.path.exists(created[0] + suffix)) + + async def test_text_file_is_rejected_and_database_untouched(self): + self._save("keep me") + upload = self._write("notes.txt", b"this is not a database") + + resp = await self._restore(upload) + + self.assertEqual(resp.status, 400) + data = await resp.json() + self.assertFalse(data["success"]) + self.assertIn("SQLite", data["error"]) + self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"]) + self.assertFalse( + any(n.startswith("prompts.db.backup_") for n in os.listdir(self.tmpdir)) + ) + + async def test_database_without_prompts_table_is_rejected(self): + self._save("keep me") + path = os.path.join(self.tmpdir, "other.db") + conn = sqlite3.connect(path) + conn.execute("CREATE TABLE things (id INTEGER)") + conn.commit() + conn.close() + + resp = await self._restore(path) + + self.assertEqual(resp.status, 400) + data = await resp.json() + self.assertFalse(data["success"]) + self.assertIn("prompts", data["error"]) + self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"]) + + async def test_wrong_field_name_is_rejected(self): + upload = self._make_db_file(["restored"]) + resp = await self._restore(upload, field="file") + self.assertEqual(resp.status, 400) + data = await resp.json() + self.assertFalse(data["success"]) + self.assertIn("database_file", data["error"]) + + async def test_empty_upload_is_rejected(self): + upload = self._write("empty.db", b"") + resp = await self._restore(upload) + self.assertEqual(resp.status, 400) + data = await resp.json() + self.assertFalse(data["success"]) + self.assertIn("empty", data["error"].lower()) + + async def test_upload_over_the_size_limit_is_rejected_while_streaming(self): + self._save("keep me") + self.api.restore_max_bytes = 1024 + upload = self._write("big.db", b"SQLite format 3\x00" + b"\x00" * 8192) + + resp = await self._restore(upload) + + self.assertEqual(resp.status, 400) + data = await resp.json() + self.assertFalse(data["success"]) + self.assertIn("too large", data["error"].lower()) + self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"]) + + async def test_non_multipart_request_is_rejected(self): + resp = await self.client.post("/prompt_manager/restore", data=b"raw bytes") + self.assertIn(resp.status, (400, 500)) + data = await resp.json() + self.assertFalse(data["success"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py new file mode 100644 index 0000000..8326b3b --- /dev/null +++ b/tests/test_database_backup.py @@ -0,0 +1,243 @@ +"""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"]) + self.assertFalse(os.path.exists(self.path + "-wal")) + self.assertFalse(os.path.exists(self.path + "-shm")) + 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()