diff --git a/database/image_collections.py b/database/image_collections.py index 43193b0b..9eeb4494 100644 --- a/database/image_collections.py +++ b/database/image_collections.py @@ -13,7 +13,9 @@ NFC name) is UNIQUE, so names are unique ignoring case and Unicode form. The two tables install like ``image_favorites`` and ``saved_views``: optional in old backups, required once a database is current, created empty with no backfill. Startup checks only the structure, so a value saved under older -rules never stops the library from opening; restore checks the values too. +rules never stops the library from opening; restore checks the values too, +but never recomputes ``name_key`` (case tables differ between Python +versions). """ import re @@ -29,7 +31,9 @@ FILE_INDEX = "idx_image_collection_items_file" TABLES = (COLLECTIONS_TABLE, ITEMS_TABLE) MAX_NAME_LENGTH = 100 _TIMESTAMP_BYTES = 64 -_CONTROL = re.compile(r"[\x00-\x1f\x7f]") +# A name key's byte bound: case folding expands a character to at most +# three, each at most 4 UTF-8 bytes. +_MAX_KEY_BYTES = MAX_NAME_LENGTH * 3 * 4 _WHITESPACE = re.compile(r"\s+") _TABLE_DDL = { @@ -72,6 +76,11 @@ class CollectionInvalid(ValueError): """A name, id or target breaks the write rules (a 400).""" +def _has_control(text: str) -> bool: + """Any Unicode control character (category Cc: C0, DEL and C1).""" + return any(unicodedata.category(char) == "Cc" for char in text) + + def normalize_collection_name(name: object) -> str: """The stored spelling of *name*: trimmed, whitespace collapsed, checked. @@ -83,7 +92,7 @@ def normalize_collection_name(name: object) -> str: collapsed = _WHITESPACE.sub(" ", name).strip() if not 1 <= len(collapsed) <= MAX_NAME_LENGTH: raise CollectionInvalid(f"Name must be 1 to {MAX_NAME_LENGTH} characters") - if _CONTROL.search(collapsed): + if _has_control(collapsed): raise CollectionInvalid("Name contains control characters") return collapsed @@ -148,17 +157,19 @@ def _value_error(conn: sqlite3.Connection) -> Optional[str]: "typeof(id) != 'integer' OR typeof(name) != 'text' " "OR typeof(name_key) != 'text' OR typeof(created_at) != 'text' " f"OR length(CAST(name AS BLOB)) > {MAX_NAME_LENGTH * 4} " + "OR length(CAST(name_key AS BLOB)) = 0 " + f"OR length(CAST(name_key AS BLOB)) > {_MAX_KEY_BYTES} " f"OR length(CAST(created_at AS BLOB)) > {_TIMESTAMP_BYTES} LIMIT 1" ).fetchone(): return "Image collections contains invalid values" - for name, name_key in conn.execute( - f"SELECT name, name_key FROM {COLLECTIONS_TABLE}" # nosec B608 - ): + # The key is not recomputed: Unicode case tables differ between Python + # versions, and a key written by another version must stay restorable. + # Its UNIQUE constraint keeps names distinct; writes compute it. + for (name,) in conn.execute(f"SELECT name FROM {COLLECTIONS_TABLE}"): # nosec B608 if ( not 1 <= len(name) <= MAX_NAME_LENGTH - or _CONTROL.search(name) + or _has_control(name) or name != name.strip() - or name_key != collection_name_key(name) ): return "Image collections contains invalid names" limit = MAX_KEY_LENGTH * 4 @@ -352,6 +363,26 @@ class ImageCollectionOperations: found.setdefault(collection_id, set()).add(path) return found + def collection_count_snapshot(self, root_keys: list) -> tuple[dict, dict]: + """``({root key: {id: relative keys}}, {id: stored rows})`` read in + one read transaction, so a write between the reads cannot skew them. + """ + conn = self.model.get_connection() + if conn.in_transaction: + raise RuntimeError("Collection counts require their own transaction") + conn.execute("BEGIN") + try: + members = {} + for root_key in dict.fromkeys(root_keys): + found: dict[int, set[str]] = {} + for collection_id, path in conn.execute(ROOT_MEMBERS_SQL, (root_key,)): + found.setdefault(collection_id, set()).add(path) + members[root_key] = found + stored = {cid: count for cid, count in conn.execute(COUNTS_SQL)} + finally: + conn.rollback() # it only read + return members, stored + def collection_counts(self) -> dict[int, int]: """``{collection id: stored rows}`` for every collection.""" rows = self.model.get_connection().execute(COUNTS_SQL) diff --git a/py/api/collections.py b/py/api/collections.py index 51fab44d..bc8046ec 100644 --- a/py/api/collections.py +++ b/py/api/collections.py @@ -49,7 +49,7 @@ BASE = "/prompt_manager/collections" # The most files one add or remove names (the bulk tagging cap, MC12). MAX_ITEMS = 500 MAX_ID = 999_999_999 -_PATH_ID = re.compile(r"[0-9]{1,9}") +_PATH_ID = re.compile(r"[1-9][0-9]{0,8}") _ROOT_INDEX = re.compile(r"[0-9]{1,9}") _ADD_TARGETS = (frozenset({"collection_id", "items"}), frozenset({"name", "items"})) _REMOVE_FIELDS = frozenset({"collection_id", "items"}) @@ -113,7 +113,7 @@ def _collection_response(failure: str) -> Callable: def _path_id(request: web.Request) -> int: raw = request.match_info["collection_id"] - if not raw.isascii() or _PATH_ID.fullmatch(raw) is None or int(raw) < 1: + if not raw.isascii() or _PATH_ID.fullmatch(raw) is None: raise _Refused(bad_request("Invalid collection id")) return int(raw) @@ -170,53 +170,66 @@ class CollectionRoutesMixin: for collection in collections: listed, stored = (None, None) if counts is None else counts count = None if listed is None else listed.get(collection["id"], 0) - missing = None if stored is None else stored.get(collection["id"], 0) + stored_rows = None if stored is None else stored.get(collection["id"], 0) rows.append( { **collection, "count": count, - "missing": None if count is None else missing - count, + # Clamped: never below zero, whatever the snapshot saw. + "missing": None if count is None else max(0, stored_rows - count), } ) return web.json_response({"success": True, "collections": rows}) async def _collection_counts(self) -> tuple[dict, dict]: - """``({id: listed members}, {id: stored rows})`` over every source.""" + """``({id: listed members}, {id: stored rows})`` over every source. + + The scans come from the gallery cache; the root keys (realpath), the + listings and one read snapshot of every member and stored count run + in the executor, so a write between two reads cannot skew them. + """ + roots = list(self._get_all_output_dirs()) + root_keys = await self._run_in_executor( + lambda: [_favorite_root_key(root) for root in roots] + ) + distinct = {} # one folder configured twice counts once + for root, root_key in zip(roots, root_keys): + distinct.setdefault(root_key, root) + scans = { + root_key: await self._get_gallery_files(root) + for root_key, root in distinct.items() + } + return await self._run_in_executor(self._count_members, distinct, scans) + + def _count_members(self, roots: dict, scans: dict) -> tuple[dict, dict]: + """Blocking: match one snapshot's members against each listing.""" + members, stored = self.db.collection_count_snapshot(list(roots)) listed: dict[int, int] = {} - seen = set() - for root in self._get_all_output_dirs(): - root_key = _favorite_root_key(root) - if root_key in seen: # one folder configured twice counts once - continue - seen.add(root_key) - files = await self._get_gallery_files(root) - listing = await self._run_in_executor(self._favorite_listing, root, files) - members = await self._run_in_executor( - self.db.root_collection_members, root_key - ) - for collection_id, keys in members.items(): + for root_key, root in roots.items(): + listing = self._favorite_listing(root, scans[root_key]) + for collection_id, keys in members.get(root_key, {}).items(): found = sum(1 for key in keys if key in listing) listed[collection_id] = listed.get(collection_id, 0) + found - stored = await self._run_in_executor(self.db.collection_counts) return listed, stored @_collection_response(READ_FAILED) async def collection_membership(self, request: web.Request) -> web.Response: """GET ``?root_index&root_scope&path``: the ids holding that file.""" query = request.query - if set(query) != {"root_index", "root_scope", "path"}: - return bad_request("Query must name root_index, root_scope and path") + if len(query) != 3 or set(query) != {"root_index", "root_scope", "path"}: + # len() counts a repeated key (?path=a&path=b) twice. + return bad_request("Query must name root_index, root_scope and path once") raw_index = query["root_index"] if not raw_index.isascii() or _ROOT_INDEX.fullmatch(raw_index) is None: return bad_request("Invalid gallery root") roots = list(self._get_all_output_dirs()) - root_key, key = self._collection_key( - { - "root_index": int(raw_index), - "root_scope": query["root_scope"], - "path": query["path"], - }, - roots, + target = { + "root_index": int(raw_index), + "root_scope": query["root_scope"], + "path": query["path"], + } + root_key, key = await self._run_in_executor( + self._collection_key, target, roots, {} ) ids = await self._run_in_executor(self.db.file_collections, root_key, key) return web.json_response({"success": True, "collection_ids": ids}) @@ -238,7 +251,7 @@ class CollectionRoutesMixin: else: target["collection_id"] = _body_id(data["collection_id"]) roots = list(self._get_all_output_dirs()) - keys = self._collection_items(data["items"], roots) + keys = await self._run_in_executor(self._collection_items, data["items"], roots) found, not_found = [], 0 listings = {} for index, root_key, key in keys: @@ -263,7 +276,7 @@ class CollectionRoutesMixin: data = await _body(request, (_REMOVE_FIELDS,)) collection_id = _body_id(data["collection_id"]) roots = list(self._get_all_output_dirs()) - keys = self._collection_items(data["items"], roots) + keys = await self._run_in_executor(self._collection_items, data["items"], roots) result = await self._run_in_executor( self.db.remove_from_collection, collection_id, @@ -294,30 +307,42 @@ class CollectionRoutesMixin: def _collection_items(self, items: object, roots: list) -> list[tuple]: """``(index, root key, relative key)`` per distinct file, in order. + Blocking (realpath per root, once per request): run in the executor. Every item is checked before anything is written: 400 for a malformed shape, 403 for an unsafe path, 409 for a stale scope. Two spellings that fold to one key (Windows) count once. """ if not isinstance(items, list) or not 1 <= len(items) <= MAX_ITEMS: raise _Refused(bad_request(f"items must list 1 to {MAX_ITEMS} images")) - keys = {} + keys, sources = {}, {} for item in items: index = item.get("root_index") if isinstance(item, dict) else None - root_key, key = self._collection_key(item, roots) + root_key, key = self._collection_key(item, roots, sources) keys.setdefault((root_key, key), index) return [(index, root_key, key) for (root_key, key), index in keys.items()] - def _collection_key(self, item: object, roots: list) -> tuple[str, str]: - """The stored keys of one target, or raise its refusal.""" + def _collection_key( + self, item: object, roots: list, sources: dict + ) -> tuple[str, str]: + """The stored keys of one target, or raise its refusal (blocking). + + *sources* caches ``{index: (scope, root key)}`` for one request, so + each root's realpath runs once however many items name it. + """ target, error = media_target(item, media_message=_MEDIA_MESSAGE) if error is not None: raise _Refused(error) index, scope, relative = target - if index >= len(roots) or not hmac.compare_digest( - scope, self._gallery_root_scope(roots[index], index) - ): + if index >= len(roots): + raise _Refused(_coded(409, "source-changed", SOURCES_CHANGED)) + if index not in sources: + sources[index] = ( + self._gallery_root_scope(roots[index], index), + _favorite_root_key(roots[index]), + ) + current, root_key = sources[index] + if not hmac.compare_digest(scope, current): raise _Refused(_coded(409, "source-changed", SOURCES_CHANGED)) - root_key = _favorite_root_key(roots[index]) key = _favorite_relative_key(relative) if max(len(root_key), len(key)) > MAX_KEY_LENGTH: # A stored key past the schema bound would fail the next restore. diff --git a/tests/test_api_collections.py b/tests/test_api_collections.py index c4dcdd86..daa0e53b 100644 --- a/tests/test_api_collections.py +++ b/tests/test_api_collections.py @@ -492,3 +492,110 @@ class TestCollectionListingFilter(CollectionRoutesTestCase): with self.assertRaises(Exception) as caught: await self.api.resolve_listing_filter({"collection": "77"}) self.assertEqual(type(caught.exception).__name__, "CollectionMissing") + + +class TestReviewFindings(CollectionRoutesTestCase): + """PR #400 review: loop work, counts snapshot, id and query rules.""" + + async def test_no_realpath_runs_on_the_event_loop(self): + import os + import threading + + names = [f"img{n}.png" for n in range(20)] + self.make(*names) + await self.call("GET", BASE) # warm the listing cache + loop_thread = threading.current_thread() + on_loop = [] + real = os.path.realpath + + def watched(*args, **kwargs): + if threading.current_thread() is loop_thread: + on_loop.append(args[0]) + return real(*args, **kwargs) + + items = [self.item(name) for name in names] + cid = None + with patch("os.path.realpath", side_effect=watched): + made = await self.add(name="Set", items=items) + cid = made["collection"]["id"] + await self.call("DELETE", ITEMS, {"collection_id": cid, "items": items[:5]}) + await self.call( + "GET", + MEMBERSHIP, + root_index="0", + root_scope=items[0]["root_scope"], + path="img9.png", + ) + await self.call("GET", BASE) + self.assertEqual(on_loop, []) + + async def test_missing_is_never_negative(self): + self.make("a.png", "b.png") + made = await self.add( + name="Set", items=[self.item("a.png"), self.item("b.png")] + ) + cid = made["collection"]["id"] + real = self.api.db.collection_count_snapshot + + def raced(root_keys): + members, stored = real(root_keys) + return members, {cid: 1} # a remove landed between two reads + + with patch.object(self.api.db, "collection_count_snapshot", side_effect=raced): + listed = await self.listed() + self.assertEqual((listed[0]["count"], listed[0]["missing"]), (2, 0)) + + async def test_path_ids_follow_the_listing_rule(self): + self.make("a.png") + made = await self.add(name="Set", items=[self.item("a.png")]) + cid = made["collection"]["id"] + for raw in ("0", "007", f"0{cid}"): + with self.subTest(raw=raw): + await self.call("PATCH", f"{BASE}/{raw}", {"name": "x"}, 400) + await self.call("DELETE", f"{BASE}/{raw}", None, 400) + self.assertEqual(len(await self.listed()), 1) + + async def test_membership_refuses_repeated_query_keys(self): + self.make("a.png", "b.png") + response = await self.client.get( + MEMBERSHIP + + f"?root_index=0&root_scope={self.scope()}&path=a.png&path=b.png" + ) + self.assertEqual(response.status, 400) + response = await self.client.get( + MEMBERSHIP + + f"?root_index=0&root_index=0&root_scope={self.scope()}&path=a.png" + ) + self.assertEqual(response.status, 400) + + async def test_deleting_a_collection_leaves_everything_else(self): + from py.api.images import _favorite_root_key + + try: + from tests.test_annotation_schema import insert_group_rows, table_rows + except ImportError: + from test_annotation_schema import insert_group_rows, table_rows + + self.make("a.png") + db = self.api.db + db.set_image_favorite(_favorite_root_key(self.output_dir), "a.png") + prompt = self._save_prompt("Linked") + db.link_image_to_prompt(prompt, str(self.output_dir / "a.png")) + with db.model.get_connection() as conn: + insert_group_rows(conn, locator=str(self.output_dir / "a.png")) + made = await self.add(name="Set", items=[self.item("a.png")]) + tables = ( + "image_favorites", + "generated_images", + "prompts", + "tags", + "media_annotation_files", + "media_annotation_bytes", + "media_image_tags", + ) + before = table_rows(db.model.get_connection(), tables) + file_bytes = (self.output_dir / "a.png").read_bytes() + await self.call("DELETE", f"{BASE}/{made['collection']['id']}") + self.assertEqual(table_rows(db.model.get_connection(), tables), before) + self.assertEqual((self.output_dir / "a.png").read_bytes(), file_bytes) + self.assertTrue(all(before[t] for t in tables)) diff --git a/tests/test_image_collections.py b/tests/test_image_collections.py index 9136c862..83307b02 100644 --- a/tests/test_image_collections.py +++ b/tests/test_image_collections.py @@ -4,6 +4,7 @@ import sqlite3 import tempfile import unittest from pathlib import Path +from unittest.mock import patch try: from tests.open_handles import assert_closed @@ -175,7 +176,7 @@ class ImageCollectionOperationsTests(unittest.TestCase): self.require() made = self.add([(ROOT, "a.png")], name=" LoRa \t testing\n ")["collection"] self.assertEqual(made["name"], "LoRa testing") - for bad in ("", " ", "x" * 101, "bell\x07", 5, None): + for bad in ("", " ", "x" * 101, "bell\x07", "c1\x81", "del\x7f", 5, None): with self.subTest(name=bad): with self.assertRaises(CollectionInvalid): self.add([(ROOT, "a.png")], name=bad) @@ -199,6 +200,50 @@ class ImageCollectionOperationsTests(unittest.TestCase): with self.assertRaises(CollectionInvalid): self.add(items, collection_id=made["id"]) + def test_a_failure_after_the_membership_move_rolls_it_back(self): + self.require() + made = self.add([(ROOT, "b.png")], name="Grids")["collection"] + move = { + "removed": [], + "kept": [], + "targets": {}, + "kept_path": "/real/output/a.png", + "favorites": [(ROOT, "b.png", ROOT, "a.png")], + "removed_locator": "/real/output/b.png", + "kept_locator": "/real/output/a.png", + "sha256": "0" * 64, + "kept_size": 1, + "kept_mtime_ns": 1, + } + with patch.object( + type(self.db), "_move_copy_tags", side_effect=RuntimeError("boom") + ): + with self.assertRaises(RuntimeError): + self.db.move_copy_records(move) + self.assertEqual(self.db.file_collections(ROOT, "b.png"), [made["id"]]) + self.assertEqual(self.db.file_collections(ROOT, "a.png"), []) + # Without the failure the same move succeeds. + self.assertEqual(self.db.move_copy_records(move)["collections"], 1) + self.assertEqual(self.db.file_collections(ROOT, "a.png"), [made["id"]]) + + def test_count_snapshot_reads_in_one_transaction(self): + self.require() + made = self.add([(ROOT, "a.png"), (OTHER, "b.png")], name="Set")["collection"] + statements = [] + conn = self.db.model.get_connection() + conn.set_trace_callback(statements.append) + try: + members, stored = self.db.collection_count_snapshot([ROOT, OTHER, ROOT]) + finally: + conn.set_trace_callback(None) + self.assertEqual( + members, {ROOT: {made["id"]: {"a.png"}}, OTHER: {made["id"]: {"b.png"}}} + ) + self.assertEqual(stored, {made["id"]: 2}) + self.assertEqual(statements[0], "BEGIN") + self.assertEqual(statements[-1], "ROLLBACK") + self.assertFalse(conn.in_transaction) + def test_reads_use_an_index(self): self.require() self.add([(ROOT, "a.png")], name="Set") diff --git a/tests/test_image_collections_schema.py b/tests/test_image_collections_schema.py index fb06ba2d..df89faf7 100644 --- a/tests/test_image_collections_schema.py +++ b/tests/test_image_collections_schema.py @@ -390,7 +390,8 @@ class ImageCollectionsSchemaTests(unittest.TestCase): ) def test_restore_refuses_a_non_text_name(self): - self.assert_refused("int-name", f"UPDATE {COLLECTIONS} SET name = 5") + # Column affinity turns a number into text; a blob stays a blob. + self.assert_refused("blob-name", f"UPDATE {COLLECTIONS} SET name = x'61'") def test_restore_refuses_an_empty_name(self): self.assert_refused("empty", f"UPDATE {COLLECTIONS} SET name = '', name_key=''") @@ -420,8 +421,43 @@ class ImageCollectionsSchemaTests(unittest.TestCase): ), ) - def test_restore_refuses_a_wrong_name_key(self): - self.assert_refused("key", f"UPDATE {COLLECTIONS} SET name_key = 'Good'") + def test_restore_accepts_a_name_key_from_other_case_tables(self): + # Python versions fold some characters differently; a key written + # by another version must never make a backup unrestorable. + self.require_module() + path = self.closed_current( + "other-key.db", + ( + f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) " + "VALUES (1, 'Good', 'Good', ?)", + (NOW,), + ), + ) + self.assertEqual( + PromptModel.verify_database_file(path, require_library_state=True), + (True, "ok"), + ) + + def test_restore_refuses_an_empty_or_non_text_name_key(self): + self.assert_refused("empty-key", f"UPDATE {COLLECTIONS} SET name_key = ''") + self.assert_refused( + "blob-name-key", f"UPDATE {COLLECTIONS} SET name_key = x'61'" + ) + self.assert_refused( + "long-name-key", + (f"UPDATE {COLLECTIONS} SET name_key = ?", ("k" * 1201,)), + ) + + def test_restore_refuses_a_c1_control_character(self): + for name in ("a\x81b", "a\x7fb"): + with self.subTest(name=repr(name)): + self.assert_refused( + f"c1-{ord(name[1])}", + ( + f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", + (name, key(name)), + ), + ) def test_restore_refuses_a_non_text_key(self): # Column affinity turns a number into text; a blob stays a blob.