feat(collections): add collection routes and the Images collection filter
This commit is contained in:
@@ -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)
|
||||
|
||||
+8
-36
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
+110
-44
@@ -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=<id>`` (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]] = {}
|
||||
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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")
|
||||
@@ -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=<id>`` 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"])
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user