nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge branch 'feat/332-image-categories-pr1' into feat/332-image-categories-pr2

Browse Source
This commit is contained in:
Vito Sansevero
2026-10-27 18:53:41 -07:00
parent a4e76b3b9f 20b6514882
commit 6cd071b8e6
7 changed files with 241 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TESTING.md
+10 -4
View File
@@ -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
database/image_categories.py
+12 -5
View File
@@ -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
+107 -31
View File
@@ -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)
py/api/images.py
+14 -10
View File
@@ -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):
tests/test_api_image_categories.py
+72
View File
@@ -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(), [])
tests/test_duplicate_image_categories.py
+15
View File
@@ -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
+11 -1
View File
@@ -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