From 0de5c0a4e76f6881efad4791c6dc7f24c8ef07c4 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Tue, 27 Oct 2026 18:31:43 -0700 Subject: [PATCH] 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 | 8 ++- database/image_categories.py | 17 +++-- py/api/image_categories.py | 92 +++++++++++++++++------- py/api/images.py | 24 ++++--- tests/test_api_image_categories.py | 21 ++++++ tests/test_duplicate_image_categories.py | 15 ++++ tests/test_image_categories.py | 12 +++- 7 files changed, 144 insertions(+), 45 deletions(-) diff --git a/TESTING.md b/TESTING.md index 26898c0c..54eba39e 100644 --- a/TESTING.md +++ b/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 diff --git a/database/image_categories.py b/database/image_categories.py index b8657f1d..a1f5d5a1 100644 --- a/database/image_categories.py +++ b/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)) diff --git a/py/api/image_categories.py b/py/api/image_categories.py index ac2182e0..381b6059 100644 --- a/py/api/image_categories.py +++ b/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) diff --git a/py/api/images.py b/py/api/images.py index dc28caec..e68255a4 100644 --- a/py/api/images.py +++ b/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): diff --git a/tests/test_api_image_categories.py b/tests/test_api_image_categories.py index 18cd93e4..5e06368a 100644 --- a/tests/test_api_image_categories.py +++ b/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 = [] diff --git a/tests/test_duplicate_image_categories.py b/tests/test_duplicate_image_categories.py index c129fcbe..79452a68 100644 --- a/tests/test_duplicate_image_categories.py +++ b/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") diff --git a/tests/test_image_categories.py b/tests/test_image_categories.py index e6560d56..9972c70e 100644 --- a/tests/test_image_categories.py +++ b/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