fix(collections): address review findings

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)

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.

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

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

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.