ComfyAssets-ComfyUI_PromptManager/tests/test_database_link_matches.py

"""The prompt-link read of the workspace Images page (#354).

``PromptDatabase.get_image_link_matches`` returns every row of its filter
SELECT in one SQLite step (one JSON value) instead of one step per row, so the
read does not have to win the GIL back once per linked image while ComfyUI
generates. Its result must equal the streamed read it replaces, which stays as
the fallback, on every filter and on odd values.
"""

import itertools
import os
import sqlite3
import sys
import tempfile
import unittest
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 import operations  # noqa: E402
from database.operations import PromptDatabase, link_filter_query  # noqa: E402

try:
    from tests.open_handles import assert_closed  # noqa: E402
except ImportError:  # discovered with tests/ as the top-level directory
    from open_handles import assert_closed  # noqa: E402

LOGGER = "prompt_manager.database"
MISSING_FUNCTION_SQL = (
    "SELECT pm_missing_function(json_group_array(json_array("
    "file_path, rating, matched, positive))) FROM ({select})"
)
QUERIES = ("", "lantern", "LANTERN title", "no-such-text")
CATEGORIES = (None, "cat-a", "cat-b", "cat-none")
TAG_SETS = ([], ["t1"], ["t1", "t2"], ["t-none"])

PROMPTS = (
    # id, text, category, rating, tags, library title
    (1, "red lantern at dusk", "cat-a", 5, ("t1", "t2"), None),
    (2, "blue harbor", "cat-a", None, ("t1",), None),
    (3, "green field", "cat-b", 3, (), "Lantern title"),
    (4, "pair positive", "cat-b", 4, ("t2",), None),
    (5, "pair negative lantern", "cat-a", 1, ("t1",), None),
)
QUOTED = '/out/it\'s "quoted".png'
BACKSLASHED = "C:\\out\\back\\slash.png"
UNICODE = "/out/ünïcødé 画像.png"
NEWLINE = "/out/new\nline.png"
IMAGES = (
    # prompt id, file path
    (1, "/out/a.png"),
    (2, QUOTED),
    (2, BACKSLASHED),
    (3, UNICODE),
    (1, NEWLINE),
    (1, "/out/ambiguous.png"),
    (2, "/out/ambiguous.png"),
    (5, "/out/pair.png"),  # linked through pair 1: its positive is prompt 4
    (3, "/out/two-pairs.png"),  # two pairs: no single pair, so prompt 3
    (1, None),  # no file path: never listed
)


