diff --git a/TESTING.md b/TESTING.md index c606358e..c1f60a29 100644 --- a/TESTING.md +++ b/TESTING.md @@ -4516,11 +4516,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 @@ -4540,7 +4542,11 @@ category; without one it shows the category its link supplies. - 409 `source-changed`: a stale `root_scope`; - 413: a body over 1 MiB. Any invalid item fails the whole request before a write. Targets that fold to one - key are written and counted once. + key are written and counted once. Instead of `items` a request may send `filter`, + the Images listing's own query ("Select all N matching", #392, keys as in + `parse_listing_filter`), resolved now: over 500 matches is a 400 + `too-many-images` ("Select at most 500 images."), a renumbered source a 409 + `source-changed`. - PUT writes only files the source's listing holds and counts the rest in `not_found`; when none is listed it answers 404 `files-not-found` and writes nothing. A file that already has the value counts as `unchanged` and keeps its 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 9b480de2..ff37ecd4 100644 --- a/py/api/image_categories.py +++ b/py/api/image_categories.py @@ -6,6 +6,8 @@ category can still be removed); ``GET /prompt_manager/images/categories`` lists the own-category values of files the sources list. Targets are the Images listing identity (``{root_index, root_scope, path}``, checked by ``media_target``) of an output source; LoRA folders are not sources here. +Instead of ``items`` a request may name ``filter``, the Images listing's own +query ("Select all N matching", #392), resolved now up to 500 images. Each file is stored under its file key (``images.image_category_key``). These routes never read or write prompts, recipes or image links. """ @@ -35,8 +37,9 @@ except ImportError: from utils.validators import validate_category CATEGORY_MEDIA_MESSAGE = "Only media files can have a category" -_PUT_FIELDS = frozenset({"items", "category"}) -_DELETE_FIELDS = frozenset({"items"}) +_PUT_FIELDS = frozenset({"category"}) +_DELETE_FIELDS = frozenset() +TOO_MANY = f"Select at most {MAX_JOB_TARGETS} images." SOURCE_CHANGED_MESSAGE = "The sources changed. Refresh the images and try again." WRITE_FAILED = "The category could not be changed." READ_FAILED = "Image categories could not be read." @@ -62,12 +65,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*. @@ -103,7 +127,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) @@ -186,20 +212,53 @@ class ImageCategoryRoutesMixin: # -- writing --------------------------------------------------------------- + async def _filter_targets(self, raw): + """``(targets, None)`` for an Images listing filter, else ``(None, error)``. + + "Select all N matching" (#392) sends the listing's own query; it is + resolved now, up to 500 images (a 400 ``too-many-images`` above). + """ + from .images import ListingSourceChanged, parse_listing_filter + + try: + selection = await self.resolve_listing_filter( + parse_listing_filter(raw), MAX_JOB_TARGETS + ) + except ListingSourceChanged: + return None, _coded(409, "source-changed", SOURCE_CHANGED_MESSAGE) + except ValueError as exc: + return None, bad_request(str(exc)) + if selection.truncated: + return None, _coded(400, "too-many-images", TOO_MANY) + return tuple(selection.targets), None + async def _category_request(self, request, fields: frozenset): - """``(data, targets, roots, None)`` or ``(None, None, None, error)``.""" + """``(data, targets, roots, None)`` or ``(None, None, None, error)``. + + The images are ``items`` (1 to 500 targets) or ``filter`` (all images + an Images listing filter admits), besides *fields*. + """ data, error = await _json_body_or_error(request) if error is not None: return None, None, None, error - if set(data) != fields: + names = set(data) - fields + if set(data) & fields != fields or names not in ({"items"}, {"filter"}): return ( None, None, None, bad_request( - "Request must contain exactly " + ", ".join(sorted(fields)) + "Request must contain " + + ", ".join(sorted(fields)) + + (" and " if fields else "") + + "either items or filter" ), ) + if "filter" in data: + targets, error = await self._filter_targets(data["filter"]) + if error is not None: + return None, None, None, error + return data, targets, list(self._get_all_output_dirs()), None items = data["items"] if not isinstance(items, list) or not 1 <= len(items) <= MAX_JOB_TARGETS: return ( @@ -211,6 +270,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): @@ -224,6 +288,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) @@ -233,22 +314,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, @@ -256,7 +333,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") @@ -279,16 +356,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 fbcd4986..f1dfd718 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -812,25 +812,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..02905267 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 = [] @@ -301,3 +322,54 @@ class TestCategoryList(ImageCategoryTestCase): data = await self.listed(status=500) self.assertEqual(data["error"], "Image categories could not be read.") self.assertNotIn("secret", json.dumps(data)) + + +class TestAllMatchingFilter(ImageCategoryTestCase): + """#392 "Select all N matching": the listing filter, resolved now.""" + + async def test_filter_sets_and_clears_every_match(self): + for name in ("bulk/a.png", "bulk/b.png", "other/c.png"): + make_png(self.output_dir / name) + data = await self.send( + "PUT", {"filter": {"subfolder": "bulk"}, "category": "Scans"} + ) + self.assertEqual((data["updated"], data["not_found"]), (2, 0)) + self.assertEqual( + sorted(row[1] for row in self.rows()), ["bulk/a.png", "bulk/b.png"] + ) + data = await self.send("DELETE", {"filter": {"subfolder": "bulk"}}) + self.assertEqual(data["cleared"], 2) + data = await self.send( + "PUT", {"filter": {"subfolder": "none"}, "category": "Scans"}, 404 + ) + self.assertEqual(data["code"], "files-not-found") + self.assertEqual( + (await self.send("DELETE", {"filter": {"q": "zz"}}))["cleared"], 0 + ) + + async def test_filter_refusals(self): + make_png(self.output_dir / "a.png") + item = self.target("a.png") + for body in ( + {"filter": {}, "items": [item], "category": "x"}, + {"filter": {"unknown": "1"}, "category": "x"}, + {"filter": "q=a", "category": "x"}, + {"filter": {"root": "0", "root_scope": "nothex"}, "category": "x"}, + ): + with self.subTest(body=body): + await self.send("PUT", body, 400) + data = await self.send( + "PUT", + {"filter": {"root": "0", "root_scope": "0" * 64}, "category": "x"}, + 409, + ) + self.assertEqual(data["code"], "source-changed") + self.assertEqual(self.rows(), []) + + async def test_more_than_500_matching_is_refused(self): + for number in range(501): + make_png(self.output_dir / f"many/{number:03}.png", size=(1, 1)) + data = await self.send("PUT", {"filter": {}, "category": "x"}, 400) + self.assertEqual(data["code"], "too-many-images") + self.assertEqual(data["error"], "Select at most 500 images.") + self.assertEqual(self.rows(), []) 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