diff --git a/py/api/__init__.py b/py/api/__init__.py index 485578c7..894673b9 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -48,6 +48,7 @@ from .review import ReviewRoutesMixin from .review_recording import ReviewRecordingMixin from .review_events import ReviewEventsMixin from .saved_views import SavedViewRoutesMixin +from .collections import CollectionRoutesMixin from .history import HistoryRoutesMixin from .images import ImageRoutesMixin from .gallery_listing import WALK_STOPPED, RootListing, WalkOutcome @@ -314,6 +315,7 @@ class PromptManagerAPI( ReviewRecordingMixin, ReviewEventsMixin, SavedViewRoutesMixin, + CollectionRoutesMixin, HistoryRoutesMixin, PromptRoutesMixin, ImageRoutesMixin, @@ -684,6 +686,7 @@ class PromptManagerAPI( self._register_review_routes(routes) self._register_review_event_routes(routes) self._register_saved_view_routes(routes) + self._register_collection_routes(routes) self._register_history_routes(routes) self._register_prompt_routes(routes) self._register_image_routes(routes) diff --git a/py/api/annotations.py b/py/api/annotations.py index 47f905e6..d9a27a0f 100644 --- a/py/api/annotations.py +++ b/py/api/annotations.py @@ -27,7 +27,12 @@ from .annotation_suggest import ( validate_options, ) from .autotag_routes import model_busy_response -from .images import _forbidden, _has_traversal, _unavailable_link, _verified_title + +# The target check is shared with the collection routes (#331); both names +# stay importable from here. +from .media_target import MAX_TARGET_PATH # noqa: F401 +from .media_target import media_target as _target +from .images import _unavailable_link, _verified_title from .prompts import bad_request from .. import model_gate @@ -40,7 +45,7 @@ try: ) from ...database.execution_history import subscribe_history_restore from ...database.library_commands import LibraryCommandError - from ...utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore + from ...utils.media_capture import CaptureError, CaptureStore except ImportError: from database.annotations import ( MAX_IMAGE_TAGS, @@ -50,15 +55,13 @@ except ImportError: ) from database.execution_history import subscribe_history_restore from database.library_commands import LibraryCommandError - from utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore + from utils.media_capture import CaptureError, CaptureStore # Request bodies of the annotation routes are small (a target plus a tag set). MAX_BODY_BYTES = 64 * 1024 -MAX_TARGET_PATH = 1024 CAPTURE_URL = "/prompt_manager/image-annotations/capture/" _PREVIEW_CHUNK = 64 * 1024 -_TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"}) _OBSERVE_FIELDS = frozenset({"target"}) _SAVE_FIELDS = frozenset( { @@ -188,37 +191,6 @@ def _capture_error(exc: CaptureError) -> web.Response: return _coded(exc.code, _MESSAGES.get(exc.code) or str(exc)) -def _target(data: object): - """``((index, scope, relative), None)`` for a target, else ``(None, error)``. - - Only the listing identity is accepted: an int root index, the root's - 64-hex scope and a bounded root-relative POSIX media path, never an - absolute, drive-qualified, backslash or ``..`` path. Nothing here touches - the filesystem or the database. - """ - if not isinstance(data, dict) or set(data) != _TARGET_FIELDS: - return None, bad_request("target must name root_index, root_scope and path") - index, scope, relative = data["root_index"], data["root_scope"], data["path"] - if type(index) is not int or index < 0: - return None, bad_request("Invalid gallery root") - if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None: - return None, bad_request("Invalid gallery root scope") - if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH: - return None, bad_request("path must be a relative media path") - parts = relative.split("/") - if ( - "\0" in relative - or "\\" in relative - or relative.startswith("/") - or _has_traversal(relative) - or any(part in ("", ".", "..") or ":" in part for part in parts) - ): - return None, _forbidden() - if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS: - return None, _forbidden("Only media files can be tagged") - return (index, scope, relative), None - - def _expires_at(seconds: float | None) -> str: moment = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta( seconds=max(0.0, seconds or 0.0) diff --git a/py/api/collections.py b/py/api/collections.py new file mode 100644 index 00000000..51fab44d --- /dev/null +++ b/py/api/collections.py @@ -0,0 +1,325 @@ +"""Collections (#331): named, hand-picked groups of images. + +Routes under ``/prompt_manager/collections``. A file is named by its listing +identity ``{root_index, root_scope, path}`` (``media_target.py``) and stored +under the favorites' keys (``_favorite_root_key``, ``_favorite_relative_key``), +which never leave the server. A stale scope (the sources were reordered) is a +409 ``source-changed`` before anything is written. Responses are +``{"success": bool, ...}``; an unexpected failure is a logged 500 with a +generic message, never the exception text. +""" + +import hmac +import re +from functools import wraps +from typing import Callable + +from aiohttp import web + +from .images import ( + _favorite_relative_key, + _favorite_root_key, + _forbidden, + _json_body_or_error, + collection_not_found, +) +from .media_target import media_target +from .prompts import bad_request + +try: + from ...database.image_collections import ( + CollectionFilesMissing, + CollectionInvalid, + CollectionMissing, + CollectionNameTaken, + normalize_collection_name, + ) + from ...database.image_favorites import MAX_KEY_LENGTH +except ImportError: + from database.image_collections import ( + CollectionFilesMissing, + CollectionInvalid, + CollectionMissing, + CollectionNameTaken, + normalize_collection_name, + ) + from database.image_favorites import MAX_KEY_LENGTH + +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}") +_ROOT_INDEX = re.compile(r"[0-9]{1,9}") +_ADD_TARGETS = (frozenset({"collection_id", "items"}), frozenset({"name", "items"})) +_REMOVE_FIELDS = frozenset({"collection_id", "items"}) +_RENAME_FIELDS = frozenset({"name"}) +_MEDIA_MESSAGE = "Only media files can be added to a collection" +READ_FAILED = "Collections could not be read." +CHANGE_FAILED = "The collection could not be changed." +SOURCES_CHANGED = "The sources changed. Refresh the images and try again." +FILES_NOT_FOUND = ( + "None of the selected images could be found. Refresh the images and try again." +) +NAME_TAKEN = "Another collection already has this name." + + +class _Refused(Exception): + """A request refused before any write; carries its response.""" + + def __init__(self, response: web.Response): + super().__init__(response.status) + self.response = response + + +def _coded(status: int, code: str, message: str, **extra) -> web.Response: + return web.json_response( + {"success": False, "code": code, "error": message, **extra}, + status=status, + headers={"Cache-Control": "no-store"}, + ) + + +def _collection_response(failure: str) -> Callable: + """Map the database's refusals to their status; anything else to a 500.""" + + def decorate(handler: Callable) -> Callable: + @wraps(handler) + async def wrapped(self, request: web.Request) -> web.Response: + try: + return await handler(self, request) + except _Refused as refused: + return refused.response + except CollectionInvalid as exc: + return bad_request(str(exc)) + except CollectionMissing: + return collection_not_found() + except CollectionFilesMissing: + return _coded(404, "files-not-found", FILES_NOT_FOUND) + except CollectionNameTaken as exc: + return _coded( + 409, "name-taken", NAME_TAKEN, collection_id=exc.collection_id + ) + except Exception: + self.logger.exception("Collection request failed") + return web.json_response( + {"success": False, "error": failure}, status=500 + ) + + return wrapped + + return decorate + + +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: + raise _Refused(bad_request("Invalid collection id")) + return int(raw) + + +def _body_id(value: object) -> int: + if type(value) is not int or not 1 <= value <= MAX_ID: + raise _Refused(bad_request("Invalid collection id")) + return value + + +async def _body(request: web.Request, allowed: tuple) -> dict: + data, error = await _json_body_or_error(request) + if error is not None: + raise _Refused(error) + if frozenset(data) not in allowed: + names = " or ".join(", ".join(sorted(fields)) for fields in allowed) + raise _Refused(bad_request(f"Request must contain exactly {names}")) + return data + + +class CollectionRoutesMixin: + """Mixin providing the collection routes (#331).""" + + def _register_collection_routes(self, routes: web.RouteTableDef) -> None: + # The fixed paths go first: "/items" must never match "/{collection_id}". + routes.get(BASE)(self.list_collections) + routes.post(BASE + "/items")(self.add_collection_items) + routes.delete(BASE + "/items")(self.remove_collection_items) + routes.get(BASE + "/membership")(self.collection_membership) + routes.patch(BASE + "/{collection_id}")(self.rename_collection) + routes.delete(BASE + "/{collection_id}")(self.delete_collection) + + # -- reads ------------------------------------------------------------- + + @_collection_response(READ_FAILED) + async def list_collections(self, request: web.Request) -> web.Response: + """Every collection with its listed ``count`` and its ``missing`` rows. + + ``count`` is the members listed under every configured source (one + primary-key range read per source, matched against its cached + listing); ``missing`` the stored rows beyond it. When a scan or a read + fails both are null for every collection and the names still return. + """ + collections = await self._run_in_executor(self.db.list_collections) + counts = None + if collections: + try: + counts = await self._collection_counts() + except Exception: + self.logger.warning( + "Collection counts could not be read", exc_info=True + ) + rows = [] + 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) + rows.append( + { + **collection, + "count": count, + "missing": None if count is None else missing - 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.""" + 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(): + 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") + 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, + ) + ids = await self._run_in_executor(self.db.file_collections, root_key, key) + return web.json_response({"success": True, "collection_ids": ids}) + + # -- writes ------------------------------------------------------------ + + @_collection_response(CHANGE_FAILED) + async def add_collection_items(self, request: web.Request) -> web.Response: + """POST ``{collection_id | name, items}``: add the listed files. + + Only files the source's listing holds are added; the rest count in + ``not_found``. A new name with nothing to add creates nothing (404 + ``files-not-found``). + """ + data = await _body(request, _ADD_TARGETS) + target = {} + if "name" in data: + target["name"] = normalize_collection_name(data["name"]) + else: + target["collection_id"] = _body_id(data["collection_id"]) + roots = list(self._get_all_output_dirs()) + keys = self._collection_items(data["items"], roots) + found, not_found = [], 0 + listings = {} + for index, root_key, key in keys: + if index not in listings: + root = roots[index] + files = await self._get_gallery_files(root) + listings[index] = await self._run_in_executor( + self._favorite_listing, root, files + ) + if key in listings[index]: + found.append((root_key, key)) + else: + not_found += 1 + result = await self._run_in_executor( + lambda: self.db.add_to_collection(items=found, **target) + ) + return web.json_response({"success": True, **result, "not_found": not_found}) + + @_collection_response(CHANGE_FAILED) + async def remove_collection_items(self, request: web.Request) -> web.Response: + """DELETE ``{collection_id, items}``: listed or not, every named file.""" + 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) + result = await self._run_in_executor( + self.db.remove_from_collection, + collection_id, + [(root_key, key) for _index, root_key, key in keys], + ) + return web.json_response({"success": True, **result}) + + @_collection_response(CHANGE_FAILED) + async def rename_collection(self, request: web.Request) -> web.Response: + """PATCH ``{name}``: a case variant of its own name is allowed.""" + collection_id = _path_id(request) + data = await _body(request, (_RENAME_FIELDS,)) + name = normalize_collection_name(data["name"]) + collection = await self._run_in_executor( + self.db.rename_collection, collection_id, name + ) + return web.json_response({"success": True, "collection": collection}) + + @_collection_response(CHANGE_FAILED) + async def delete_collection(self, request: web.Request) -> web.Response: + """DELETE: the collection and its members; files are never touched.""" + collection_id = _path_id(request) + await self._run_in_executor(self.db.delete_collection, collection_id) + return web.json_response({"success": True, "deleted": True}) + + # -- request steps ----------------------------------------------------- + + def _collection_items(self, items: object, roots: list) -> list[tuple]: + """``(index, root key, relative key)`` per distinct file, in order. + + 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 = {} + for item in items: + index = item.get("root_index") if isinstance(item, dict) else None + root_key, key = self._collection_key(item, roots) + 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.""" + 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) + ): + 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. + raise _Refused(_forbidden("Path is too long to add to a collection")) + return root_key, key diff --git a/py/api/image_tag_jobs.py b/py/api/image_tag_jobs.py index 1d611745..66f148f5 100644 --- a/py/api/image_tag_jobs.py +++ b/py/api/image_tag_jobs.py @@ -58,7 +58,9 @@ from .annotations import ( from .autotag_routes import BATCH_CANCELLED, MODEL_BUSY_ERROR, model_busy_response from .images import ( MAX_LISTING_SELECTION, + CollectionMissing, ListingSourceChanged, + collection_not_found, _current_fingerprint, _listing_posix, _public_path, @@ -685,6 +687,8 @@ class ImageTagJobsMixin: selection = await self.resolve_listing_filter(spec.filter, cap) except ListingSourceChanged: return None, _coded("source-changed") + except CollectionMissing: + return None, collection_not_found() except ValueError as exc: return None, bad_request(str(exc)) except Exception: diff --git a/py/api/images.py b/py/api/images.py index 439a2800..10e70070 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -18,7 +18,7 @@ import urllib.parse from bisect import bisect_left, bisect_right, insort from collections.abc import Mapping from pathlib import Path -from typing import Callable, NamedTuple +from typing import Callable, Iterable, NamedTuple from aiohttp import web from PIL import Image, UnidentifiedImageError @@ -58,6 +58,7 @@ from .inspection_service import ( ) try: + from ...database.image_collections import CollectionMissing from ...database.image_favorites import MAX_KEY_LENGTH from ...database.library import display_title from ...database.models import normalize_image_path @@ -78,6 +79,7 @@ try: reveal_in_file_manager, ) except ImportError: + from database.image_collections import CollectionMissing from database.image_favorites import MAX_KEY_LENGTH from database.library import display_title from database.models import normalize_image_path @@ -104,6 +106,21 @@ IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS _CARD_RECORDS = ("prompt", "workflow") _PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" _LISTING_FAILED = "The image list could not be filtered. Check server logs for details." +_COLLECTION_ID = re.compile(r"[1-9][0-9]{0,8}") +COLLECTION_NOT_FOUND = "That collection no longer exists." + + +def collection_not_found() -> web.Response: + """The 404 a write or a listing answers for an unknown collection (#331).""" + return web.json_response( + { + "success": False, + "code": "collection-not-found", + "error": COLLECTION_NOT_FOUND, + }, + status=404, + headers={"Cache-Control": "no-store"}, + ) def _card_header(path: Path) -> tuple: @@ -1056,8 +1073,9 @@ def _workspace_filters(query: Mapping) -> dict: the linked positive's text, title and tag names. ``category`` is exact, ``tags`` is a comma list that must all be present, ``sort=rating`` ranks by the linked positive, and ``favorites=1`` keeps only favorites (the - Favorites page, #249 D51). Raises ValueError for any unknown or oversized - value. + Favorites page, #249 D51). ``collection=`` (1 to 9 digits) keeps only + that collection's members (#331); it wins over ``favorites`` (a page never + sets both). Raises ValueError for any unknown or oversized value. """ text = query.get("q", "").strip() folded = text.casefold() @@ -1071,8 +1089,10 @@ def _workspace_filters(query: Mapping) -> dict: {tag.strip() for tag in query.get("tags", "").split(",") if tag.strip()} ) favorites = query.get("favorites", "0") + collection = query.get("collection", "") if ( favorites not in ("0", "1") + or (collection != "" and _COLLECTION_ID.fullmatch(collection) is None) or len(text) > 500 or len(folded) > 500 or media not in ("all", "image", "video") @@ -1095,7 +1115,8 @@ def _workspace_filters(query: Mapping) -> dict: "path": exact.casefold(), "category": category, "tags": tags, - "favorites": favorites == "1", + "favorites": favorites == "1" and not collection, + "collection": int(collection) if collection else None, } @@ -1113,6 +1134,8 @@ LISTING_FILTER_KEYS = frozenset( "media", "sort", "favorites", + # A collection page's filter (#331): bulk actions act on its members. + "collection", } ) MAX_LISTING_FILTER_VALUE = 2048 @@ -1169,7 +1192,7 @@ def _admitted_entries( entries: list, filters: dict, links: dict | None, - favorites: frozenset | None = None, + members: frozenset | None = None, image_tags: dict | None = None, ): """Yield ``(position, entry, link match)`` for each entry the filters admit. @@ -1177,11 +1200,11 @@ def _admitted_entries( *links* maps normalized paths to (hit, rating, positive id) for every unambiguously linked image (see PromptDatabase.get_image_link_matches); a category or tag filter admits only those, and ``q`` admits a filename match - or a linked-positive hit. *favorites*, when given, admits only those listed - paths (``favorites=1``). *image_tags* maps normalized paths to their - _ImageTagMatch: a tag filter also admits a file whose current image tags - (with its linked positive's tags) carry every tag, and ``q`` also admits an - image tag name hit. + or a linked-positive hit. *members*, when given, admits only those listed + paths (the favorites of ``favorites=1`` or a ``collection``'s members). + *image_tags* maps normalized paths to their _ImageTagMatch: a tag filter + also admits a file whose current image tags (with its linked positive's + tags) carry every tag, and ``q`` also admits an image tag name hit. """ query, media, folder = filters["query"], filters["media"], filters["folder"] members_only = filters["category"] is not None or bool(filters["tags"]) @@ -1197,7 +1220,7 @@ def _admitted_entries( continue if media != "all" and (media == "video") != entry.video: continue - if favorites is not None and entry.path not in favorites: + if members is not None and entry.path not in members: continue match = None if links is None else links.get(entry.normalized) image = None if image_tags is None else image_tags.get(entry.normalized) @@ -1219,7 +1242,7 @@ def _filtered_page( size: int, filters: dict, links: dict | None, - favorites: frozenset | None = None, + members: frozenset | None = None, image_tags: dict | None = None, ) -> tuple[int, list, int | None]: """Count the *entries* the filters admit and take *size* from *start* on. @@ -1230,7 +1253,7 @@ def _filtered_page( """ total, page, positives = 0, [], set() for position, entry, match in _admitted_entries( - entries, filters, links, favorites, image_tags + entries, filters, links, members, image_tags ): total += 1 if match is not None: @@ -1934,6 +1957,7 @@ class ImageRoutesMixin: filters["category"], filters["tags"], filters["favorites"], + filters["collection"], ] ).encode() ).hexdigest() @@ -1944,17 +1968,19 @@ class ImageRoutesMixin: except ValueError as exc: return bad_request(str(exc)) - listings = [ - (index, root, await self._get_gallery_files(root)) - for index, root in enumerate(output_dirs) - if root_index is None or index == root_index - ] try: + if filters["collection"] is not None and not await self._run_in_executor( + self.db.collection_exists, filters["collection"] + ): + raise CollectionMissing("No such collection") + listings = await self._workspace_listings(output_dirs, root_index) page, total, has_more, next_cursor, positive_count = ( await self._run_in_executor( self._select_workspace_page, listings, filters, after, limit, scope ) ) + except CollectionMissing: + return collection_not_found() except Exception: # A filter that needs the prompt links or the image tags cannot be # served without them. Never echo the exception: it can carry SQL @@ -2215,9 +2241,11 @@ class ImageRoutesMixin: ) -> tuple: """The ordered entries and what admits them (blocking; run in executor). - Returns ``(entries, key, links, image_tags, favorites)``: the listing in - the filters' order, its keyset key, and the link, image-tag and favorite - maps _admitted_entries reads. One SELECT (get_image_link_matches) reads + Returns ``(entries, key, links, image_tags, members)``: the listing in + the filters' order, its keyset key, and the link, image-tag and member + maps _admitted_entries reads (members: a collection's, else the + favorites, else None). A collection deleted since the route checked it + raises CollectionMissing. One SELECT (get_image_link_matches) reads every linked positive when *read_links*. A failed read raises when the filters need it (category, tags, the rating sort) or *strict*; else the listing stays usable without links (a search then matches filenames @@ -2290,10 +2318,13 @@ class ImageRoutesMixin: def key(entry): return _workspace_key(entry, order) - favorites = ( - self._workspace_favorite_paths(listings) if filters["favorites"] else None + return ( + entries, + key, + links, + image_tags, + self._workspace_members(listings, filters), ) - return entries, key, links, image_tags, favorites def _select_workspace_page( self, @@ -2322,7 +2353,7 @@ class ImageRoutesMixin: or bool(filters["tags"]) or order == "rating" ) - entries, key, links, image_tags, favorites = self._workspace_admission( + entries, key, links, image_tags, members = self._workspace_admission( listings, filters, read_links=needed or after is None ) start = 0 @@ -2333,7 +2364,7 @@ class ImageRoutesMixin: else: start = bisect_right(entries, after, key=key) total, page, positive_count = _filtered_page( - entries, start, limit + 1, filters, links, favorites, image_tags + entries, start, limit + 1, filters, links, members, image_tags ) if not needed and after is not None: positive_count = _COUNT_OMITTED @@ -2355,12 +2386,12 @@ class ImageRoutesMixin: The links are read strictly: the positives must be the listing's own. """ - entries, _key, links, image_tags, favorites = self._workspace_admission( + entries, _key, links, image_tags, members = self._workspace_admission( listings, filters, read_links=True, strict=True ) targets, positives, unlinked, scopes = [], {}, 0, {} for _position, entry, match in _admitted_entries( - entries, filters, links, favorites, image_tags + entries, filters, links, members, image_tags ): if len(targets) == cap: return ListingSelection(targets, list(positives), unlinked, True) @@ -2382,9 +2413,11 @@ class ImageRoutesMixin: ) -> ListingSelection: """The images a parse_listing_filter result admits right now (#392). - Raises ValueError for a root outside the current roots, and + Raises ValueError for a root outside the current roots, ListingSourceChanged when the filter's ``root_scope`` no longer names - its root (the sources were renumbered since the listing). + its root (the sources were renumbered since the listing), and + CollectionMissing when its ``collection`` no longer exists (checked + before any scan and again after the members are read). """ output_dirs = list(self._get_all_output_dirs()) root_index = _workspace_root(params, output_dirs) @@ -2394,11 +2427,11 @@ class ImageRoutesMixin: raise ListingSourceChanged() raise ValueError("Invalid gallery root scope") filters = _workspace_filters(params) - listings = [ - (index, root, await self._get_gallery_files(root)) - for index, root in enumerate(output_dirs) - if root_index is None or index == root_index - ] + if filters["collection"] is not None and not await self._run_in_executor( + self.db.collection_exists, filters["collection"] + ): + raise CollectionMissing("No such collection") + listings = await self._workspace_listings(output_dirs, root_index) return await self._run_in_executor( self._select_workspace_matches, listings, filters, cap ) @@ -3583,25 +3616,58 @@ class ImageRoutesMixin: self._favorite_listing_cache[cache_key] = (files, listing) return listing - def _workspace_favorite_paths(self, listings: list[tuple]) -> frozenset: - """The listed paths of every favorite in the selected roots (blocking). + def _workspace_members(self, listings: list[tuple], filters: dict): + """The member set the filters ask for, or None (blocking).""" + collection = filters.get("collection") + if collection is not None: + members = self._workspace_member_paths( + listings, + lambda root_key: self.db.collection_members(collection, root_key), + ) + # Checked after the reads: a collection deleted before or while + # they ran is a 404, never an empty page passing for it. + if not self.db.collection_exists(collection): + raise CollectionMissing("No such collection") + return members + if filters["favorites"]: + return self._workspace_favorite_paths(listings) + return None - One primary-key read per root (get_root_favorites), matched against the - root's cached ``{relative key: listing spelling}``, so the work follows - the favorites rather than the files, and a key folds case exactly as - the toggle stored it (Windows). A favorite whose file is gone matches - nothing. A failed read raises: an unfiltered page must never pass for - the favorites. + def _workspace_favorite_paths(self, listings: list[tuple]) -> frozenset: + """The listed paths of every favorite in the selected roots (blocking).""" + return self._workspace_member_paths(listings, self.db.get_root_favorites) + + def _workspace_member_paths( + self, listings: list[tuple], keys_for_source: Callable[[str], Iterable[str]] + ) -> frozenset: + """The listed paths of the stored keys of each selected root (blocking). + + One read per root (``keys_for_source(root key)``: the favorites' or a + collection's primary-key range), matched against the root's cached + ``{relative key: listing spelling}``, so the work follows the members + rather than the files, and a key folds case exactly as it was stored + (Windows). A member whose file is gone matches nothing. A failed read + raises: an unfiltered page must never pass for the members. """ paths = set() for _index, root, files in listings: listing = self._favorite_listing(root, files) - for key in self.db.get_root_favorites(_favorite_root_key(root)): + for key in keys_for_source(_favorite_root_key(root)): listed = listing.get(key) if listed is not None: paths.add(root / listed) return frozenset(paths) + async def _workspace_listings( + self, output_dirs: list[Path], root_index: int | None + ) -> list[tuple]: + """``(index, root, scan)`` of every selected root (cached scans).""" + return [ + (index, root, await self._get_gallery_files(root)) + for index, root in enumerate(output_dirs) + if root_index is None or index == root_index + ] + def _workspace_favorites(self, page: list[tuple]) -> set[tuple[int, str]]: """``(root index, listing path)`` of the favorites among *page* entries.""" by_root: dict[tuple[int, Path], dict[str, str]] = {} diff --git a/py/api/media_target.py b/py/api/media_target.py new file mode 100644 index 00000000..ee1fcadf --- /dev/null +++ b/py/api/media_target.py @@ -0,0 +1,54 @@ +"""The listing identity of one media file, as a request names it. + +Shared by the image-annotation routes and the collection routes (#331): a +target is ``{root_index, root_scope, path}``, checked here without touching +the filesystem or the database. The scope is compared with the current roots +by the caller. +""" + +import os +import re + +from .images import _forbidden, _has_traversal +from .prompts import bad_request + +try: + from ...utils.media_capture import CAPTURE_EXTENSIONS +except ImportError: + from utils.media_capture import CAPTURE_EXTENSIONS + +MAX_TARGET_PATH = 1024 +TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"}) +_HEX64 = re.compile(r"[0-9a-f]{64}") +TAGGING_MEDIA_MESSAGE = "Only media files can be tagged" + + +def media_target(data: object, *, media_message: str = TAGGING_MEDIA_MESSAGE): + """``((index, scope, relative), None)`` for a target, else ``(None, error)``. + + Only the listing identity is accepted: an int root index, the root's + 64-hex scope and a bounded root-relative POSIX media path, never an + absolute, drive-qualified, backslash or ``..`` path (a 400 for a malformed + shape, a 403 for an unsafe or non-media path, carrying *media_message*). + """ + if not isinstance(data, dict) or set(data) != TARGET_FIELDS: + return None, bad_request("target must name root_index, root_scope and path") + index, scope, relative = data["root_index"], data["root_scope"], data["path"] + if type(index) is not int or index < 0: + return None, bad_request("Invalid gallery root") + if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None: + return None, bad_request("Invalid gallery root scope") + if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH: + return None, bad_request("path must be a relative media path") + parts = relative.split("/") + if ( + "\0" in relative + or "\\" in relative + or relative.startswith("/") + or _has_traversal(relative) + or any(part in ("", ".", "..") or ":" in part for part in parts) + ): + return None, _forbidden() + if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS: + return None, _forbidden(media_message) + return (index, scope, relative), None diff --git a/py/api/prompts.py b/py/api/prompts.py index e48296ea..8f5a1554 100644 --- a/py/api/prompts.py +++ b/py/api/prompts.py @@ -1227,7 +1227,9 @@ class PromptRoutesMixin: """ from .images import ( MAX_LISTING_SELECTION, + CollectionMissing, ListingSourceChanged, + collection_not_found, parse_listing_filter, ) @@ -1252,6 +1254,8 @@ class PromptRoutesMixin: }, status=409, ) + except CollectionMissing: + return collection_not_found() except ValueError as exc: return refused(str(exc)) if selection.truncated: diff --git a/tests/browser/server.py b/tests/browser/server.py index 71ea5669..95deff78 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -262,6 +262,8 @@ def create_app(root: Path) -> web.Application: with target.db.model.get_connection() as conn: for table in ( "saved_views", + "image_collection_items", + "image_collections", "image_favorites", "media_review_decisions", "media_review_items", diff --git a/tests/test_api_collections.py b/tests/test_api_collections.py new file mode 100644 index 00000000..c4dcdd86 --- /dev/null +++ b/tests/test_api_collections.py @@ -0,0 +1,494 @@ +"""Collections (#331): the routes under /prompt_manager/collections.""" + +import json +import ntpath +import sqlite3 +from unittest.mock import patch + +try: + from test_api_images import ImageAPITestCase, make_png + from open_handles import assert_closed +except ImportError: + from tests.test_api_images import ImageAPITestCase, make_png + from tests.open_handles import assert_closed + +BASE = "/prompt_manager/collections" +ITEMS = BASE + "/items" +MEMBERSHIP = BASE + "/membership" +SOURCES_CHANGED = "The sources changed. Refresh the images and try again." + + +class CollectionRoutesTestCase(ImageAPITestCase): + async def tearDownAsync(self): + self.api.db.close_all() + assert_closed(self, self._temp_db.name, self._output_tmp.name) + await super().tearDownAsync() + + def make(self, *names, root=None): + for name in names: + make_png((root or self.output_dir) / name) + + def scope(self, index=0): + roots = list(self.api._get_all_output_dirs()) + return self.api._gallery_root_scope(roots[index], index) + + def item(self, path, index=0, scope=None): + return { + "root_index": index, + "root_scope": scope or self.scope(index), + "path": path, + } + + def two_sources(self): + second = self.comfy_dir / "second" + second.mkdir() + self.api._get_all_output_dirs = lambda: [self.output_dir, second] + return second + + async def call(self, method, url, body=None, status=200, **params): + response = await self.client.request(method, url, json=body, params=params) + text = await response.text() + self.assertEqual(response.status, status, text) + data = json.loads(text) + self.assertNotIn(self.comfy_dir.as_posix(), text) + self.assertNotIn(str(self.comfy_dir).replace("\\", "\\\\"), text) + return data + + async def add(self, status=200, **body): + return await self.call("POST", ITEMS, body, status) + + async def listed(self): + return (await self.call("GET", BASE))["collections"] + + +class TestCollectionRoutes(CollectionRoutesTestCase): + async def test_each_route_succeeds(self): + self.make("a.png", "b.png", "c.png") + self.assertEqual( + await self.call("GET", BASE), {"success": True, "collections": []} + ) + made = await self.add( + name=" LoRa testing ", items=[self.item("a.png"), self.item("b.png")] + ) + grids = made["collection"] + self.assertEqual(grids["name"], "LoRa testing") + self.assertEqual( + made, + { + "success": True, + "collection": grids, + "created": True, + "added": 2, + "already": 0, + "not_found": 0, + }, + ) + again = await self.add( + collection_id=grids["id"], items=[self.item("b.png"), self.item("c.png")] + ) + self.assertEqual( + (again["created"], again["added"], again["already"]), (False, 1, 1) + ) + self.assertEqual( + await self.listed(), + [{"id": grids["id"], "name": "LoRa testing", "count": 3, "missing": 0}], + ) + membership = await self.call( + "GET", MEMBERSHIP, root_index="0", root_scope=self.scope(), path="a.png" + ) + self.assertEqual(membership, {"success": True, "collection_ids": [grids["id"]]}) + removed = await self.call( + "DELETE", + ITEMS, + {"collection_id": grids["id"], "items": [self.item("a.png")]}, + ) + self.assertEqual(removed, {"success": True, "collection": grids, "removed": 1}) + renamed = await self.call( + "PATCH", f"{BASE}/{grids['id']}", {"name": "LoRA grids"} + ) + self.assertEqual( + renamed, + {"success": True, "collection": {"id": grids["id"], "name": "LoRA grids"}}, + ) + deleted = await self.call("DELETE", f"{BASE}/{grids['id']}") + self.assertEqual(deleted, {"success": True, "deleted": True}) + self.assertEqual(await self.listed(), []) + # Files are never touched. + self.assertTrue((self.output_dir / "a.png").is_file()) + + async def test_names_are_unique_ignoring_case(self): + self.make("a.png", "b.png") + first = (await self.add(name="LoRa testing", items=[self.item("a.png")]))[ + "collection" + ] + second = await self.add(name="lora TESTING", items=[self.item("b.png")]) + self.assertEqual(second["collection"], first) + self.assertFalse(second["created"]) + self.assertEqual(len(await self.listed()), 1) + + async def test_bad_requests_are_400(self): + self.make("a.png") + made = await self.add(name="Set", items=[self.item("a.png")]) + cid = made["collection"]["id"] + good = [self.item("a.png")] + for body in ( + {"items": good}, + {"name": "x", "collection_id": cid, "items": good}, + {"name": "x"}, + {"name": "x", "items": good, "extra": 1}, + {"collection_id": "1", "items": good}, + {"collection_id": True, "items": good}, + {"collection_id": 0, "items": good}, + {"collection_id": 10**10, "items": good}, + {"name": "", "items": good}, + {"name": "x" * 101, "items": good}, + {"name": "bell\x07", "items": good}, + {"name": 5, "items": good}, + {"name": "x", "items": []}, + {"name": "x", "items": "a.png"}, + {"name": "x", "items": [self.item("a.png")] * 501}, + {"name": "x", "items": [{"root_index": 0, "path": "a.png"}]}, + {"name": "x", "items": [{**self.item("a.png"), "root_index": "0"}]}, + {"name": "x", "items": [{**self.item("a.png"), "root_scope": "xyz"}]}, + ): + with self.subTest(body=body): + data = await self.add(400, **body) + self.assertFalse(data["success"]) + for body in ({"items": good}, {"collection_id": cid}): + with self.subTest(remove=body): + await self.call("DELETE", ITEMS, body, 400) + for raw in ("abc", "0", "1234567890", "-1", "1.0"): + with self.subTest(path_id=raw): + await self.call("PATCH", f"{BASE}/{raw}", {"name": "x"}, 400) + await self.call("DELETE", f"{BASE}/{raw}", None, 400) + await self.call("PATCH", f"{BASE}/{cid}", {"name": ""}, 400) + await self.call("PATCH", f"{BASE}/{cid}", {"label": "x"}, 400) + for params in ( + {"root_index": "x", "root_scope": self.scope(), "path": "a.png"}, + {"root_index": "0", "path": "a.png"}, + {"root_index": "0", "root_scope": "nope", "path": "a.png"}, + ): + with self.subTest(membership=params): + await self.call("GET", MEMBERSHIP, status=400, **params) + self.assertEqual((await self.listed())[0]["count"], 1) + + async def test_unsafe_paths_are_403(self): + self.make("a.png") + for path in ( + "../a.png", + "sub/../../a.png", + "C:/a.png", + "a\\b.png", + "/a.png", + "notes.txt", + ): + with self.subTest(path=path): + await self.add(403, name="Set", items=[self.item(path)]) + await self.call( + "GET", + MEMBERSHIP, + status=403, + root_index="0", + root_scope=self.scope(), + path=path, + ) + self.assertEqual(await self.listed(), []) + + async def test_a_body_over_the_cap_is_413(self): + response = await self.client.post( + ITEMS, data=b'{"name":"x","items":[' + b" " * (1024 * 1024 + 1) + b"]}" + ) + self.assertEqual(response.status, 413) + + async def test_missing_collection_and_files_are_404(self): + self.make("a.png") + data = await self.add(404, collection_id=999, items=[self.item("a.png")]) + self.assertEqual(data["code"], "collection-not-found") + data = await self.call( + "DELETE", ITEMS, {"collection_id": 999, "items": [self.item("a.png")]}, 404 + ) + self.assertEqual(data["code"], "collection-not-found") + data = await self.call("PATCH", f"{BASE}/999", {"name": "x"}, 404) + self.assertEqual(data["code"], "collection-not-found") + data = await self.call("DELETE", f"{BASE}/999", None, 404) + self.assertEqual(data["code"], "collection-not-found") + data = await self.add(404, name="New", items=[self.item("never.png")]) + self.assertEqual(data["code"], "files-not-found") + self.assertEqual(await self.listed(), []) + + async def test_unlisted_files_count_as_not_found(self): + self.make("a.png") + made = await self.add( + name="Set", items=[self.item("a.png"), self.item("later.png")] + ) + self.assertEqual((made["added"], made["not_found"]), (1, 1)) + again = await self.add( + collection_id=made["collection"]["id"], + items=[self.item("a.png"), self.item("later.png")], + ) + self.assertEqual( + (again["added"], again["already"], again["not_found"]), (0, 1, 1) + ) + + async def test_name_taken_is_409(self): + self.make("a.png") + first = (await self.add(name="First", items=[self.item("a.png")]))["collection"] + second = (await self.add(name="Second", items=[self.item("a.png")]))[ + "collection" + ] + data = await self.call( + "PATCH", f"{BASE}/{second['id']}", {"name": "FIRST"}, 409 + ) + self.assertEqual( + (data["code"], data["collection_id"]), ("name-taken", first["id"]) + ) + renamed = await self.call("PATCH", f"{BASE}/{second['id']}", {"name": "SECOND"}) + self.assertEqual(renamed["collection"]["name"], "SECOND") + + async def test_stale_scope_is_409_after_reordering_sources(self): + second = self.two_sources() + self.make("a.png") + self.make("b.png", root=second) + stale = self.item("b.png", 1) + made = await self.add(name="Set", items=[self.item("a.png")]) + cid = made["collection"]["id"] + self.api._get_all_output_dirs = lambda: [second, self.output_dir] + for method, body in ( + ("POST", {"collection_id": cid, "items": [self.item("b.png", 0), stale]}), + ("DELETE", {"collection_id": cid, "items": [stale]}), + ): + with self.subTest(method=method): + data = await self.call(method, ITEMS, body, 409) + self.assertEqual(data["code"], "source-changed") + self.assertEqual(data["error"], SOURCES_CHANGED) + data = await self.call( + "GET", + MEMBERSHIP, + status=409, + root_index="1", + root_scope=stale["root_scope"], + path="b.png", + ) + self.assertEqual(data["code"], "source-changed") + # Nothing was written by the refused add; the member follows its folder. + listed = await self.listed() + self.assertEqual((listed[0]["count"], listed[0]["missing"]), (1, 0)) + membership = await self.call( + "GET", MEMBERSHIP, root_index="1", root_scope=self.scope(1), path="a.png" + ) + self.assertEqual(membership["collection_ids"], [cid]) + + async def test_counts_exclude_a_missing_member(self): + second = self.two_sources() + self.make("a.png", "b.png") + self.make("c.png", root=second) + made = await self.add( + name="Grids", + items=[self.item("a.png"), self.item("b.png"), self.item("c.png", 1)], + ) + other = await self.add(name="Other", items=[self.item("c.png", 1)]) + (self.output_dir / "b.png").unlink() + self.api.invalidate_gallery_cache() + listed = await self.listed() + self.assertEqual( + listed, + [ + { + "id": made["collection"]["id"], + "name": "Grids", + "count": 2, + "missing": 1, + }, + { + "id": other["collection"]["id"], + "name": "Other", + "count": 1, + "missing": 0, + }, + ], + ) + # The member is removable although it is no longer listed. + removed = await self.call( + "DELETE", + ITEMS, + {"collection_id": made["collection"]["id"], "items": [self.item("b.png")]}, + ) + self.assertEqual(removed["removed"], 1) + + async def test_counts_are_null_when_a_source_scan_fails(self): + self.make("a.png") + await self.add(name="Set", items=[self.item("a.png")]) + + async def broken(_root): + raise OSError("scan failed at /secret/place") + + with patch.object(self.api, "_get_gallery_files", side_effect=broken): + data = await self.call("GET", BASE) + self.assertEqual( + data["collections"][0] | {"id": 0}, + {"id": 0, "name": "Set", "count": None, "missing": None}, + ) + self.assertNotIn("secret", json.dumps(data)) + + async def test_windows_case_spellings_are_one_member(self): + self.make("Sub/A.png") + with patch("py.api.images._path_flavor", ntpath): + made = await self.add( + name="Set", items=[self.item("Sub/A.png"), self.item("SUB/a.PNG")] + ) + self.assertEqual((made["added"], made["already"]), (1, 0)) + (self.output_dir / "Sub" / "A.png").rename( + self.output_dir / "Sub" / "a.png" + ) + (self.output_dir / "Sub").rename(self.output_dir / "sub") + self.api.invalidate_gallery_cache() + listed = await self.listed() + self.assertEqual((listed[0]["count"], listed[0]["missing"]), (1, 0)) + membership = await self.call( + "GET", + MEMBERSHIP, + root_index="0", + root_scope=self.scope(), + path="sub/a.png", + ) + self.assertEqual(membership["collection_ids"], [made["collection"]["id"]]) + response = await self.client.get( + "/prompt_manager/images/output", + params={ + "mode": "workspace", + "collection": str(made["collection"]["id"]), + }, + ) + page = await response.json() + self.assertEqual([i["path"] for i in page["images"]], ["sub/a.png"]) + + async def test_posix_case_distinct_files_stay_two_members(self): + self.make("A.png", "a.png") + if len(list(self.output_dir.iterdir())) != 2: + self.skipTest("case-insensitive filesystem") + made = await self.add( + name="Set", items=[self.item("A.png"), self.item("a.png")] + ) + self.assertEqual(made["added"], 2) + + async def test_a_500_never_carries_the_exception_text(self): + self.make("a.png") + made = await self.add(name="Set", items=[self.item("a.png")]) + cid = made["collection"]["id"] + + def broken(*_args, **_kwargs): + raise sqlite3.OperationalError("disk I/O error at /secret/prompts.db") + + cases = ( + ( + "list_collections", + "GET", + BASE, + None, + {}, + "Collections could not be read.", + ), + ( + "file_collections", + "GET", + MEMBERSHIP, + None, + {"root_index": "0", "root_scope": self.scope(), "path": "a.png"}, + "Collections could not be read.", + ), + ( + "add_to_collection", + "POST", + ITEMS, + {"collection_id": cid, "items": [self.item("a.png")]}, + {}, + "The collection could not be changed.", + ), + ( + "remove_from_collection", + "DELETE", + ITEMS, + {"collection_id": cid, "items": [self.item("a.png")]}, + {}, + "The collection could not be changed.", + ), + ( + "rename_collection", + "PATCH", + f"{BASE}/{cid}", + {"name": "x"}, + {}, + "The collection could not be changed.", + ), + ( + "delete_collection", + "DELETE", + f"{BASE}/{cid}", + None, + {}, + "The collection could not be changed.", + ), + ) + for method_name, method, url, body, params, message in cases: + with self.subTest(method=method_name): + with patch.object(self.api.db, method_name, side_effect=broken): + data = await self.call(method, url, body, 500, **params) + self.assertEqual(data, {"success": False, "error": message}) + + +class TestCollectionListingFilter(CollectionRoutesTestCase): + """#392's "all N matching" filter carries ``collection`` like the listing.""" + + async def test_a_filter_with_a_collection_resolves_only_members(self): + second = self.two_sources() + self.make("a.png", "b.png") + self.make("c.png", root=second) + made = await self.add( + name="Set", items=[self.item("a.png"), self.item("c.png", 1)] + ) + cid = str(made["collection"]["id"]) + selection = await self.api.resolve_listing_filter({"collection": cid}) + self.assertEqual( + sorted((index, path) for index, _scope, path in selection.targets), + [(0, "a.png"), (1, "c.png")], + ) + selection = await self.api.resolve_listing_filter( + {"collection": cid, "root": "1", "root_scope": self.scope(1)} + ) + self.assertEqual([t[2] for t in selection.targets], ["c.png"]) + for bad in ("0", "x", "1234567890"): + with self.subTest(bad=bad): + with self.assertRaises(ValueError): + from py.api.images import parse_listing_filter + + parse_listing_filter({"collection": bad}) + + async def test_bulk_category_acts_on_collection_members_only(self): + self.make("a.png", "b.png") + inside = self._save_prompt("Inside", category="Old") + outside = self._save_prompt("Outside", category="Old") + self.api.db.link_image_to_prompt(inside, str(self.output_dir / "a.png")) + self.api.db.link_image_to_prompt(outside, str(self.output_dir / "b.png")) + made = await self.add(name="Set", items=[self.item("a.png")]) + cid = str(made["collection"]["id"]) + data = await self.call( + "POST", + "/prompt_manager/bulk/category", + {"images_filter": {"collection": cid}, "category": "New"}, + ) + self.assertEqual((data["updated_count"], data["prompt_count"]), (1, 1)) + self.assertEqual(self.api.db.get_prompt_by_id(inside)["category"], "New") + self.assertEqual(self.api.db.get_prompt_by_id(outside)["category"], "Old") + + async def test_a_deleted_collection_in_a_bulk_filter_is_a_404(self): + self.make("a.png") + data = await self.call( + "POST", + "/prompt_manager/bulk/category", + {"images_filter": {"collection": "77"}, "category": "New"}, + status=404, + ) + self.assertEqual(data["code"], "collection-not-found") + with self.assertRaises(Exception) as caught: + await self.api.resolve_listing_filter({"collection": "77"}) + self.assertEqual(type(caught.exception).__name__, "CollectionMissing") diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index f30f3947..26211a5b 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -2494,3 +2494,171 @@ class TestListingFilterSelection(ImageAPITestCase): self.assertEqual(response.status, 409) self.assertEqual((await response.json())["code"], "source-changed") + + +class TestWorkspaceCollectionFilter(ImageAPITestCase): + """#331: ``collection=`` keeps only that collection's listed members.""" + + media = TestWorkspaceImages.media + get = TestWorkspaceImages.get + _names = TestWorkspaceImages._names + + async def tearDownAsync(self): + self.api.db.close_all() + assert_closed(self, self._temp_db.name, self._output_tmp.name) + await super().tearDownAsync() + + def keys(self, root, relative): + from py.api.images import _favorite_relative_key, _favorite_root_key + + return (_favorite_root_key(root), _favorite_relative_key(relative)) + + def collect(self, name, *members): + """A collection holding ``(root, relative)`` members; its id.""" + result = self.api.db.add_to_collection( + name=name, items=[self.keys(root, rel) for root, rel in members] + ) + return result["collection"]["id"] + + def two_sources(self): + second = self.comfy_dir / "second" + second.mkdir() + self.api._get_all_output_dirs = lambda: [self.output_dir, second] + return second + + async def test_listing_keeps_only_collection_members(self): + second = self.two_sources() + self.media("a.png", 300) + self.media("b.png", 200) + self.media("sub/c.png", 100, root=second) + self.media("d.png", 50, root=second) + grids = self.collect( + "Grids", + (self.output_dir, "a.png"), + (second, "sub/c.png"), + (self.output_dir, "gone.png"), # stored, never listed + ) + self.collect("Other", (self.output_dir, "b.png")) + response, data = await self.get(collection=str(grids)) + self.assertEqual(response.status, 200, data) + self.assertEqual(self._names(data), ["a.png", "c.png"]) + self.assertEqual(data["total"], 2) + _, data = await self.get(collection=str(grids), root="1") + self.assertEqual(self._names(data), ["c.png"]) + # Without the filter the listing is unchanged. + _, data = await self.get() + self.assertEqual(data["total"], 4) + + async def test_collection_filter_combines_with_the_other_filters(self): + night = self._save_prompt("Neon zebra", category="Night", tags=["a"]) + day = self._save_prompt("Morning field", category="Day", tags=["a"]) + files = { + name: self.media(name, 100 + index) + for index, name in enumerate( + ("night.png", "day.png", "zebra_plain.png", "clip.mp4", "out.png") + ) + } + self.api.db.link_image_to_prompt(night, str(files["night.png"])) + self.api.db.link_image_to_prompt(day, str(files["day.png"])) + members = [(self.output_dir, name) for name in files if name != "out.png"] + grids = str(self.collect("Grids", *members)) + cases = ( + ({}, ["clip.mp4", "day.png", "night.png", "zebra_plain.png"]), + ({"q": "zebra"}, ["night.png", "zebra_plain.png"]), + ({"category": "Night"}, ["night.png"]), + ({"tags": "a"}, ["day.png", "night.png"]), + ({"media": "video"}, ["clip.mp4"]), + ({"media": "image", "q": "day"}, ["day.png"]), + ) + for params, expected in cases: + with self.subTest(**params): + response, data = await self.get(collection=grids, sort="name", **params) + self.assertEqual(response.status, 200, data) + self.assertEqual(self._names(data), expected) + self.assertEqual(data["total"], len(expected)) + + async def test_invalid_collection_is_a_400(self): + for value in ("x", "0", "-1", "1.5", "1234567890", " 1"): + with self.subTest(value=value): + response, data = await self.get(collection=value) + self.assertEqual(response.status, 400, data) + + async def test_unknown_collection_is_a_404(self): + self.media("a.png") + response, data = await self.get(collection="77") + self.assertEqual(response.status, 404) + self.assertEqual(data["code"], "collection-not-found") + self.assertNotIn("images", data) + self.api._get_all_output_dirs = lambda: [] + response, data = await self.get(collection="77") + self.assertEqual((response.status, data["code"]), (404, "collection-not-found")) + + async def test_a_collection_deleted_during_the_read_is_a_404(self): + self.media("a.png") + grids = self.collect("Grids", (self.output_dir, "a.png")) + real = self.api.db.collection_members + + def deleted_meanwhile(collection_id, root_key): + self.api.db.delete_collection(collection_id) + return real(collection_id, root_key) + + with patch.object( + self.api.db, "collection_members", side_effect=deleted_meanwhile + ): + response, data = await self.get(collection=str(grids)) + self.assertEqual((response.status, data["code"]), (404, "collection-not-found")) + + async def test_a_failed_members_read_is_a_500_never_an_unfiltered_page(self): + self.media("a.png") + self.media("b.png") + grids = self.collect("Grids", (self.output_dir, "a.png")) + + def broken(*_args, **_kwargs): + raise sqlite3.OperationalError("database is locked") + + with patch.object(self.api.db, "collection_members", side_effect=broken): + response, data = await self.get(collection=str(grids)) + self.assertEqual(response.status, 500) + self.assertNotIn("images", data) + self.assertNotIn("locked", json.dumps(data)) + + async def test_reordering_sources_keeps_members(self): + second = self.two_sources() + self.media("a.png", root=second) + self.media("b.png") + grids = str(self.collect("Grids", (second, "a.png"))) + _, data = await self.get(collection=grids) + self.assertEqual( + [(i["root_index"], i["path"]) for i in data["images"]], [(1, "a.png")] + ) + self.api._get_all_output_dirs = lambda: [second, self.output_dir] + _, data = await self.get(collection=grids) + self.assertEqual( + [(i["root_index"], i["path"]) for i in data["images"]], [(0, "a.png")] + ) + + async def test_cursor_from_another_collection_is_refused(self): + for index, name in enumerate(("a.png", "b.png", "c.png")): + self.media(name, 100 + index) + everything = [(self.output_dir, name) for name in ("a.png", "b.png", "c.png")] + first = str(self.collect("First", *everything)) + second = str(self.collect("Second", *everything)) + _, page = await self.get(collection=first, limit=1) + token = page["next_cursor"] + self.assertIsInstance(token, str) + response, _ = await self.get(collection=first, limit=1, cursor=token) + self.assertEqual(response.status, 200) + for params in ({"collection": second}, {}, {"favorites": "1"}): + with self.subTest(**params): + response, _ = await self.get(cursor=token, limit=1, **params) + self.assertEqual(response.status, 400) + + async def test_collection_wins_over_favorites(self): + self.media("a.png") + self.media("b.png") + grids = str(self.collect("Grids", (self.output_dir, "a.png"))) + from py.api.images import _favorite_root_key + + self.api.db.set_image_favorite(_favorite_root_key(self.output_dir), "b.png") + _, data = await self.get(collection=grids, favorites="1") + self.assertEqual(self._names(data), ["a.png"]) diff --git a/tests/test_image_tag_jobs.py b/tests/test_image_tag_jobs.py index 0af74981..f593aaf5 100644 --- a/tests/test_image_tag_jobs.py +++ b/tests/test_image_tag_jobs.py @@ -884,6 +884,16 @@ class TestFilterStart(ImageTagJobTestCase): self.assertEqual((status, data["code"]), (400, "no-images")) self.assertEqual(self.api._jobs.list(), []) + async def test_filter_with_a_deleted_collection_is_a_404_without_a_job(self): + self.png("a.png") + + status, data = await self.start( + {"mode": "add", "filter": {"collection": "77"}, "tags": ["x"]} + ) + + self.assertEqual((status, data["code"]), (404, "collection-not-found")) + self.assertEqual(self.api._jobs.list(), []) + async def test_filter_over_the_cap_is_refused_without_a_job(self): from py.api import image_tag_jobs