fix(images): key symlinked output folders and keep key work off the event loop (#332)

Review fixes: a source matches by its configured spelling too, so a
detected output folder that is a symlink keys its files (routes and
duplicate cleanup). File keys are computed in the executor. A path no
source holds is a plain 403, a lone surrogate a 400. The install repairs
a table that lacks its category index instead of refusing to open.

TESTING.md

@@ -4484,11 +4484,13 @@ category; without one it shows the category its link supplies.
 - The **file key** is `image_category_key(sources, absolute path)` in
   `py/api/images.py`: the deepest output source that contains the path (compared
   lexically, never resolving the path, so a symlinked subfolder keeps its listed
-  spelling) and the path below it, keyed with the favorites helpers. Without
-  nested sources it equals a favorite's key. Every write and read uses it. A path
+  spelling) and the path below it, keyed with the favorites helpers. A source
+  matches by its real path or its configured spelling (a detected output folder
+  can be a symlink). Without nested sources it equals a favorite's key. Every write and read uses it. A path
   that no source contains has no key, so it has no own category.
 - `image_categories_schema_error` checks the structure at startup (a value saved
-  under older rules never stops the library opening). Restore also checks the
+  under older rules never stops the library opening; a table without its index,
+  which this code never leaves, is repaired by the install). Restore also checks the
   values: text columns, keys within the favorites bound, `category` 1 to 100
   characters with no character below U+0020 and no leading or trailing space, and
   `updated_at` of at most 64 bytes. Backups made before #332 restore without the

database/image_categories.py

@@ -65,7 +65,9 @@ def _compact(sql: str) -> str:
     return re.sub(r"\s+", "", sql or "").lower()
 
 
