ComfyAssets-ComfyUI_PromptManager/tests/test_database_backup.py

"""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
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__))))

from database.library_commands import LibraryCommandError
from database.models import PromptModel
from database.operations import PromptDatabase
from tests.open_handles import assert_closed
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_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"])

        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_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"])

        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), [])
        conn = sqlite3.connect(self.path)
        try:
            self.assertEqual(
                [r[1] for r in conn.execute("PRAGMA table_info(prompt_library_state)")],
                ["prompt_id", "title"],
            )
        finally:
            conn.close()
        ok, _ = PromptModel.verify_database_file(self.path, require_library_state=True)
        self.assertTrue(ok)

    def test_backup_round_trips_library_state_from_the_wal(self):
        prompt_id = self._save(self.db, "titled")
        with self.db.model.get_connection() as conn:
            conn.execute(
                "INSERT INTO prompt_library_state VALUES (?, 'WAL title')",
                (prompt_id,),
            )
        dst = os.path.join(self.tmpdir, "backup.db")
        self.assertTrue(self.db.model.backup_database(dst))
        conn = sqlite3.connect(dst)
        try:
            self.assertEqual(
                conn.execute("SELECT title FROM prompt_library_state").fetchall(),
                [("WAL title",)],
            )
        finally:
            conn.close()
        self.assertTrue(
            PromptModel.verify_database_file(dst, require_library_state=True)[0]
        )

    def test_restore_old_backup_without_favorites(self):
        self.db.set_image_favorite("output", "live.png")
        src = self._make_source_db(["predates favorites"])
        with closing(sqlite3.connect(src)) as conn:
            conn.execute("DROP TABLE image_favorites")
            conn.commit()
        self.assertTrue(PromptModel.verify_database_file(src)[0])

        self.db.model.restore_from_file(src)

        self.assertEqual(_texts(self.path), ["predates favorites"])
        self.assertEqual(self.db.get_image_favorites("output", ["live.png"]), set())
        self.assertTrue(self.db.set_image_favorite("output", "new.png"))
        ok, reason = PromptModel.verify_database_file(
            self.path, require_library_state=True
        )
        self.assertTrue(ok, reason)
        with closing(sqlite3.connect(src)) as conn:  # the source is untouched
            self.assertIsNone(
                conn.execute(
                    "SELECT 1 FROM sqlite_master WHERE name='image_favorites'"
                ).fetchone()
            )

    def test_restore_old_backup_without_annotations(self):
        tables = (
            "media_image_tags",
            "media_annotation_bytes",
            "media_annotation_files",
            "media_annotation_metadata",
        )
        src = self._make_source_db(["predates image annotations"])
        with closing(sqlite3.connect(src)) as conn:
            conn.execute("PRAGMA foreign_keys=OFF")
            for table in tables:
                conn.execute(f"DROP TABLE {table}")
            conn.commit()
        self.assertTrue(PromptModel.verify_database_file(src)[0])
        ok, reason = PromptModel.verify_database_file(src, require_library_state=True)
        self.assertFalse(ok)
        self.assertEqual(reason, "Missing image annotation schema")

        self.db.model.restore_from_file(src)

        self.assertEqual(_texts(self.path), ["predates image annotations"])
        with closing(sqlite3.connect(self.path)) as conn:
            self.assertEqual(
                conn.execute("SELECT * FROM media_annotation_metadata").fetchall(),
                [(1, 1, 0)],
            )
        ok, reason = PromptModel.verify_database_file(
            self.path, require_library_state=True
        )
        self.assertTrue(ok, reason)
        with closing(sqlite3.connect(src)) as conn:  # the source is untouched
            self.assertEqual(
                conn.execute(
                    "SELECT COUNT(*) FROM sqlite_master WHERE name LIKE 'media_%'"
                ).fetchone()[0],
                0,
            )

    def test_restore_preserves_favorites_and_closes_handles(self):
        self.db.set_image_favorite("output", "live.png")
        with tempfile.TemporaryDirectory() as directory:
            src_path = os.path.join(directory, "source.db")
            src = PromptDatabase(src_path)
            try:
                self._save(src, "with favorites")
                src.set_image_favorite("output", "sub/kept.png")
                src.set_image_favorite("extra/root", "kept.png")
            finally:
                src.close_all()
            # The favorites still sit in the WAL; a backup must carry them.
            backup = os.path.join(directory, "backup.db")
            reopened = PromptDatabase(src_path)
            try:
                self.assertTrue(reopened.model.backup_database(backup))
            finally:
                reopened.close_all()

            self.db.model.restore_from_file(backup)

            self.assertEqual(
                self.db.get_image_favorites("output", ["sub/kept.png", "live.png"]),
                {"sub/kept.png"},
            )
            self.assertEqual(
                self.db.get_image_favorites("extra/root", ["kept.png"]), {"kept.png"}
            )
            PromptModel.close_all_instances()
            # Windows refuses to delete a file a restore left open.
            assert_closed(self, directory)
        backups = list(Path(self.tmpdir).glob("prompts.db.backup_*"))
        self.assertEqual(len(backups), 1)
        with closing(sqlite3.connect(str(backups[0]))) as conn:
            self.assertEqual(
                conn.execute(
                    "SELECT root_key, relative_path FROM image_favorites"
                ).fetchall(),
                [("output", "live.png")],
            )
        assert_closed(self, str(backups[0]))

    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 TestRetainedTags(BackupTestCase):
    """Unused registry tags are user data: they outlive restore and cleanup."""

    def _tag_names(self, db):
        return sorted(db.get_all_tags())

    def test_unused_tag_survives_backup_restore_and_restore_rotates_epoch(self):
        old_epoch = self.db.get_execution_epoch()
        self.db.create_vocabulary_tag("keeper", expected_epoch=old_epoch)
        backup = os.path.join(self.tmpdir, "backup.db")
        self.assertTrue(self.db.model.backup_database(backup))
        self.db.close_all()
        target_path = os.path.join(self.tmpdir, "target.db")
        target = PromptDatabase(target_path)
        try:
            target.model.restore_from_file(backup)
            self.assertIn("keeper", self._tag_names(target))
            new_epoch = target.get_execution_epoch()
            self.assertNotEqual(new_epoch, old_epoch)
            with self.assertRaises(LibraryCommandError) as caught:
                target.create_vocabulary_tag("other", expected_epoch=old_epoch)
            self.assertEqual(caught.exception.code, "storage-changed")
            self.assertNotIn("other", self._tag_names(target))
        finally:
            target.close_all()
        assert_closed(self, self.tmpdir)

    def test_unused_tag_survives_startup_maintenance(self):
        epoch = self.db.get_execution_epoch()
        self.db.create_vocabulary_tag("keeper", expected_epoch=epoch)
        owner = self.db.save_prompt(
            text="owner", prompt_hash=generate_prompt_hash("owner"), tags=["used"]
        )
        self.assertTrue(self.db.delete_prompt(owner))
        self.db.cleanup_duplicates()
        self.db.prune_orphaned_prompts()
        self.assertEqual(self._tag_names(self.db), ["keeper", "used"])
        self.db.close_all()
        assert_closed(self, self.tmpdir)


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"])