diff --git a/TESTING.md b/TESTING.md index 61516fce..3e3b2153 100644 --- a/TESTING.md +++ b/TESTING.md @@ -391,6 +391,10 @@ A spec performance budget gets two tests: `test_filter_budget` (Images filters, spec D12: p95 < 150 ms at 50,000 files and 20,000 links). +A hot SQLite read that returns many rows aggregates them into one step +(`json_group_array`) instead of stepping once per row, so it does not wait for +the GIL once per row while ComfyUI generates (#354, `get_image_link_matches`). + The `perf budget` job in `.github/workflows/test.yml` is the reference measurement: `ubuntu-latest`, one Python, no coverage, only the wall-clock tests. It runs on every push to `main` and on `workflow_dispatch`, never on pull @@ -4656,8 +4660,18 @@ Filters (`mode=workspace`; anything invalid is 400 "Invalid gallery filters"): rating 5 to 1, then linked but unrated, then unlinked and ambiguous; ties break newest first. A signed cursor covers every filter and the sort, so reusing it with a change is 400. -- One link read (`get_image_link_matches`, streamed and bounded to the listed - paths) serves the filters, the rating sort and the count. +- One link read (`get_image_link_matches`) serves the filters, the rating sort + and the count. SQLite returns its rows as one JSON value in one step (#354): + Python's `sqlite3` releases and retakes the GIL around every step, and while + ComfyUI generates each retake can wait on the busy thread, so a read that + stepped once per linked row took seconds. The result keeps only the listed + paths; peak memory grows with the database's matching linked rows (about + 530 bytes each, freed on return). When the JSON step fails (a SQLite without + JSON functions, a BLOB value, a JSON text over SQLite's length limit) the read + streams the rows as before and logs one warning per process; a missing JSON + function stops the attempt for the rest of the process. An interrupted read + (Move library) raises as before. `tests/test_database_link_matches.py` proves both reads + equal on every filter and on odd values. `positive_count` is the number of distinct linked positives over the full matching set, not the page: diff --git a/database/operations.py b/database/operations.py index cb37accb..3087d615 100644 --- a/database/operations.py +++ b/database/operations.py @@ -7,7 +7,7 @@ import json import math import datetime import os -from typing import Optional, List, Dict, Any, Set, Tuple +from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple try: from ..utils.inspection_response import inherited_category @@ -251,8 +251,9 @@ WHERE gi.image_path IN ({marks}) OR gi.file_path IN ({marks}){fallback}""" # writer since the 3.2.4 migration stores file_path, so it is the path key. # The prompt-level conditions are uncorrelated IN sets SQLite builds once, and # the CASE evaluates the per-row ambiguity probe only for admitted rows. -_LINK_FILTER_SQL = """SELECT gi.file_path, p.rating, {hit}, p.id, -pp.id IS NOT NULL, pp.category, p.category +_LINK_FILTER_SQL = """SELECT gi.file_path AS file_path, p.rating AS rating, + {hit} AS matched, p.id AS positive, pp.id IS NOT NULL AS has_pair, + pp.category AS pair_category, p.category AS category FROM generated_images gi LEFT JOIN prompt_pairs pp ON pp.id = ( SELECT CASE WHEN COUNT(DISTINCT ip.pair_id) = 1 THEN MIN(ip.pair_id) END @@ -261,6 +262,25 @@ JOIN prompts p ON p.id = COALESCE(pp.positive_prompt_id, gi.prompt_id) WHERE gi.file_path IS NOT NULL AND CASE WHEN {admit} THEN NOT EXISTS ( SELECT 1 FROM generated_images other WHERE other.file_path = gi.file_path AND other.id <> gi.id) ELSE 0 END""" +# _LINK_FILTER_SQL's rows as one JSON array, built by SQLite in one step (#354). +# Python's sqlite3 releases and retakes the GIL around every step, so a read +# that steps once per linked row waits on a busy ComfyUI thread thousands of +# times. A BLOB value (no writer stores one) makes the value NULL (or +# json_array refuses it), and the read streams instead. The raw categories are +# put in only for a category read (#332), so plain reads stay as small as #354 +# measured them. +_LINK_MATCHES_JSON_SQL = """SELECT CASE + WHEN max({blob}) THEN NULL + ELSE json_group_array(json_array({columns})) END +FROM ({select})""" +_LINK_PLAIN_COLUMNS = ("file_path", "rating", "matched", "positive") +_LINK_CATEGORY_COLUMNS = _LINK_PLAIN_COLUMNS + ("has_pair", "pair_category", "category") +# The columns a BLOB can reach: matched, positive and has_pair are computed. +_LINK_STORED_COLUMNS = ("file_path", "rating", "pair_category", "category") +# Set when this SQLite has no JSON functions: every later read streams. +_one_step_unsupported = False +# One warning per process when a read falls back to streaming. +_one_step_warned = False # The linked positive of each given path (one unambiguous link, pair scope as # in _LINK_FILTER_SQL) with the raw categories of its inherited category; {hit} # as in link_filter_query. Used for the few files whose current image tags meet @@ -339,6 +359,50 @@ def link_filter_query( return sql, params +def link_matches_json_query(sql: str, categories: bool = False) -> str: + """*sql* from link_filter_query, wrapped to return its rows as one JSON value. + + With *categories* each row also carries the raw values the inherited + category is decided from; otherwise only path, rating, hit and positive. + """ + columns = _LINK_CATEGORY_COLUMNS if categories else _LINK_PLAIN_COLUMNS + blob = " OR ".join( + f"typeof({name}) = 'blob'" for name in _LINK_STORED_COLUMNS if name in columns + ) + return _LINK_MATCHES_JSON_SQL.format( + blob=blob, columns=", ".join(columns), select=sql + ) + + +def _link_matches_from_rows( + rows: Iterable[Sequence[Any]], categories: bool, wanted: Optional[Set[str]] +) -> Dict[str, LinkMatch]: + """LinkMatch per wanted path from link rows (one-step or streamed). + + Without *categories* the inherited category is None, as the caller then + never compares it. + """ + matches: Dict[str, LinkMatch] = {} + for path, rating, matched, positive, *raw in rows: + if wanted is None or path in wanted: + if categories: + matches[path] = _link_match(rating, matched, positive, *raw[:3]) + else: + matches[path] = (bool(matched), rating, positive, None) + return matches + + +def _streamed_link_matches( + conn: sqlite3.Connection, + sql: str, + params: Dict[str, Any], + wanted: Optional[Set[str]], + categories: bool = False, +) -> Dict[str, LinkMatch]: + """The link read stepped once per row (the fallback of the one-step read).""" + return _link_matches_from_rows(conn.execute(sql, params), categories, wanted) + + _ASCII_FOLD = str.maketrans("ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz") @@ -2588,30 +2652,72 @@ class PromptDatabase( names its negative variant and is not searched. Without tags only hits are returned, unless *every_link* asks for every linked image (the rating sort and the positive count). - The positive id is the card's linked positive. One SELECT whose rows - are streamed; with *wanted* (the listing's normalized paths) only those - paths are kept, so memory is bounded by the listing, not the database. - With *category* only links whose raw recipe or positive category is - that value are read (a superset of the links that inherit it; the - caller compares the inherited value) and the inherited category is - decided per row; otherwise it is None, which keeps the plain reads - cheap. + The positive id is the card's linked positive. With *category* only + links whose raw recipe or positive category is that value are read (a + superset of the links that inherit it; the caller compares the + inherited value) and the inherited category is decided per row; + otherwise it is None, which keeps the plain reads cheap. + + One SELECT, returned as one JSON value (#354): SQLite steps once, so + the read takes the GIL back a constant two or three times instead of + once per linked row while ComfyUI generates. Peak memory therefore + grows with the database's matching linked rows (about 530 bytes each, + more with the raw categories of a category read, freed on return); the + result is bounded by *wanted* (the listing's normalized paths) when + given. Where the JSON step fails (no JSON functions, a BLOB value, a + JSON text over SQLite's length limit) the rows are streamed as before. """ sql, params = link_filter_query( q, tags, every_link=every_link, category=category ) - matches: Dict[str, LinkMatch] = {} + categories = category is not None with self.model.get_connection() as conn: - rows = conn.execute(sql, params) - if category is not None: - for path, *row in rows: - if wanted is None or path in wanted: - matches[path] = _link_match(*row) - else: - for path, rating, matched, positive, *_categories in rows: - if wanted is None or path in wanted: - matches[path] = (bool(matched), rating, positive, None) - return matches + rows = None + if not _one_step_unsupported: + rows = self._one_step_link_rows(conn, sql, params, categories) + if rows is None: + return _streamed_link_matches(conn, sql, params, wanted, categories) + return _link_matches_from_rows(rows, categories, wanted) + + def _one_step_link_rows( + self, + conn: sqlite3.Connection, + sql: str, + params: Dict[str, Any], + categories: bool = False, + ) -> Optional[List[List[Any]]]: + """*sql*'s rows from one JSON value, or None when this read must stream. + + An interrupted read (Move library retiring the connection) is raised + as before, never retried row by row. + """ + global _one_step_unsupported, _one_step_warned + error: Optional[str] = None + try: + cursor = conn.execute(link_matches_json_query(sql, categories), params) + try: + ((document,),) = cursor.fetchall() + finally: + cursor.close() # resets the statement: its snapshot ends now + except (sqlite3.OperationalError, sqlite3.DataError) as failure: + # DataError: the JSON text is longer than SQLITE_LIMIT_LENGTH. + error = str(failure) + if error == "interrupted": + raise + if error.startswith("no such function"): + _one_step_unsupported = True + else: + if document is not None: + return json.loads(document) + error = "a BLOB value" + if not _one_step_warned: + _one_step_warned = True + self.logger.warning( + "Reading image links row by row (slower while ComfyUI " + "generates); the one-step read failed: %s", + error, + ) + return None def get_image_link_facts( self, paths: List[str], q: str = "" diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index 0097e7ca..604d4ade 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -1636,7 +1636,7 @@ class TestWorkspaceImages(ImageAPITestCase): import math from py.api import images as images_module - from database.operations import link_filter_query + from database.operations import link_filter_query, link_matches_json_query listed = await asyncio.get_running_loop().run_in_executor( None, self._budget_fixture, 5_000, 2_000 @@ -1765,24 +1765,42 @@ class TestWorkspaceImages(ImageAPITestCase): self.assertEqual(self.api._workspace_listed_cache, {}) self.assertEqual(self.api._workspace_order_cache, {}) # The per-row ambiguity probe is an index search, never a scan. + # The read runs it wrapped in one JSON aggregate (#354): same plan. sql, params = link_filter_query("lantern", ["tag-1"], every_link=True) - with self.api.db.model.get_connection() as conn: - plan = [row[3] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, 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")]) - # The category's narrowing clause (#332) keeps every join an index search. + for query in (sql, link_matches_json_query(sql)): + with self.subTest(wrapped=query is not sql): + with self.api.db.model.get_connection() as conn: + plan = [ + row[3] + for row in conn.execute("EXPLAIN QUERY PLAN " + query, 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")] + ) + # The category's narrowing clause (#332) keeps every join an index + # search, plain and wrapped with the raw categories. sql, params = link_filter_query( "", ["tag-1"], every_link=True, category="cat-3" ) - with self.api.db.model.get_connection() as conn: - plan = [row[3] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, params)] - self.assertIn( - "SEARCH other USING COVERING INDEX idx_image_file_path (file_path=?)", plan - ) - scans = [step for step in plan if step.startswith("SCAN")] - self.assertTrue(all(step.startswith("SCAN gi") for step in scans), plan) + for query in (sql, link_matches_json_query(sql, categories=True)): + with self.subTest(category_wrapped=query is not sql): + with self.api.db.model.get_connection() as conn: + plan = [ + row[3] + for row in conn.execute("EXPLAIN QUERY PLAN " + query, params) + ] + self.assertIn( + "SEARCH other USING COVERING INDEX idx_image_file_path " + "(file_path=?)", + plan, + ) + scans = [step for step in plan if step.startswith("SCAN")] + self.assertTrue(all(step.startswith("SCAN gi") for step in scans), plan) @unittest.skipUnless( PERF_BUDGET and not TRACED, diff --git a/tests/test_database_link_matches.py b/tests/test_database_link_matches.py new file mode 100644 index 00000000..8b50b3e3 --- /dev/null +++ b/tests/test_database_link_matches.py @@ -0,0 +1,456 @@ +"""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()