-def _structure_error(conn: sqlite3.Connection, allow_missing: bool) -> str | None:
+def _structure_error(
+    conn: sqlite3.Connection, allow_missing: bool, allow_missing_index: bool = False
+) -> str | None:
     objects = conn.execute(
         "SELECT type,name,sql FROM sqlite_master WHERE name=? "
         "OR (type IN ('index','trigger') AND tbl_name=?)",
@@ -75,12 +77,17 @@ def _structure_error(conn: sqlite3.Connection, allow_missing: bool) -> str | Non
         return None if allow_missing else "Missing image categories schema"
     tables = [sql for kind, _name, sql in objects if kind == "table"]
     others = {(kind, name) for kind, name, _sql in objects if kind != "table"}
-    if len(tables) != 1 or others != {("index", _AUTOINDEX), ("index", INDEX)}:
+    expected = {("index", _AUTOINDEX), ("index", INDEX)}
+    # Before its install, a table without the category index is repaired by
+    # CREATE INDEX IF NOT EXISTS rather than refused.
+    if allow_missing_index and others == {("index", _AUTOINDEX)}:
+        expected = others
+    if len(tables) != 1 or others != expected:
         return "Image categories name is reserved by an unsupported schema object"
     if _compact(tables[0]) != _compact(_DDL):
         return "Image categories has an invalid table definition"
-    index_sql = next(sql for kind, name, sql in objects if name == INDEX)
-    if _compact(index_sql) != _compact(_INDEX_DDL):
+    index_sql = next((sql for kind, name, sql in objects if name == INDEX), None)
+    if ("index", INDEX) in expected and _compact(index_sql) != _compact(_INDEX_DDL):
         return "Image categories has an invalid index"
     autoindex_sql = next(sql for kind, name, sql in objects if name == _AUTOINDEX)
     if autoindex_sql is not None:
@@ -148,7 +155,7 @@ def migrate_image_categories(conn: sqlite3.Connection) -> None:
     and writes only its own two objects (no backfill)."""
     conn.execute("SAVEPOINT image_categories_schema")
     try:
-        error = image_categories_schema_error(conn, check_values=False)
+        error = _structure_error(conn, True, allow_missing_index=True)
         if error:
             raise sqlite3.DatabaseError(error)
         conn.execute(_DDL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1))

py/api/image_categories.py

@@ -60,12 +60,33 @@ def _category_value(raw) -> tuple[str | None, web.Response | None]:
             f"Category cannot exceed {MAX_CATEGORY_LENGTH} characters"
         )
     try:
+        value.encode("utf-8")
         validate_category(value)
+    except UnicodeEncodeError:
+        return None, bad_request("category must be valid text")
     except ValueError as exc:
         return None, bad_request(str(exc))
     return value, None
 
 
+class _KeyRefused(Exception):
+    """A target has no storable file key (outside every source, or too long)."""
+
+    def __init__(self, message: str = "Access denied"):
+        super().__init__(message)
+        self.message = message
+
+
+def _file_key(roots, root, relative) -> tuple[str, str]:
+    """The target's file key (blocking: resolves the sources)."""
+    key = image_category_key(roots, Path(root) / relative)
+    if key is None:
+        raise _KeyRefused()
+    if max(len(key[0]), len(key[1])) > MAX_KEY_LENGTH:
+        raise _KeyRefused("Path is too long")
+    return key
+
+
 def _nested_sources(root: Path, sources) -> list[tuple[str, str]]:
     """``[(relative key prefix, source root key)]`` of sources inside *root*.
 
@@ -101,7 +122,9 @@ class ImageCategoryRoutesMixin:
 
     def _register_image_category_routes(self, routes):
         # Fixed paths beside DELETE /prompt_manager/images/{image_id}, as the
-        # favorite routes are: matched before the dynamic resource.
+        # favorite routes are. Registered after it, they still win: aiohttp
+        # (3.10+) looks fixed paths up before dynamic resources
+        # (test_delete_category_never_reaches_delete_image).
         @routes.put("/prompt_manager/images/category")
         async def set_image_category_route(request):
             return await self.set_image_category(request)
@@ -198,6 +221,11 @@ class ImageCategoryRoutesMixin:
         targets, error = _parse_targets(items, media_message=CATEGORY_MEDIA_MESSAGE)
         if error is not None:
             return None, None, None, error
+        try:
+            for _index, _scope, relative in targets:
+                relative.encode("utf-8")
+        except UnicodeEncodeError:
+            return None, None, None, bad_request("path must be valid text")
         roots = list(self._get_all_output_dirs())
         for index, scope in dict.fromkeys((t[0], t[1]) for t in targets):
             if index >= len(roots):
@@ -211,6 +239,23 @@ class ImageCategoryRoutesMixin:
                 return None, None, None, error
         return data, targets, roots, None
 
+    def _listed_file_keys(self, roots, targets, files) -> tuple[list, list]:
+        """``(listed keys, unlisted keys)`` of the targets (blocking).
+
+        A target counts as listed when its source's listing holds it; its key
+        then follows the listed spelling. Targets folding to one key count
+        once, and a key that is listed through any target is listed.
+        """
+        listings, listed, missing = {}, {}, {}
+        for index, _scope, relative in targets:
+            root = roots[index]
+            if index not in listings:
+                listings[index] = self._favorite_listing(root, files[index])
+            spelled = listings[index].get(_favorite_relative_key(relative))
+            key = _file_key(roots, root, spelled or relative)
+            (missing if spelled is None else listed)[key] = True
+        return list(listed), [key for key in missing if key not in listed]
+
     async def set_image_category(self, request):
         """PUT ``{items, category}``: set the own category of listed files."""
         data, targets, roots, error = await self._category_request(request, _PUT_FIELDS)
@@ -220,22 +265,18 @@ class ImageCategoryRoutesMixin:
         if error is not None:
             return error
         try:
-            listings = {}
-            listed_keys, missing_keys = {}, {}
-            for index, _scope, relative in targets:
-                root = roots[index]
-                if index not in listings:
-                    files = await self._get_gallery_files(root)
-                    listings[index] = await self._run_in_executor(
-                        self._favorite_listing, root, files
-                    )
-                listed = listings[index].get(_favorite_relative_key(relative))
-                key = image_category_key(roots, root / (listed or relative))
-                if key is None or max(len(key[0]), len(key[1])) > MAX_KEY_LENGTH:
-                    return _forbidden("Path is too long")
-                (missing_keys if listed is None else listed_keys)[key] = True
-            for key in listed_keys:
-                missing_keys.pop(key, None)
+            files = {}
+            for index in dict.fromkeys(t[0] for t in targets):
+                files[index] = await self._get_gallery_files(roots[index])
+            listed_keys, missing_keys = await self._run_in_executor(
+                self._listed_file_keys, roots, targets, files
+            )
+        except _KeyRefused as refused:
+            return _forbidden(refused.message)
+        except Exception:
+            self.logger.exception("Image category update failed")
+            return _coded(500, "category-failed", WRITE_FAILED)
+        try:
             if not listed_keys:
                 return _coded(
                     404,
@@ -243,7 +284,7 @@ class ImageCategoryRoutesMixin:
                     "None of the selected images could be found.",
                 )
             counts = await self._run_in_executor(
-                self.db.set_image_categories, list(listed_keys), category
+                self.db.set_image_categories, listed_keys, category
             )
         except Exception:
             self.logger.exception("Image category update failed")
@@ -266,16 +307,15 @@ class ImageCategoryRoutesMixin:
         )
         if error is not None:
             return error
-        keys = {}
-        for index, _scope, relative in targets:
-            key = image_category_key(roots, roots[index] / relative)
-            if key is None or max(len(key[0]), len(key[1])) > MAX_KEY_LENGTH:
-                return _forbidden("Path is too long")
-            keys[key] = True
         try:
-            cleared = await self._run_in_executor(
-                self.db.clear_image_categories, list(keys)
+            keys = await self._run_in_executor(
+                lambda: list(
+                    dict.fromkeys(_file_key(roots, roots[i], r) for i, _s, r in targets)
+                )
             )
+            cleared = await self._run_in_executor(self.db.clear_image_categories, keys)
+        except _KeyRefused as refused:
+            return _forbidden(refused.message)
         except Exception:
             self.logger.exception("Image category removal failed")
             return _coded(500, "category-failed", WRITE_FAILED)

py/api/images.py

@@ -807,25 +807,29 @@ def _source_depth(source: str) -> int:
 def deepest_source(sources, absolute_path):
     """``(source, POSIX path below it)`` of the deepest source holding the path.
 
-    *sources* are compared by their real path (they are stored resolved); the
-    path is compared lexically after ``normpath`` and never resolved, so a
+    A source matches by its real path or by its spelling as configured (a
+    detected output folder can be a symlink), and the deepest real path wins.
+    The path is compared lexically after ``normpath`` and never resolved, so a
     symlinked subfolder keeps its listed spelling. Containment folds case
-    where the OS does. None when no source contains the path.
+    where the OS does. *source* is the real path. None when no source
+    contains the path.
     """
     path = os.path.normpath(str(absolute_path))
     folded = os.path.normcase(path)
     best = None
     for source in sources:
         real = os.path.normpath(os.path.realpath(str(source)))
-        prefix = os.path.join(os.path.normcase(real), "")
-        if not folded.startswith(prefix):
-            continue
-        if best is None or _source_depth(real) > _source_depth(best):
-            best = real
+        for spelling in dict.fromkeys((real, os.path.normpath(str(source)))):
+            prefix = os.path.join(os.path.normcase(spelling), "")
+            if not folded.startswith(prefix):
+                continue
+            if best is None or _source_depth(real) > _source_depth(best[0]):
+                best = (real, spelling)
+            break
     if best is None:
         return None
-    relative = os.path.relpath(path, best)
-    return best, relative.replace(os.sep, "/")
+    relative = os.path.relpath(path, best[1])
+    return best[0], relative.replace(os.sep, "/")
 
 
 def image_category_key(sources, absolute_path):

tests/test_api_image_categories.py

@@ -201,6 +201,27 @@ class TestSetAndClear(ImageCategoryTestCase):
             data = await self.delete(["a.png"], status=500)
         self.assertNotIn("secret", json.dumps(data))
 
+    async def test_a_symlinked_output_folder_keys_like_its_favorites(self):
+        """A detected output folder can be a symlink: keys use its real path."""
+        real = self.comfy_dir / "real-output"
+        make_png(real / "a.png")
+        link = self.comfy_dir / "linked-output"
+        self._make_symlink(real, link, target_is_directory=True)
+        self.api._get_all_output_dirs = lambda: [link]
+        data = await self.put(["a.png"], "Scans")
+        self.assertEqual(data["updated"], 1)
+        self.assertEqual(self.rows(), [(_favorite_root_key(link), "a.png", "Scans")])
+        self.assertEqual((await self.listed())["categories"], ["Scans"])
+        self.assertEqual((await self.delete(["a.png"]))["cleared"], 1)
+
+    async def test_lone_surrogates_are_400_not_500(self):
+        make_png(self.output_dir / "a.png")
+        await self.send(
+            "PUT", {"items": [self.target("a.png")], "category": "x\ud800"}, 400
+        )
+        await self.send("DELETE", {"items": [self.target("a\ud800.png")]}, 400)
+        self.assertEqual(self.rows(), [])
+
     async def test_delete_category_never_reaches_delete_image(self):
         make_png(self.output_dir / "a.png")
         calls = []

tests/test_duplicate_image_categories.py

@@ -134,3 +134,18 @@ class TestDuplicateCleanupKeepsCategories(DuplicateFilesTestCase):
         await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
         self.assertFalse((self.output_dir / "sub" / "b.png").exists())
         self.assertEqual(self.api.db.image_category(sub_key, "a.png"), "Portraits")
+
+    async def test_a_symlinked_output_folder_still_moves_the_category(self):
+        link = self.output_dir.parent / "linked-output"
+        try:
+            os.symlink(self.output_dir, link, target_is_directory=True)
+        except (OSError, NotImplementedError):
+            self.skipTest("symlinks are not available")
+        self.trio()
+        with patch.object(self.api, "_get_all_output_dirs", return_value=[link]):
+            key = image_category_key([link], link / "b.png")
+            self.api.db.set_image_categories([key], "Portraits")
+            await self.run_cleanup()
+            kept = image_category_key([link], link / "a.png")
+        self.assertFalse((self.output_dir / "b.png").exists())
+        self.assertEqual(self.api.db.image_category(*kept), "Portraits")

tests/test_image_categories.py

@@ -195,7 +195,6 @@ class TestImageCategoriesSchema(_TempDir):
                 _INDEX_DDL,
             ),
             "wrong index": (_DDL, f"CREATE INDEX {INDEX} ON {TABLE}(updated_at)"),
-            "no index": (_DDL,),
             "extra index": (
                 _DDL,
                 _INDEX_DDL,
@@ -231,6 +230,17 @@ class TestImageCategoriesSchema(_TempDir):
                 )
                 with self.assertRaises(sqlite3.DatabaseError):
                     migrate_image_categories(conn)
+        # A table without its category index (never written by this code) is
+        # refused by the check but repaired by the install.
+        with closing(sqlite3.connect(":memory:")) as conn:
+            conn.execute(_DDL)
+            conn.execute(f"INSERT INTO {TABLE} VALUES ('r', 'a.png', 'Scans', 'now')")
+            self.assertIsNotNone(image_categories_schema_error(conn))
+            migrate_image_categories(conn)
+            self.assertIsNone(image_categories_schema_error(conn, allow_missing=False))
+            self.assertEqual(
+                conn.execute(f"SELECT COUNT(*) FROM {TABLE}").fetchone()[0], 1
+            )
         with closing(sqlite3.connect(":memory:")) as conn:
             migrate_image_categories(conn)
             migrate_image_categories(conn)  # idempotent