diff --git a/TESTING.md b/TESTING.md index ce4bd6b2..26898c0c 100644 --- a/TESTING.md +++ b/TESTING.md @@ -4470,6 +4470,64 @@ concept fixture stars the concept's favorites (each other prompt's first image). restore. Old backups without the table restore, and backup plus restore keeps WAL-resident favorites and closes every handle (`tests/test_database_backup.py`). +Image categories (#332, store and routes): an image file may carry its own +category; without one it shows the category its link supplies. + +- The table is `image_categories(root_key, relative_path, category, updated_at)` + with primary key `(root_key, relative_path)` and the index + `idx_image_categories_category`. `migrate_image_categories` creates both empty + inside a savepoint, with no backfill. It reads only `sqlite_master` and its own + table's PRAGMAs, and writes only its two objects. `tests/test_image_categories.py` + checks this with an SQLite authorizer on a 23,000-link library, and checks that + `prompts` and `prompt_pairs` `(id, category, updated_at)` are byte-for-byte equal + afterwards. +- The **file key** is `image_category_key(sources, absolute path)` in + `py/api/images.py`: the deepest output source that contains the path (compared + lexically, never resolving the path, so a symlinked subfolder keeps its listed + spelling) and the path below it, keyed with the favorites helpers. Without + nested sources it equals a favorite's key. Every write and read uses it. A path + that no source contains has no key, so it has no own category. +- `image_categories_schema_error` checks the structure at startup (a value saved + under older rules never stops the library opening). Restore also checks the + values: text columns, keys within the favorites bound, `category` 1 to 100 + characters with no character below U+0020 and no leading or trailing space, and + `updated_at` of at most 64 bytes. Backups made before #332 restore without the + table; a migrated restore stage must have it (`require_library_state`). +- `PUT /prompt_manager/images/category` `{items, category}` answers + `{success, category, updated, unchanged, not_found}`. + `DELETE /prompt_manager/images/category` `{items}` answers `{success, cleared}`. + `GET /prompt_manager/images/categories` answers `{success, categories}`: the + sorted distinct own values whose file is in some source's listing. +- `items` are 1 to 500 Images targets `{root_index, root_scope, path}`, checked by + `py/api/media_target.py` (shared with Collections, #331). Refusals: + - 400: a malformed body, a field other than `items` and `category`, a + `root_index` outside the output sources (a LoRA folder), or a category that is + blank after stripping, over 100 characters, or holds a control character; + - 403: traversal, a drive, a backslash, a non-media name ("Only media files can + have a category") or a key over the stored bound; + - 409 `source-changed`: a stale `root_scope`; + - 413: a body over 1 MiB. + Any invalid item fails the whole request before a write. Targets that fold to one + key are written and counted once. +- PUT writes only files the source's listing holds and counts the rest in + `not_found`; when none is listed it answers 404 `files-not-found` and writes + nothing. A file that already has the value counts as `unchanged` and keeps its + `updated_at`. DELETE accepts any well-formed target, so a missing file's + category can still be removed. An unexpected failure is a logged 500 ("The + category could not be changed." / "Image categories could not be read.") without + the exception text. The routes never read or write `prompts`, `prompt_pairs` or + `generated_images`; `POST /prompt_manager/bulk/category` is unchanged. +- Duplicate cleanup (System › Maintenance): the kept copy's own category wins. A + kept copy without one receives the removed copy's (`COPY_SQL`, part of the + cleanup's query-plan check). The removed copy's row is never deleted, so a + give-back (the kept copy changed) and a failed delete lose nothing. The job log + says "Kept N image categories on the kept copies". + `tests/test_duplicate_image_categories.py` covers nested sources in both orders. +- A moved, renamed or deleted file's row stays and is hidden; a file that later + appears at the same path shows it. The real-host check is the last test in + `tests/comfyui/image-attribution.spec.mjs`: PUT, list, DELETE on a generated + output. + Image annotations schema (#249, `database/annotation_schema.py`): the group `media_annotation_metadata` (one row: schema version, `image_index_revision`), `media_annotation_files` (one private locator per physical file, diff --git a/py/api/__init__.py b/py/api/__init__.py index def2186f..0d24ec0c 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -72,6 +72,7 @@ from .diagnostics_report import DiagnosticsReportMixin from .autotag_routes import AutotagRoutesMixin from .lora_integration import LoraIntegrationMixin from .lora_import_routes import LoraImportRoutesMixin +from .image_categories import ImageCategoryRoutesMixin try: from ...database.operations import PromptDatabase @@ -336,6 +337,7 @@ class PromptManagerAPI( AutotagRoutesMixin, LoraIntegrationMixin, LoraImportRoutesMixin, + ImageCategoryRoutesMixin, ): """REST API handler for PromptManager operations and web interface. @@ -700,6 +702,7 @@ class PromptManagerAPI( self._register_autotag_routes(routes) self._register_lora_routes(routes) self._register_lora_import_routes(routes) + self._register_image_category_routes(routes) # Register gzip compression middleware (once) global _gzip_registered 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/image_categories.py b/py/api/image_categories.py new file mode 100644 index 00000000..ac2182e0 --- /dev/null +++ b/py/api/image_categories.py @@ -0,0 +1,285 @@ +"""Image category routes (#332): an image file's own category. + +``PUT /prompt_manager/images/category`` sets the own category of up to 500 +listed files; ``DELETE`` removes it (listed or not, so a missing file's +category can still be removed); ``GET /prompt_manager/images/categories`` +lists the own-category values of files the sources list. Targets are the +Images listing identity (``{root_index, root_scope, path}``, checked by +``media_target``) of an output source; LoRA folders are not sources here. +Each file is stored under its file key (``images.image_category_key``). +These routes never read or write prompts, recipes or image links. +""" + +import os +from pathlib import Path + +from aiohttp import web + +from .image_tag_jobs import MAX_JOB_TARGETS, _parse_targets +from .images import ( + _favorite_relative_key, + _favorite_root_key, + _forbidden, + _json_body_or_error, + image_category_key, +) +from .prompts import bad_request + +try: + from ...database.image_categories import MAX_CATEGORY_LENGTH, MAX_KEY_LENGTH + from ...utils.validators import validate_category +except ImportError: + from database.image_categories import MAX_CATEGORY_LENGTH, MAX_KEY_LENGTH + from utils.validators import validate_category + +CATEGORY_MEDIA_MESSAGE = "Only media files can have a category" +_PUT_FIELDS = frozenset({"items", "category"}) +_DELETE_FIELDS = frozenset({"items"}) +SOURCE_CHANGED_MESSAGE = "The sources changed. Refresh the images and try again." +WRITE_FAILED = "The category could not be changed." +READ_FAILED = "Image categories could not be read." + + +def _coded(status: int, code: str, message: str) -> web.Response: + return web.json_response( + {"success": False, "code": code, "error": message}, + status=status, + headers={"Cache-Control": "no-store"}, + ) + + +def _category_value(raw) -> tuple[str | None, web.Response | None]: + """The stripped category, or a 400 (an empty value is a 400: use DELETE).""" + if not isinstance(raw, str): + return None, bad_request("category must be a string") + value = raw.strip() + if not value: + return None, bad_request("category must not be empty") + if len(value) > MAX_CATEGORY_LENGTH: + return None, bad_request( + f"Category cannot exceed {MAX_CATEGORY_LENGTH} characters" + ) + try: + validate_category(value) + except ValueError as exc: + return None, bad_request(str(exc)) + return value, None + + +def _nested_sources(root: Path, sources) -> list[tuple[str, str]]: + """``[(relative key prefix, source root key)]`` of sources inside *root*. + + Deepest first, so the first prefix that matches is the file's source. + """ + base = os.path.join(os.path.normcase(os.path.realpath(str(root))), "") + nested = [] + for source in sources: + real = os.path.realpath(str(source)) + if not os.path.normcase(real).startswith(base): + continue + relative = os.path.relpath(real, os.path.realpath(str(root))) + nested.append( + ( + _favorite_relative_key(relative.replace(os.sep, "/")), + _favorite_root_key(source), + ) + ) + nested.sort(key=lambda item: item[0].count("/"), reverse=True) + return nested + + +def _source_prefix(relative_key: str, nested) -> str | None: + """The prefix of the deepest nested source holding *relative_key*.""" + for prefix, _source_key in nested: + if relative_key.startswith(prefix + "/"): + return prefix + return None + + +class ImageCategoryRoutesMixin: + """Routes for images' own categories (#332).""" + + def _register_image_category_routes(self, routes): + # Fixed paths beside DELETE /prompt_manager/images/{image_id}, as the + # favorite routes are: matched before the dynamic resource. + @routes.put("/prompt_manager/images/category") + async def set_image_category_route(request): + return await self.set_image_category(request) + + @routes.delete("/prompt_manager/images/category") + async def clear_image_category_route(request): + return await self.clear_image_category(request) + + @routes.get("/prompt_manager/images/categories") + async def list_image_categories_route(request): + return await self.list_image_categories(request) + + # -- reading --------------------------------------------------------------- + + def _own_category_map(self, sources, listings) -> dict: + """``{root / listed path: own category}`` for the listed files (blocking). + + One primary-key range read per source; each stored row is matched to + the listing of every selected root that holds its source, so the work + follows the stored rows, not the files. A file under a nested source + is keyed by the deepest one (``image_category_key``). A row whose file + is gone matches nothing. A failed read raises. + """ + sources = list(sources) + rows = self.db.own_categories_of(_favorite_root_key(s) for s in sources) + if not rows: + return {} + by_source: dict[str, dict[str, str]] = {} + for (root_key, relative), category in rows.items(): + by_source.setdefault(root_key, {})[relative] = category + found = {} + for _index, root, files in listings: + listing = self._favorite_listing(root, files) + root_key = _favorite_root_key(root) + nested = [ + item for item in _nested_sources(root, sources) if item[1] != root_key + ] + for relative, category in by_source.get(root_key, {}).items(): + listed = listing.get(relative) + if listed is not None and _source_prefix(relative, nested) is None: + found[root / listed] = category + for prefix, source_key in nested: + for relative, category in by_source.get(source_key, {}).items(): + full = f"{prefix}/{relative}" + listed = listing.get(full) + if listed is not None and _source_prefix(full, nested) == prefix: + found[root / listed] = category + return found + + async def _source_listings(self, roots) -> list[tuple]: + return [ + (index, root, await self._get_gallery_files(root)) + for index, root in enumerate(roots) + ] + + async def list_image_categories(self, request): + """GET: the own-category values of listed files, sorted.""" + try: + roots = list(self._get_all_output_dirs()) + listings = await self._source_listings(roots) + found = await self._run_in_executor(self._own_category_map, roots, listings) + except Exception: + self.logger.exception("Image categories read failed") + return _coded(500, "categories-failed", READ_FAILED) + return web.json_response( + {"success": True, "categories": sorted(set(found.values()))}, + headers={"Cache-Control": "no-store"}, + ) + + # -- writing --------------------------------------------------------------- + + async def _category_request(self, request, fields: frozenset): + """``(data, targets, roots, None)`` or ``(None, None, None, error)``.""" + data, error = await _json_body_or_error(request) + if error is not None: + return None, None, None, error + if set(data) != fields: + return ( + None, + None, + None, + bad_request( + "Request must contain exactly " + ", ".join(sorted(fields)) + ), + ) + items = data["items"] + if not isinstance(items, list) or not 1 <= len(items) <= MAX_JOB_TARGETS: + return ( + None, + None, + None, + bad_request(f"items must list 1 to {MAX_JOB_TARGETS} images"), + ) + targets, error = _parse_targets(items, media_message=CATEGORY_MEDIA_MESSAGE) + if error is not None: + return None, None, None, error + roots = list(self._get_all_output_dirs()) + for index, scope in dict.fromkeys((t[0], t[1]) for t in targets): + if index >= len(roots): + return None, None, None, bad_request("Invalid gallery root") + error = self._gallery_scope_error( + {"root_scope": scope, "root": str(index)}, roots + ) + if error is not None: + if error.status == 409: + error = _coded(409, "source-changed", SOURCE_CHANGED_MESSAGE) + return None, None, None, error + return data, targets, roots, None + + async def set_image_category(self, request): + """PUT ``{items, category}``: set the own category of listed files.""" + data, targets, roots, error = await self._category_request(request, _PUT_FIELDS) + if error is not None: + return error + category, error = _category_value(data["category"]) + if error is not None: + return error + try: + listings = {} + listed_keys, missing_keys = {}, {} + for index, _scope, relative in targets: + root = roots[index] + if index not in listings: + files = await self._get_gallery_files(root) + listings[index] = await self._run_in_executor( + self._favorite_listing, root, files + ) + listed = listings[index].get(_favorite_relative_key(relative)) + key = image_category_key(roots, root / (listed or relative)) + if key is None or max(len(key[0]), len(key[1])) > MAX_KEY_LENGTH: + return _forbidden("Path is too long") + (missing_keys if listed is None else listed_keys)[key] = True + for key in listed_keys: + missing_keys.pop(key, None) + if not listed_keys: + return _coded( + 404, + "files-not-found", + "None of the selected images could be found.", + ) + counts = await self._run_in_executor( + self.db.set_image_categories, list(listed_keys), category + ) + except Exception: + self.logger.exception("Image category update failed") + return _coded(500, "category-failed", WRITE_FAILED) + return web.json_response( + { + "success": True, + "category": category, + "updated": counts["updated"], + "unchanged": counts["unchanged"], + "not_found": len(missing_keys), + }, + headers={"Cache-Control": "no-store"}, + ) + + async def clear_image_category(self, request): + """DELETE ``{items}``: remove the own category, listed or not.""" + _data, targets, roots, error = await self._category_request( + request, _DELETE_FIELDS + ) + if error is not None: + return error + keys = {} + for index, _scope, relative in targets: + key = image_category_key(roots, roots[index] / relative) + if key is None or max(len(key[0]), len(key[1])) > MAX_KEY_LENGTH: + return _forbidden("Path is too long") + keys[key] = True + try: + cleared = await self._run_in_executor( + self.db.clear_image_categories, list(keys) + ) + except Exception: + self.logger.exception("Image category removal failed") + return _coded(500, "category-failed", WRITE_FAILED) + return web.json_response( + {"success": True, "cleared": cleared}, + headers={"Cache-Control": "no-store"}, + ) diff --git a/py/api/image_tag_jobs.py b/py/api/image_tag_jobs.py index 1d611745..6baafd4c 100644 --- a/py/api/image_tag_jobs.py +++ b/py/api/image_tag_jobs.py @@ -193,13 +193,17 @@ class _Spec: return scope -def _parse_targets(raw): - """Unique validated targets in request order, or ``(None, error)``.""" +def _parse_targets(raw, **target_options): + """Unique validated targets in request order, or ``(None, error)``. + + ``target_options`` reach the target check (the image category routes + pass their own ``media_message``, #332). + """ if not isinstance(raw, list) or not 1 <= len(raw) <= MAX_JOB_TARGETS: return None, bad_request(f"targets must list 1 to {MAX_JOB_TARGETS} images") seen, targets = set(), [] for item in raw: - target, error = _target(item) + target, error = _target(item, **target_options) if error is not None: return None, error if target not in seen: 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/tests/comfyui/image-attribution.spec.mjs b/tests/comfyui/image-attribution.spec.mjs index 64f14e68..65ffd655 100644 --- a/tests/comfyui/image-attribution.spec.mjs +++ b/tests/comfyui/image-attribution.spec.mjs @@ -140,3 +140,26 @@ test('an output saved without metadata links once its recorded run names a store expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true); } }); + +// #332 on the real host: the migration ran at startup and the image category routes +// set, list and remove an own category on a generated output. +test('an output gets its own image category and loses it again', async ({ execution }, testInfo) => { + const request = execution.request; + const filename = await saved(request, graph({ text: `category host ${randomUUID()}`, prefix: 'pm332/category' })); + let image; + await expect.poll(async () => { + const listing = await (await request.get('/prompt_manager/images/output?mode=workspace&q=pm332&rescan=1')).json(); + image = (listing.images || []).find(entry => entry.filename === filename || entry.path?.endsWith(`/${filename}`)); + return Boolean(image); + }, { timeout: 20000, message: `${filename} is listed by Images` }).toBe(true); + const items = [{ root_index: image.root_index, root_scope: image.root_scope, path: image.path }]; + const value = `PM332 host ${testInfo.project.name}`; + const set = await request.put('/prompt_manager/images/category', { data: { items, category: value } }); + expect(set.status()).toBe(200); + expect(await set.json()).toEqual({ success: true, category: value, updated: 1, unchanged: 0, not_found: 0 }); + expect((await (await request.get('/prompt_manager/images/categories')).json()).categories).toContain(value); + const cleared = await request.delete('/prompt_manager/images/category', { data: { items } }); + expect(cleared.status()).toBe(200); + expect(await cleared.json()).toEqual({ success: true, cleared: 1 }); + expect((await (await request.get('/prompt_manager/images/categories')).json()).categories).not.toContain(value); +}); diff --git a/tests/test_api_image_categories.py b/tests/test_api_image_categories.py new file mode 100644 index 00000000..18cd93e4 --- /dev/null +++ b/tests/test_api_image_categories.py @@ -0,0 +1,303 @@ +"""Image category routes (#332): PUT/DELETE /images/category, GET /images/categories.""" + +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 + from test_image_categories import category_checksum +except ImportError: + from tests.test_api_images import ImageAPITestCase, make_png + from tests.open_handles import assert_closed + from tests.test_image_categories import category_checksum + +from database.image_categories import MAX_KEY_LENGTH, TABLE +from py.api.images import _favorite_root_key, image_category_key + +CATEGORY = "/prompt_manager/images/category" +CATEGORIES = "/prompt_manager/images/categories" + + +class ImageCategoryTestCase(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 target(self, path, index=0): + roots = self.api._get_all_output_dirs() + return { + "root_index": index, + "root_scope": self.api._gallery_root_scope(roots[index], index), + "path": path, + } + + async def send(self, method, body, status=200): + response = await self.client.request(method, CATEGORY, json=body) + text = await response.text() + self.assertEqual(response.status, status, text) + return json.loads(text) + + async def put(self, paths, category, status=200): + items = [p if isinstance(p, dict) else self.target(p) for p in paths] + return await self.send("PUT", {"items": items, "category": category}, status) + + async def delete(self, paths, status=200): + items = [p if isinstance(p, dict) else self.target(p) for p in paths] + return await self.send("DELETE", {"items": items}, status) + + async def listed(self, status=200): + response = await self.client.get(CATEGORIES) + text = await response.text() + self.assertEqual(response.status, status, text) + return json.loads(text) + + def rows(self): + return [ + tuple(row) + for row in self.api.db.model.get_connection().execute( + f"SELECT root_key, relative_path, category FROM {TABLE} " + "ORDER BY root_key, relative_path" + ) + ] + + def key(self, relative, root=None): + roots = self.api._get_all_output_dirs() + return image_category_key(roots, (root or self.output_dir) / relative) + + +class TestSetAndClear(ImageCategoryTestCase): + async def test_put_sets_listed_and_counts_missing(self): + make_png(self.output_dir / "a.png") + make_png(self.output_dir / "sub" / "b.png") + data = await self.put(["a.png", "sub/b.png", "missing.png"], " Portraits ") + self.assertEqual( + data, + { + "success": True, + "category": "Portraits", + "updated": 2, + "unchanged": 0, + "not_found": 1, + }, + ) + root_key = _favorite_root_key(self.output_dir) + self.assertEqual( + self.rows(), + [(root_key, "a.png", "Portraits"), (root_key, "sub/b.png", "Portraits")], + ) + data = await self.put(["a.png", "sub/b.png"], "Portraits") + self.assertEqual((data["updated"], data["unchanged"]), (0, 2)) + data = await self.put(["a.png"], "Scans") + self.assertEqual((data["updated"], data["unchanged"]), (1, 0)) + self.assertEqual(self.api.db.image_category(*self.key("a.png")), "Scans") + + async def test_put_with_no_listed_file_is_404_and_writes_nothing(self): + data = await self.put(["missing.png", "gone/x.png"], "Scans", status=404) + self.assertEqual(data["code"], "files-not-found") + self.assertEqual(self.rows(), []) + + async def test_duplicate_targets_are_written_and_counted_once(self): + make_png(self.output_dir / "a.png") + item = self.target("a.png") + data = await self.put([item, dict(item)], "Scans") + self.assertEqual((data["updated"], data["not_found"]), (1, 0)) + + async def test_delete_clears_listed_and_missing_files(self): + make_png(self.output_dir / "a.png") + await self.put(["a.png"], "Scans") + self.api.db.set_image_categories([self.key("gone.png")], "Scans") + data = await self.delete(["a.png", "gone.png", "never.png"]) + self.assertEqual(data, {"success": True, "cleared": 2}) + self.assertEqual(self.rows(), []) + self.assertEqual((await self.delete(["a.png"]))["cleared"], 0) + + async def test_put_never_writes_prompts(self): + make_png(self.output_dir / "a.png") + make_png(self.output_dir / "b.png") + positive = self._save_prompt("castle at dusk", category="Landscapes") + self.api.db.link_image_to_prompt(positive, str(self.output_dir / "a.png")) + negative = self._save_prompt("blurry", category="Negatives") + self.api.db.save_prompt_pair("castle at dusk", "blurry", category="Recipes") + conn = self.api.db.model.get_connection() + before = category_checksum(conn) + links = conn.execute("SELECT * FROM generated_images").fetchall() + await self.put(["a.png", "b.png"], "Portraits") + await self.listed() + await self.delete(["a.png"]) + self.assertEqual(category_checksum(conn), before) + self.assertEqual( + conn.execute("SELECT * FROM generated_images").fetchall(), links + ) + self.assertIsNotNone(negative) + + async def test_two_sources_in_one_request(self): + second = self.comfy_dir / "renders" + make_png(self.output_dir / "a.png") + make_png(second / "a.png") + self.api._get_all_output_dirs = lambda: [self.output_dir, second] + data = await self.put([self.target("a.png", 0), self.target("a.png", 1)], "X") + self.assertEqual(data["updated"], 2) + self.assertEqual( + {row[0] for row in self.rows()}, + {_favorite_root_key(self.output_dir), _favorite_root_key(second)}, + ) + self.assertEqual((await self.delete([self.target("a.png", 1)]))["cleared"], 1) + + async def test_nested_sources_key_by_the_deepest(self): + inner = self.output_dir / "nested" + make_png(inner / "a.png") + self.api._get_all_output_dirs = lambda: [self.output_dir, inner] + # Named through the outer source, stored under the inner one. + data = await self.put([self.target("nested/a.png", 0)], "Scans") + self.assertEqual(data["updated"], 1) + self.assertEqual(self.rows(), [(_favorite_root_key(inner), "a.png", "Scans")]) + # The same file through the inner source is the same key. + data = await self.put([self.target("a.png", 1)], "Scans") + self.assertEqual(data["unchanged"], 1) + self.assertEqual((await self.listed())["categories"], ["Scans"]) + self.assertEqual((await self.delete([self.target("a.png", 1)]))["cleared"], 1) + + async def test_video_files_can_have_a_category(self): + (self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42") + data = await self.put(["clip.mp4"], "Clips") + self.assertEqual(data["updated"], 1) + + async def test_windows_case_folds_the_key(self): + make_png(self.output_dir / "Sub" / "A.png") + with patch("py.api.images._path_flavor", ntpath): + data = await self.put(["Sub/A.png"], "Scans") + self.assertEqual(data["updated"], 1) + data = await self.put(["sub/a.PNG"], "Scans") + self.assertEqual(data["unchanged"], 1) + self.assertEqual(self.rows()[0][1], "sub/a.png") + self.assertEqual((await self.delete(["SUB/A.png"]))["cleared"], 1) + + async def test_a_key_beyond_the_stored_bound_is_403(self): + make_png(self.output_dir / "a.png") + too_long = "/" + "r" * MAX_KEY_LENGTH + with patch("py.api.images._favorite_root_key", return_value=too_long): + data = await self.put(["a.png"], "Scans", status=403) + self.assertNotIn(too_long, json.dumps(data)) + await self.delete(["a.png"], status=403) + self.assertEqual(self.rows(), []) + + async def test_an_unexpected_failure_is_a_logged_500(self): + make_png(self.output_dir / "a.png") + + def broken(*_args, **_kwargs): + raise sqlite3.OperationalError("secret detail") + + self.api.db.set_image_categories = broken + self.api.db.clear_image_categories = broken + with self.assertLogs("prompt_manager", level="ERROR"): + data = await self.put(["a.png"], "Scans", status=500) + self.assertEqual(data["error"], "The category could not be changed.") + self.assertNotIn("secret", json.dumps(data)) + with self.assertLogs("prompt_manager", level="ERROR"): + data = await self.delete(["a.png"], status=500) + self.assertNotIn("secret", json.dumps(data)) + + async def test_delete_category_never_reaches_delete_image(self): + make_png(self.output_dir / "a.png") + calls = [] + + async def recorded(request): + calls.append(request.match_info) + raise AssertionError("delete_image reached") + + self.api.delete_image = recorded + await self.put(["a.png"], "Scans") + data = await self.delete(["a.png"]) + self.assertEqual(data["cleared"], 1) + self.assertEqual(calls, []) + + +class TestRefusals(ImageCategoryTestCase): + async def test_malformed_bodies_are_400(self): + make_png(self.output_dir / "a.png") + item = self.target("a.png") + cases = [ + {"items": [item]}, + {"category": "x"}, + {"items": [item], "category": "x", "extra": 1}, + {"items": item, "category": "x"}, + {"items": [], "category": "x"}, + {"items": [item] * 501, "category": "x"}, + {"items": [item], "category": 3}, + {"items": [item], "category": None}, + {"items": [item], "category": " \t "}, + {"items": [item], "category": "x" * 101}, + {"items": [item], "category": "a\x07b"}, + {"items": [{**item, "root_index": "0"}], "category": "x"}, + {"items": [{**item, "root_scope": "nothex"}], "category": "x"}, + {"items": [{"root_index": 0, "path": "a.png"}], "category": "x"}, + {"items": [{**item, "root_index": 1}], "category": "x"}, + ] + for body in cases: + with self.subTest(body=str(body)[:80]): + data = await self.send("PUT", body, status=400) + self.assertFalse(data["success"]) + await self.send("DELETE", {"items": [item], "category": "x"}, status=400) + await self.send("DELETE", {"items": []}, status=400) + response = await self.client.request( + "PUT", CATEGORY, data="not json", headers={"Content-Type": "text/plain"} + ) + self.assertEqual(response.status, 400) + self.assertEqual(self.rows(), []) + + async def test_unsafe_paths_are_403(self): + (self.output_dir / "notes.txt").write_text("not media") + for path in ("../outside.png", "/a.png", "C:/a.png", "a\\b.png", "a/./b.png"): + with self.subTest(path=path): + await self.put([self.target(path)], "x", status=403) + data = await self.put([self.target("notes.txt")], "x", status=403) + self.assertEqual(data["error"], "Only media files can have a category") + await self.delete([self.target("notes.txt")], status=403) + + async def test_a_stale_scope_is_409(self): + make_png(self.output_dir / "a.png") + item = {**self.target("a.png"), "root_scope": "0" * 64} + data = await self.put([item], "x", status=409) + self.assertEqual(data["code"], "source-changed") + data = await self.delete([item], status=409) + self.assertEqual(data["code"], "source-changed") + + async def test_an_oversized_body_is_413(self): + response = await self.client.request( + "PUT", + CATEGORY, + data=b'{"items": [], "category": "' + b"x" * (2 * 1024 * 1024) + b'"}', + headers={"Content-Type": "application/json"}, + ) + self.assertEqual(response.status, 413) + + +class TestCategoryList(ImageCategoryTestCase): + async def test_lists_only_values_whose_file_is_listed(self): + make_png(self.output_dir / "a.png") + make_png(self.output_dir / "c.png") + await self.put(["a.png"], "B") + await self.put(["c.png"], "A") + self.api.db.set_image_categories([self.key("gone.png")], "Z") + self.assertEqual( + await self.listed(), {"success": True, "categories": ["A", "B"]} + ) + + async def test_empty_library_lists_nothing(self): + self.assertEqual((await self.listed())["categories"], []) + + async def test_a_failed_read_is_500(self): + make_png(self.output_dir / "a.png") + + def broken(*_args, **_kwargs): + raise sqlite3.OperationalError("secret detail") + + self.api.db.own_categories_of = broken + with self.assertLogs("prompt_manager", level="ERROR"): + data = await self.listed(status=500) + self.assertEqual(data["error"], "Image categories could not be read.") + self.assertNotIn("secret", json.dumps(data))