class LinkMatchesTestCase(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.root = Path(os.path.realpath(self.temp.name))
        self.db = PromptDatabase(str(self.root / "prompts.db"))
        for name in ("_one_step_unsupported", "_one_step_warned"):
            patcher = patch.object(operations, name, False)
            patcher.start()
            self.addCleanup(patcher.stop)
        self.statements = []
        original = self.db.model.get_connection

        def traced_connection():
            connection = original()
            connection.set_trace_callback(self.statements.append)
            return connection

        patcher = patch.object(
            self.db.model, "get_connection", side_effect=traced_connection
        )
        patcher.start()
        self.addCleanup(patcher.stop)

    def tearDown(self):
        self.db.close_all()
        assert_closed(self, self.root)
        self.temp.cleanup()

    def seed(self):
        with self.db.model.get_connection() as conn:
            for prompt_id, text, category, rating, tags, title in PROMPTS:
                conn.execute(
                    "INSERT INTO prompts(id,text,category,rating,hash) "
                    "VALUES(?,?,?,?,?)",
                    (prompt_id, text, category, rating, f"hash-{prompt_id}"),
                )
                for tag in tags:
                    conn.execute("INSERT OR IGNORE INTO tags(name) VALUES(?)", (tag,))
                    conn.execute(
                        "INSERT INTO prompt_tags(prompt_id,tag_id) "
                        "SELECT ?, id FROM tags WHERE name=?",
                        (prompt_id, tag),
                    )
                if title:
                    conn.execute(
                        "INSERT INTO prompt_library_state(prompt_id,title) "
                        "VALUES(?,?)",
                        (prompt_id, title),
                    )
            ids = {}
            for prompt_id, path in IMAGES:
                cursor = conn.execute(
                    "INSERT INTO generated_images"
                    "(prompt_id,image_path,filename,file_path) VALUES(?,?,?,?)",
                    (prompt_id, path or "/elsewhere/x.png", "x.png", path),
                )
                ids.setdefault(path, cursor.lastrowid)
            conn.execute(
                "INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id) "
                "VALUES(1,4,5),(2,3,5),(3,3,NULL)"
            )
            conn.executemany(
                "INSERT INTO image_prompt_pairs(image_id,pair_id,provenance) "
                "VALUES(?,?,'{}')",
                [
                    (ids["/out/pair.png"], 1),
                    (ids["/out/two-pairs.png"], 2),
                    (ids["/out/two-pairs.png"], 3),
                ],
            )

    def add_image(self, prompt_id, path):
        with self.db.model.get_connection() as conn:
            conn.execute(
                "INSERT INTO generated_images(prompt_id,image_path,filename,file_path)"
                " VALUES(?,?,?,?)",
                (prompt_id, "/blob/x.png", "x.png", path),
            )

    def streamed(self, q, category, tags, every_link=False, wanted=None):
        """Today's read, called directly (the fallback)."""
        sql, params = link_filter_query(
            q, tags, every_link=every_link, category=category
        )
        with self.db.model.get_connection() as conn:
            return operations._streamed_link_matches(
                conn, sql, params, wanted, category is not None
            )

    def selects(self):
        found = [s for s in self.statements if s.lstrip().upper().startswith("SELECT")]
        self.statements.clear()
        return found


class TestOneStepLinkRead(LinkMatchesTestCase):
    def test_plain_read_returns_every_unambiguous_link(self):
        self.seed()
        matches = self.db.get_image_link_matches("", [], every_link=True)
        self.assertEqual(
            matches,
            {
                "/out/a.png": (True, 5, 1, None),
                QUOTED: (True, None, 2, None),
                BACKSLASHED: (True, None, 2, None),
                UNICODE: (True, 3, 3, None),
                NEWLINE: (True, 5, 1, None),
                "/out/pair.png": (True, 4, 4, None),
                "/out/two-pairs.png": (True, 3, 3, None),
            },
        )
        self.assertIsNone(matches[QUOTED][1])
        self.assertIsInstance(matches["/out/a.png"][1], int)
        self.assertIsInstance(matches["/out/a.png"][2], int)

    def test_search_hits_follow_text_title_and_tags(self):
        self.seed()
        self.assertEqual(
            self.db.get_image_link_matches("lantern", []),
            {
                "/out/a.png": (True, 5, 1, None),
                NEWLINE: (True, 5, 1, None),
                UNICODE: (True, 3, 3, None),
                "/out/two-pairs.png": (True, 3, 3, None),
            },
        )
        self.assertEqual(
            self.db.get_image_link_matches("lantern", [], category="cat-a"),
            {
                "/out/a.png": (True, 5, 1, "cat-a"),
                NEWLINE: (True, 5, 1, "cat-a"),
                QUOTED: (False, None, 2, "cat-a"),
                BACKSLASHED: (False, None, 2, "cat-a"),
            },
        )

    def test_every_filter_equals_the_streamed_read(self):
        self.seed()
        wanted_sets = (None, {"/out/a.png", UNICODE, "/out/pair.png", "/absent.png"})
        for q, category, tags, every_link, wanted in itertools.product(
            QUERIES, CATEGORIES, TAG_SETS, (False, True), wanted_sets
        ):
            with self.subTest(
                q=q, category=category, tags=tags, every=every_link, wanted=wanted
            ):
                expected = self.streamed(q, category, tags, every_link, wanted)
                self.statements.clear()
                actual = self.db.get_image_link_matches(
                    q,
                    list(tags),
                    every_link=every_link,
                    wanted=wanted,
                    category=category,
                )
                self.assertEqual(actual, expected)
                self.assertEqual(
                    {path: tuple(map(type, row)) for path, row in actual.items()},
                    {path: tuple(map(type, row)) for path, row in expected.items()},
                )
                found = self.selects()
                self.assertEqual(len(found), 1, found)
                self.assertIn("json_group_array", found[0])

    def test_empty_database_returns_nothing_in_one_statement(self):
        self.assertEqual(self.db.get_image_link_matches("", [], every_link=True), {})
        self.assertEqual(self.db.get_image_link_matches("x", ["t"], category="c"), {})
        found = self.selects()
        self.assertEqual(len(found), 2, found)
        self.assertTrue(all("json_group_array" in query for query in found))

    def test_wanted_limits_the_result_to_the_listing(self):
        self.seed()
        self.assertEqual(
            self.db.get_image_link_matches(
                "", [], every_link=True, wanted={"/out/pair.png", "/x.png"}
            ),
            {"/out/pair.png": (True, 4, 4, None)},
        )
        self.assertEqual(
            self.db.get_image_link_matches("", [], every_link=True, wanted=set()),
            {},
        )

    def test_category_read_carries_the_inherited_category_in_one_step(self):
        # A recipe's own category wins over its positive's (#332 MC4); the
        # read only narrows to links holding the value, the caller compares.
        self.seed()
        with self.db.model.get_connection() as conn:
            conn.execute("UPDATE prompt_pairs SET category='cat-a' WHERE id=1")
        self.statements.clear()
        for category, expected in (
            ("cat-a", {"/out/pair.png": (True, 4, 4, "cat-a")}),
            ("cat-b", {"/out/pair.png": (True, 4, 4, "cat-a")}),
        ):
            with self.subTest(category=category):
                actual = self.db.get_image_link_matches(
                    "", [], every_link=True, category=category, wanted={"/out/pair.png"}
                )
                self.assertEqual(actual, expected)
                self.assertEqual(
                    actual, self.streamed("", category, [], True, {"/out/pair.png"})
                )
        found = [query for query in self.selects() if "json_group_array" in query]
        self.assertEqual(len(found), 2, found)
        self.assertTrue(all("pair_category" in query for query in found))

    def test_plain_read_leaves_the_categories_out_of_the_json(self):
        self.seed()
        self.statements.clear()
        self.db.get_image_link_matches("", [], every_link=True)
        (query,) = self.selects()
        aggregate = query[query.index("json_array(") :].split(")", 1)[0]
        self.assertEqual(aggregate, "json_array(file_path, rating, matched, positive")

    def test_read_releases_its_snapshot(self):
        # A statement left open after the read would pin its WAL snapshot, and
        # a TRUNCATE checkpoint from another connection would report busy.
        self.seed()
        self.assertTrue(self.db.get_image_link_matches("", [], every_link=True))
        other = sqlite3.connect(str(self.root / "prompts.db"))
        try:
            other.execute("UPDATE prompts SET rating=2 WHERE id=1")
            other.commit()
            busy, _log, _done = other.execute(
                "PRAGMA wal_checkpoint(TRUNCATE)"
            ).fetchone()
        finally:
            other.close()
        self.assertEqual(busy, 0)
        self.assertEqual(os.path.getsize(self.root / "prompts.db-wal"), 0)

    def test_wrapped_query_keeps_the_ambiguity_probe_an_index_search(self):
        self.seed()
        sql, params = link_filter_query(
            "lantern", ["t1"], every_link=True, category="cat-a"
        )
        with self.db.model.get_connection() as conn:
            plan = [
                row[3]
                for row in conn.execute(
                    "EXPLAIN QUERY PLAN "
                    + operations.link_matches_json_query(sql, categories=True),
                    params,
                )
            ]
        self.assertIn(
            "SEARCH other USING COVERING INDEX idx_image_file_path (file_path=?)", plan
        )
        self.assertFalse([step for step in plan if step.startswith("SCAN other")])


class TestStreamedFallback(LinkMatchesTestCase):
    def test_missing_json_functions_stream_and_stop_trying(self):
        self.seed()
        expected = self.streamed("", None, [], True)
        self.statements.clear()
        wrapper = operations.link_matches_json_query
        with (
            patch.object(operations, "_LINK_MATCHES_JSON_SQL", MISSING_FUNCTION_SQL),
            patch.object(
                operations, "link_matches_json_query", side_effect=wrapper
            ) as attempts,
            self.assertLogs(LOGGER, "WARNING") as logged,
        ):
            first = self.db.get_image_link_matches("", [], every_link=True)
            found = self.selects()
            self.assertEqual(attempts.call_count, 1)
            second = self.db.get_image_link_matches("", [], every_link=True)
            again = self.selects()
        self.assertEqual(first, expected)
        self.assertEqual(second, expected)
        self.assertEqual(len(logged.records), 1, logged.output)
        self.assertIn("no such function: pm_missing_function", logged.output[0])
        self.assertTrue(operations._one_step_unsupported)
        # SQLite refuses the wrapped statement before it runs (not traced).
        self.assertEqual(len(found), 1, found)
        self.assertNotIn("json_group_array", found[0])
        # The next read skips the attempt for the rest of the process.
        self.assertEqual(attempts.call_count, 1)
        self.assertEqual(len(again), 1, again)
        self.assertNotIn("json_group_array", again[0])

    def test_blob_path_streams_that_read_and_keeps_trying(self):
        self.seed()
        self.add_image(2, b"/out/blob.png")
        expected = self.streamed("", None, [], True)
        self.assertEqual(expected[b"/out/blob.png"], (True, None, 2, None))
        self.statements.clear()
        with self.assertLogs(LOGGER, "WARNING") as logged:
            first = self.db.get_image_link_matches("", [], every_link=True)
            second = self.db.get_image_link_matches("", [], every_link=True)
        self.assertEqual(first, expected)
        self.assertEqual(second, expected)
        self.assertEqual(len(logged.records), 1, logged.output)
        self.assertFalse(operations._one_step_unsupported)
        found = self.selects()
        self.assertEqual(sum("json_group_array" in query for query in found), 2)
        self.assertEqual(len(found), 4, found)

    def test_blob_that_reads_as_json_still_streams(self):
        # SQLite 3.45+ reads a BLOB that is valid JSONB as JSON (b"\x00" is
        # null) instead of refusing it; the streamed read keeps the bytes.
        self.seed()
        self.add_image(2, b"\x00")
        expected = self.streamed("", None, [], True)
        self.assertIn(b"\x00", expected)
        with self.assertLogs(LOGGER, "WARNING"):
            self.assertEqual(
                self.db.get_image_link_matches("", [], every_link=True), expected
            )
        self.assertFalse(operations._one_step_unsupported)

    def test_blob_rating_streams(self):
        self.seed()
        with self.db.model.get_connection() as conn:
            # Only a database written around the CHECK can hold one.
            conn.execute("PRAGMA ignore_check_constraints=ON")
            conn.execute("UPDATE prompts SET rating=? WHERE id=3", (b"\x00",))
            conn.execute("PRAGMA ignore_check_constraints=OFF")
        expected = self.streamed("", None, [], True)
        self.assertEqual(expected[UNICODE], (True, b"\x00", 3, None))
        with self.assertLogs(LOGGER, "WARNING"):
            self.assertEqual(
                self.db.get_image_link_matches("", [], every_link=True), expected
            )

    def test_blob_category_streams_a_category_read(self):
        self.seed()
        with self.db.model.get_connection() as conn:
            # Admitted by its recipe's text category; the positive's is a BLOB.
            conn.execute("UPDATE prompts SET category=? WHERE id=4", (b"cat-b",))
            conn.execute("UPDATE prompt_pairs SET category='cat-b' WHERE id=1")
        expected = self.streamed("", "cat-b", [], True)
        self.assertEqual(expected["/out/pair.png"], (True, 4, 4, "cat-b"))
        with self.assertLogs(LOGGER, "WARNING") as logged:
            actual = self.db.get_image_link_matches(
                "", [], every_link=True, category="cat-b"
            )
        self.assertEqual(actual, expected)
        self.assertIn("BLOB", logged.output[0])
        self.assertFalse(operations._one_step_unsupported)

    def test_json_text_over_the_length_limit_streams(self):
        self.seed()
        expected = self.streamed("", None, [], True)
        conn = self.db.model.get_connection()
        limit = conn.setlimit(sqlite3.SQLITE_LIMIT_LENGTH, 200)
        try:
            with self.assertLogs(LOGGER, "WARNING") as logged:
                actual = self.db.get_image_link_matches("", [], every_link=True)
        finally:
            conn.setlimit(sqlite3.SQLITE_LIMIT_LENGTH, limit)
        self.assertEqual(actual, expected)
        self.assertIn("too big", logged.output[0])
        self.assertFalse(operations._one_step_unsupported)

    def test_an_interrupted_read_raises_without_a_fallback(self):
        # Move library interrupts a retired connection: the read fails as
        # before, and the one warning is kept for a real fallback.
        self.seed()
        conn = self.db.model.get_connection()
        conn.set_progress_handler(lambda: 1, 1)
        self.statements.clear()
        try:
            with (
                self.assertNoLogs(LOGGER, "WARNING"),
                self.assertRaisesRegex(sqlite3.OperationalError, "^interrupted$"),
            ):
                self.db.get_image_link_matches("", [], every_link=True)
        finally:
            conn.set_progress_handler(None, 1)
        self.assertFalse(operations._one_step_warned)
        self.assertFalse(operations._one_step_unsupported)
        self.assertFalse([s for s in self.statements if "json_group_array" not in s])

    def test_the_maintenance_link_rows_reader_is_not_shadowed(self):
        from database.maintenance import MaintenanceOperations

        self.assertIs(PromptDatabase._link_rows, MaintenanceOperations._link_rows)

    def test_a_failed_streamed_read_still_raises(self):
        with (
            patch.object(operations, "_LINK_MATCHES_JSON_SQL", MISSING_FUNCTION_SQL),
            patch.object(operations, "_LINK_FILTER_SQL", "SELECT broken FROM nowhere"),
            self.assertLogs(LOGGER, "WARNING"),
        ):
            with self.assertRaises(sqlite3.OperationalError):
                self.db.get_image_link_matches("", [], every_link=True)


if __name__ == "__main__":
    unittest.main()