From b5a56bee0b5c2df9d3f483f0f7a6b133e604a4c9 Mon Sep 17 00:00:00 2001 From: Vito Date: Wed, 28 Oct 2026 15:29:41 +0000 Subject: [PATCH] feat(images): store images' own categories and set them by route (#332, PR 1) (#401) * feat(images): store images' own categories in the library database (#332) One new table, image_categories, and its category index, created empty by a forward-only migration that reads and writes only its own objects. Startup checks the structure; restore also checks the values. The store sets, clears and reads own categories by file key (the deepest output source holding the file and the path below it), and copies a removed duplicate's category to a kept copy that has none. * feat(images): duplicate cleanup keeps images' own categories (#332) The kept copy's own category wins; a kept copy without one receives the removed copy's, keyed by the deepest output source holding each file. The removed copy's row is never deleted, so a give-back or a failed delete loses nothing. The job log reports how many were kept. * feat(images): set, remove and list images' own categories (#332) PUT /prompt_manager/images/category sets the own category of up to 500 listed files, DELETE removes it (listed or not), and GET /prompt_manager/images/categories lists the values whose file is listed. Targets are the Images listing identity of an output source, checked by media_target (shared with Collections). The routes never touch prompts, recipes or links; POST /prompt_manager/bulk/category is unchanged. * fix(images): key symlinked output folders and keep key work off the event loop (#332) Review fixes: a source matches by its configured spelling too, so a detected output folder that is a symlink keys its files (routes and duplicate cleanup). File keys are computed in the executor. A path no source holds is a plain 403, a lone surrogate a 400. The install repairs a table that lacks its category index instead of refusing to open. * feat(images): set or remove a category on all N matching images (#332) The category routes also accept the Images listing filter of "Select all N matching" (#392) instead of explicit items, resolved now and capped at 500 images like an explicit selection. * test(vendor): count the #332 host case in the execution matrix image-attribution.spec.mjs gains the image category route case, so the execution config lists 92 tests (3 attribution cases per project). * test(legacy): record the image category routes in the legacy contract snapshot (#332) #332 adds PUT/DELETE /prompt_manager/images/category and GET /prompt_manager/images/categories; the legacy pages' routes are unchanged. --- TESTING.md | 64 +++ database/duplicate_copies.py | 15 +- database/image_categories.py | 249 +++++++++ database/models.py | 13 +- database/operations.py | 2 + py/api/__init__.py | 3 + py/api/annotations.py | 44 +- py/api/duplicate_files.py | 4 + py/api/duplicate_spellings.py | 25 +- py/api/image_categories.py | 361 ++++++++++++ py/api/image_tag_jobs.py | 10 +- py/api/images.py | 47 ++ py/api/media_target.py | 54 ++ tests/browser/server.py | 1 + tests/browser/workspace_concept_fixture.py | 1 + .../workspace_library_editor_fixture.py | 1 + tests/browser/workspace_review_fixture.py | 1 + tests/browser/workspace_tagging_fixture.py | 1 + tests/browser/workspace_visual_fixture.py | 1 + tests/browser/workspace_vocabulary_fixture.py | 1 + tests/comfyui/image-attribution.spec.mjs | 23 + tests/fixtures/legacy_contract_snapshot.json | 3 + tests/test_api_image_categories.py | 375 +++++++++++++ tests/test_duplicate_image_categories.py | 151 +++++ tests/test_image_categories.py | 523 ++++++++++++++++++ tests/vendor/comfyui-config.test.mjs | 7 +- 26 files changed, 1933 insertions(+), 47 deletions(-) create mode 100644 database/image_categories.py create mode 100644 py/api/image_categories.py create mode 100644 py/api/media_target.py create mode 100644 tests/test_api_image_categories.py create mode 100644 tests/test_duplicate_image_categories.py create mode 100644 tests/test_image_categories.py diff --git a/TESTING.md b/TESTING.md index d958cabe..36dc168b 100644 --- a/TESTING.md +++ b/TESTING.md @@ -4676,6 +4676,70 @@ 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. A source + matches by its real path or its configured spelling (a detected output folder + can be a symlink). Without nested sources it equals a favorite's key. Every write and read uses it. A path + that no source contains has no key, so it has no own category. +- `image_categories_schema_error` checks the structure at startup (a value saved + under older rules never stops the library opening; a table without its index, + which this code never leaves, is repaired by the install). Restore also checks the + values: text columns, keys within the favorites bound, `category` 1 to 100 + characters with no character below U+0020 and no leading or trailing space, and + `updated_at` of at most 64 bytes. Backups made before #332 restore without the + 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. Instead of `items` a request may send `filter`, + the Images listing's own query ("Select all N matching", #392, keys as in + `parse_listing_filter`), resolved now: over 500 matches is a 400 + `too-many-images` ("Select at most 500 images."), a renumbered source a 409 + `source-changed`. +- PUT writes only files the source's listing holds and counts the rest in + `not_found`; when none is listed it answers 404 `files-not-found` and writes + nothing. A file that already has the value counts as `unchanged` and keeps its + `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/database/duplicate_copies.py b/database/duplicate_copies.py index 2baea72b..22129786 100644 --- a/database/duplicate_copies.py +++ b/database/duplicate_copies.py @@ -11,6 +11,8 @@ IMMEDIATE`` transaction, so no row ever names a deleted file: - its favorite; - its image tags (``media_image_tags`` of its byte record), joined onto the kept copy's byte record as their union, in order; +- its own image category (#332), copied to the kept copy when that copy has + none; the removed copy's row is kept, so a give-back loses nothing; - its pending Needs review items, settled: only an item still ``pending`` at revision 0 with no decision row is removed (``review._SETTLE_ITEM``), so a decided item is never touched. @@ -22,6 +24,8 @@ import os from .annotation_schema import BYTES_TABLE, FILES_TABLE from .annotations import _BUMP_INDEX +from .image_categories import COPY_SQL as COPY_CATEGORY_SQL +from .image_categories import copy_image_category from .models import normalize_image_path, utc_now_iso from .review import _BUMP_REVIEW_INDEX, _PENDING_OF_LOCATOR, _SETTLE_ITEM from .review_changes import notify_review_change @@ -76,6 +80,7 @@ FILE_LINK_SQL = ( COPY_BYTE_SQL, COPY_TAG_COUNT_SQL, PENDING_COUNT_SQL, + COPY_CATEGORY_SQL, ) @@ -137,8 +142,10 @@ class DuplicateCopyOperations: removed key, kept root key, kept key)`` per listed root holding the file), ``removed_locator``, ``kept_locator``, ``sha256`` and the kept copy's ``kept_size`` and - ``kept_mtime_ns``. Returns ``{moved, merged, favorites, tags, - review}``. + ``kept_mtime_ns``; ``categories``: ``(removed file key, kept file + key)`` pairs of the image categories (#332). Returns ``{moved, + merged, favorites, tags, review, categories}`` (``categories``: the + image categories copied to the kept copy). """ counts = dict.fromkeys(("moved", "merged", "favorites", "tags", "review"), 0) kept_keys = [normalize_image_path(path) for path in move["kept"]] @@ -174,6 +181,10 @@ class DuplicateCopyOperations: for favorite in move.get("favorites", ()): counts["favorites"] += _move_favorite(conn, *favorite) counts["tags"] = self._move_copy_tags(conn, move) + counts["categories"] = sum( + copy_image_category(conn, removed, kept) + for removed, kept in move.get("categories", ()) + ) for (item_id,) in conn.execute( _PENDING_OF_LOCATOR, (move["removed_locator"],) ).fetchall(): diff --git a/database/image_categories.py b/database/image_categories.py new file mode 100644 index 00000000..a1f5d5a1 --- /dev/null +++ b/database/image_categories.py @@ -0,0 +1,249 @@ +"""Image categories (#332): an image file's own category. + +An image shows its own category when it has one, otherwise the category its +link supplies (its recipe's, else its positive prompt's); this table holds +only the own values. A row is keyed by the file key: the deepest configured +output source that contains the file (its real, case-normalized path) and the +path below it, case-folded where the OS folds case +(``py/api/images.py`` ``image_category_key``). Without nested sources the key +equals a favorite's key. Nothing follows a moved file: its row stays and is +hidden until a file appears at the same path again. + +Writes validate strictly (the route strips the value and applies +``validate_category``). Startup checks the structure only, so a value saved +under older rules never stops the library opening; restore also checks the +values. Optional in old backups, required once a database is current. +""" + +import re +import sqlite3 +from typing import Iterable + +TABLE = "image_categories" +INDEX = "idx_image_categories_category" +MAX_KEY_LENGTH = 4096 +MAX_CATEGORY_LENGTH = 100 +_DDL = ( + f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT NOT NULL, " + "category TEXT NOT NULL, updated_at TEXT NOT NULL, " + "PRIMARY KEY (root_key, relative_path))" +) +_INDEX_DDL = f"CREATE INDEX {INDEX} ON {TABLE}(category)" +_AUTOINDEX = f"sqlite_autoindex_{TABLE}_1" +_COLUMNS = ( + ("root_key", "TEXT", 1, 1), + ("relative_path", "TEXT", 1, 2), + ("category", "TEXT", 1, 0), + ("updated_at", "TEXT", 1, 0), +) +# A category may not hold a character below U+0020 (validate_category). +_CONTROL_GLOB = "*[" + "".join(chr(code) for code in range(1, 32)) + "]*" + +CATEGORY_OF_SQL = ( + f"SELECT category FROM {TABLE} WHERE root_key = ? AND relative_path = ?" +) +CATEGORIES_OF_ROOT_SQL = ( + f"SELECT relative_path, category FROM {TABLE} WHERE root_key = ?" +) +UPSERT_SQL = ( + f"INSERT INTO {TABLE} (root_key, relative_path, category, updated_at) " + "VALUES (?, ?, ?, ?) ON CONFLICT(root_key, relative_path) " + "DO UPDATE SET category = excluded.category, updated_at = excluded.updated_at" +) +DELETE_SQL = f"DELETE FROM {TABLE} WHERE root_key = ? AND relative_path = ?" +# The duplicate cleanup's copy (#332 MC11): the removed copy's row is copied +# to the kept copy only when the kept copy has none; it is never deleted. +COPY_SQL = ( + f"INSERT OR IGNORE INTO {TABLE} (root_key, relative_path, category, updated_at) " + f"SELECT ?, ?, category, updated_at FROM {TABLE} " + "WHERE root_key = ? AND relative_path = ?" +) +STATEMENTS = (CATEGORY_OF_SQL, CATEGORIES_OF_ROOT_SQL, UPSERT_SQL, DELETE_SQL, COPY_SQL) + + +def _compact(sql: str) -> str: + return re.sub(r"\s+", "", sql or "").lower() + + +def _structure_error( + conn: sqlite3.Connection, allow_missing: bool, allow_missing_index: bool = False +) -> str | None: + objects = conn.execute( + "SELECT type,name,sql FROM sqlite_master WHERE name=? " + "OR (type IN ('index','trigger') AND tbl_name=?)", + (TABLE, TABLE), + ).fetchall() + if not objects: + return None if allow_missing else "Missing image categories schema" + tables = [sql for kind, _name, sql in objects if kind == "table"] + others = {(kind, name) for kind, name, _sql in objects if kind != "table"} + expected = {("index", _AUTOINDEX), ("index", INDEX)} + # Before its install, a table without the category index is repaired by + # CREATE INDEX IF NOT EXISTS rather than refused. + if allow_missing_index and others == {("index", _AUTOINDEX)}: + expected = others + if len(tables) != 1 or others != expected: + return "Image categories name is reserved by an unsupported schema object" + if _compact(tables[0]) != _compact(_DDL): + return "Image categories has an invalid table definition" + index_sql = next((sql for kind, name, sql in objects if name == INDEX), None) + if ("index", INDEX) in expected and _compact(index_sql) != _compact(_INDEX_DDL): + return "Image categories has an invalid index" + autoindex_sql = next(sql for kind, name, sql in objects if name == _AUTOINDEX) + if autoindex_sql is not None: + return "Image categories has an invalid index" + columns = tuple( + (row[1], row[2].upper(), row[3], row[5]) + for row in conn.execute(f"PRAGMA table_info({TABLE})") + ) + if columns != _COLUMNS: + return "Image categories has invalid columns" + return None + + +def _values_error(conn: sqlite3.Connection) -> str | None: + # Bytes bound characters from above (UTF-8 is at most 4 bytes each). + limit = MAX_KEY_LENGTH * 4 + if conn.execute( + f"SELECT 1 FROM {TABLE} WHERE typeof(root_key) != 'text' " # nosec B608 + "OR typeof(relative_path) != 'text' OR typeof(category) != 'text' " + "OR typeof(updated_at) != 'text' " + f"OR length(CAST(root_key AS BLOB)) > {limit} " + f"OR length(CAST(relative_path AS BLOB)) > {limit} " + "OR length(CAST(updated_at AS BLOB)) > 64 " + "OR instr(CAST(category AS BLOB), x'00') > 0 " + f"OR length(category) < 1 OR length(category) > {MAX_CATEGORY_LENGTH} " + "OR category != TRIM(category) OR category GLOB ? LIMIT 1", + (_CONTROL_GLOB,), + ).fetchone(): + return "Image categories contains invalid values" + for root, path in conn.execute( + f"SELECT CAST(root_key AS BLOB), CAST(relative_path AS BLOB) " # nosec B608 + f"FROM {TABLE} WHERE length(CAST(root_key AS BLOB)) > {MAX_KEY_LENGTH} " + f"OR length(CAST(relative_path AS BLOB)) > {MAX_KEY_LENGTH}" + ): + for raw in (root, path): + if len(bytes(raw).decode("utf-8", "replace")) > MAX_KEY_LENGTH: + return "Image categories contains invalid values" + return None + + +def image_categories_schema_error( + conn: sqlite3.Connection, + *, + allow_missing: bool = True, + check_values: bool = True, +) -> str | None: + """Reject an unsupported image categories schema (and, optionally, values). + + Reads only. A missing table is acceptable for backups that predate it + unless ``allow_missing`` is false. ``check_values`` also refuses a stored + value the write path could never have written. + """ + error = _structure_error(conn, allow_missing) + if error is not None or not check_values: + return error + exists = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (TABLE,) + ).fetchone() + return _values_error(conn) if exists else None + + +def migrate_image_categories(conn: sqlite3.Connection) -> None: + """Install the table and its index atomically; structure is checked + before and after. Reads only ``sqlite_master`` and this table's PRAGMAs + and writes only its own two objects (no backfill).""" + conn.execute("SAVEPOINT image_categories_schema") + try: + error = _structure_error(conn, True, allow_missing_index=True) + if error: + raise sqlite3.DatabaseError(error) + conn.execute(_DDL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1)) + conn.execute( + _INDEX_DDL.replace("CREATE INDEX", "CREATE INDEX IF NOT EXISTS", 1) + ) + error = image_categories_schema_error( + conn, allow_missing=False, check_values=False + ) + if error: + raise sqlite3.DatabaseError(error) + conn.execute("RELEASE image_categories_schema") + except BaseException: + conn.execute("ROLLBACK TO image_categories_schema") + conn.execute("RELEASE image_categories_schema") + raise + + +def copy_image_category(conn: sqlite3.Connection, removed_key, kept_key) -> int: + """Give the kept copy the removed copy's category when it has none. + + Runs inside the caller's transaction; never deletes the removed copy's + row, so a give-back or a failed delete loses nothing. Returns 1 when a + row was inserted, else 0. + """ + if tuple(removed_key) == tuple(kept_key): + return 0 + return max(conn.execute(COPY_SQL, (*kept_key, *removed_key)).rowcount, 0) + + +class ImageCategoryOperations: + """PromptDatabase mixin: set, clear and read images' own categories. + + Never reads or writes ``prompts``, ``prompt_pairs`` or + ``generated_images``. + """ + + def image_category(self, root_key: str, relative_key: str) -> str | None: + """One file's own category, or None.""" + row = ( + self.model.get_connection() + .execute(CATEGORY_OF_SQL, (root_key, relative_key)) + .fetchone() + ) + return row[0] if row else None + + def own_categories_of(self, root_keys: Iterable[str]) -> dict: + """``{(root key, relative key): category}`` under the given sources. + + One primary-key range read per source. + """ + found = {} + conn = self.model.get_connection() + for root_key in dict.fromkeys(root_keys): + for relative, category in conn.execute(CATEGORIES_OF_ROOT_SQL, (root_key,)): + found[(root_key, relative)] = category + return found + + def set_image_categories(self, file_keys, category: str) -> dict: + """Set *category* on every file key, in one transaction. + + ``file_keys``: ``(root key, relative key)`` pairs from any number of + sources. Returns ``{updated, unchanged}``; a file that already had + exactly this value is ``unchanged`` and keeps its ``updated_at``. + """ + from .models import utc_now_iso # models imports this module's schema + + keys = list(dict.fromkeys(tuple(key) for key in file_keys)) + counts = {"updated": 0, "unchanged": 0} + now = utc_now_iso() + with self._pair_transaction() as conn: + for key in keys: + row = conn.execute(CATEGORY_OF_SQL, key).fetchone() + if row is not None and row[0] == category: + counts["unchanged"] += 1 + continue + conn.execute(UPSERT_SQL, (*key, category, now)) + counts["updated"] += 1 + return counts + + def clear_image_categories(self, file_keys) -> int: + """Remove the own category of every file key (one transaction). + + Returns the number of rows removed. + """ + keys = list(dict.fromkeys(tuple(key) for key in file_keys)) + removed = 0 + with self._pair_transaction() as conn: + for key in keys: + removed += conn.execute(DELETE_SQL, key).rowcount + return removed diff --git a/database/models.py b/database/models.py index dad9cf07..25e98e44 100644 --- a/database/models.py +++ b/database/models.py @@ -34,6 +34,7 @@ from .review_schema import ( review_without_annotations_error, ) from .saved_views import migrate_saved_views, saved_views_schema_error +from .image_categories import image_categories_schema_error, migrate_image_categories from .workflow_documents import ( DOCUMENT_ID_COLUMNS, migrate_workflow_documents, @@ -350,6 +351,8 @@ class PromptModel: migrate_annotation_state(conn) migrate_review_state(conn) migrate_autotag_state(conn) + # Image categories (#332): independent of every table above. + migrate_image_categories(conn) # Last: the older generated_images rebuilds run only before file_path. migrate_workflow_documents(conn) @@ -1155,8 +1158,9 @@ class PromptModel: Args: path: File to inspect require_library_state: Also require the tables every current - database has (Library state, image favorites, saved views, - image annotations and the review queue), as a migrated + database has (Library state, image favorites, image + categories, saved views, image annotations and the review + queue), as a migrated restore stage must; old backups may lack them. documents: Read every workflow document and refuse a damaged one. A restore checks its upload without it and repairs @@ -1220,6 +1224,11 @@ class PromptModel: ) if favorites_error: return False, favorites_error + categories_error = image_categories_schema_error( + conn, allow_missing=not require_library_state + ) + if categories_error: + return False, categories_error views_error = saved_views_schema_error( conn, allow_missing=not require_library_state ) diff --git a/database/operations.py b/database/operations.py index 0f45097f..4460fcf3 100644 --- a/database/operations.py +++ b/database/operations.py @@ -27,6 +27,7 @@ from .health import HealthOperations, iter_consistency_issues from .duplicate_copies import DuplicateCopyOperations from .workflow_storage import WorkflowStorageOperations from .category_import import CategoryImportOperations, EmptyReplacement # noqa: F401 +from .image_categories import ImageCategoryOperations from .workflow_documents import ( DocumentReader, documents_of_rows, @@ -368,6 +369,7 @@ class PromptDatabase( DuplicateCopyOperations, WorkflowStorageOperations, CategoryImportOperations, + ImageCategoryOperations, ): """Database operations class for managing prompts.""" 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/duplicate_files.py b/py/api/duplicate_files.py index 6edf265d..dc7f511c 100644 --- a/py/api/duplicate_files.py +++ b/py/api/duplicate_files.py @@ -589,6 +589,7 @@ class DuplicateFilesMixin: spellings = await self._run_in_executor(self._source_spellings) removed, failure, keepers = 0, None, {} totals = dict.fromkeys(("moved", "merged", "favorites", "tags", "review"), 0) + totals["categories"] = 0 # image categories kept on the kept copies (#332) for done, target in enumerate(targets, start=1): if failure is not None: self._record_target(job, "kept", target, NOT_ATTEMPTED) @@ -635,6 +636,9 @@ class DuplicateFilesMixin: job.log_line( f"Added {_plural(totals['tags'], 'image tag')} to the kept copies" ) + if totals.get("categories"): + kept = _plural(totals["categories"], "image category", "image categories") + job.log_line(f"Kept {kept} on the kept copies") if totals["review"]: job.log_line( f"Settled {_plural(totals['review'], 'Needs review item')} " diff --git a/py/api/duplicate_spellings.py b/py/api/duplicate_spellings.py index 31e0f0b9..0e197aa4 100644 --- a/py/api/duplicate_spellings.py +++ b/py/api/duplicate_spellings.py @@ -18,7 +18,7 @@ Blocking: each lookup stats the file's folders once (cached). import os from pathlib import Path -from .images import _favorite_relative_key, _favorite_root_key +from .images import _favorite_relative_key, _favorite_root_key, image_category_key def _identity(path): @@ -78,6 +78,15 @@ class Spellings: ) return spelled + @property + def listed_roots(self): + """The listed roots (the output sources), in order.""" + return [root for root, _identity in self._listed] + + def category_key(self, path): + """*path*'s image category file key (#332), or None outside every source.""" + return image_category_key(self.listed_roots, path) + def favorites(self, path): """``{listed root: path below it (POSIX)}`` for each listed root holding it.""" path = os.path.abspath(str(path)) @@ -125,9 +134,23 @@ def copy_move(target, kept_stat, spellings): "targets": targets, "kept_path": target["keeper_abs"], "favorites": favorites, + "categories": _category_pairs(target, spellings), "removed_locator": target["locator"], "kept_locator": target["keeper_locator"], "sha256": target["hash"], "kept_size": kept_stat.st_size, "kept_mtime_ns": kept_stat.st_mtime_ns, } + + +def _category_pairs(target, spellings): + """``[(removed file key, kept file key)]`` for the image category (#332). + + From the two absolute paths, not the favorites tuples: those follow the + listed root, which differs from the file key when sources nest. + """ + removed = spellings.category_key(target["abs"]) + kept = spellings.category_key(target["keeper_abs"]) + if removed is None or kept is None: + return [] + return [(removed, kept)] diff --git a/py/api/image_categories.py b/py/api/image_categories.py new file mode 100644 index 00000000..0b497bae --- /dev/null +++ b/py/api/image_categories.py @@ -0,0 +1,361 @@ +"""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. +Instead of ``items`` a request may name ``filter``, the Images listing's own +query ("Select all N matching", #392), resolved now up to 500 images. +Each file is stored under its file key (``images.image_category_key``). +These routes never read or write prompts, recipes or image links. +""" + +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({"category"}) +_DELETE_FIELDS = frozenset() +TOO_MANY = f"Select at most {MAX_JOB_TARGETS} images." +SOURCE_CHANGED_MESSAGE = "The sources changed. Refresh the images and try again." +WRITE_FAILED = "The category could not be changed." +READ_FAILED = "Image categories could not be read." + + +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: + value.encode("utf-8") + validate_category(value) + except UnicodeEncodeError: + return None, bad_request("category must be valid text") + except ValueError as exc: + return None, bad_request(str(exc)) + return value, None + + +class _KeyRefused(Exception): + """A target has no storable file key (outside every source, or too long).""" + + def __init__(self, message: str = "Access denied"): + super().__init__(message) + self.message = message + + +def _file_key(roots, root, relative) -> tuple[str, str]: + """The target's file key (blocking: resolves the sources).""" + key = image_category_key(roots, Path(root) / relative) + if key is None: + raise _KeyRefused() + if max(len(key[0]), len(key[1])) > MAX_KEY_LENGTH: + raise _KeyRefused("Path is too long") + return key + + +def _nested_sources(root: Path, sources) -> list[tuple[str, str]]: + """``[(relative key prefix, source root key)]`` of sources inside *root*. + + 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. Registered after it, they still win: aiohttp + # (3.10+) looks fixed paths up before dynamic resources + # (test_delete_category_never_reaches_delete_image). + @routes.put("/prompt_manager/images/category") + async def set_image_category_route(request): + return await self.set_image_category(request) + + @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 _filter_targets(self, raw): + """``(targets, None)`` for an Images listing filter, else ``(None, error)``. + + "Select all N matching" (#392) sends the listing's own query; it is + resolved now, up to 500 images (a 400 ``too-many-images`` above). + """ + from .images import ListingSourceChanged, parse_listing_filter + + try: + selection = await self.resolve_listing_filter( + parse_listing_filter(raw), MAX_JOB_TARGETS + ) + except ListingSourceChanged: + return None, _coded(409, "source-changed", SOURCE_CHANGED_MESSAGE) + except ValueError as exc: + return None, bad_request(str(exc)) + if selection.truncated: + return None, _coded(400, "too-many-images", TOO_MANY) + return tuple(selection.targets), None + + async def _category_request(self, request, fields: frozenset): + """``(data, targets, roots, None)`` or ``(None, None, None, error)``. + + The images are ``items`` (1 to 500 targets) or ``filter`` (all images + an Images listing filter admits), besides *fields*. + """ + data, error = await _json_body_or_error(request) + if error is not None: + return None, None, None, error + names = set(data) - fields + if set(data) & fields != fields or names not in ({"items"}, {"filter"}): + return ( + None, + None, + None, + bad_request( + "Request must contain " + + ", ".join(sorted(fields)) + + (" and " if fields else "") + + "either items or filter" + ), + ) + if "filter" in data: + targets, error = await self._filter_targets(data["filter"]) + if error is not None: + return None, None, None, error + return data, targets, list(self._get_all_output_dirs()), None + items = data["items"] + if not isinstance(items, list) or not 1 <= len(items) <= MAX_JOB_TARGETS: + return ( + 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 + try: + for _index, _scope, relative in targets: + relative.encode("utf-8") + except UnicodeEncodeError: + return None, None, None, bad_request("path must be valid text") + roots = list(self._get_all_output_dirs()) + for index, scope in dict.fromkeys((t[0], t[1]) for t in targets): + if index >= len(roots): + 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 + + def _listed_file_keys(self, roots, targets, files) -> tuple[list, list]: + """``(listed keys, unlisted keys)`` of the targets (blocking). + + A target counts as listed when its source's listing holds it; its key + then follows the listed spelling. Targets folding to one key count + once, and a key that is listed through any target is listed. + """ + listings, listed, missing = {}, {}, {} + for index, _scope, relative in targets: + root = roots[index] + if index not in listings: + listings[index] = self._favorite_listing(root, files[index]) + spelled = listings[index].get(_favorite_relative_key(relative)) + key = _file_key(roots, root, spelled or relative) + (missing if spelled is None else listed)[key] = True + return list(listed), [key for key in missing if key not in listed] + + async def set_image_category(self, request): + """PUT ``{items, category}``: set the own category of listed files.""" + data, targets, roots, error = await self._category_request(request, _PUT_FIELDS) + if error is not None: + return error + category, error = _category_value(data["category"]) + if error is not None: + return error + try: + files = {} + for index in dict.fromkeys(t[0] for t in targets): + files[index] = await self._get_gallery_files(roots[index]) + listed_keys, missing_keys = await self._run_in_executor( + self._listed_file_keys, roots, targets, files + ) + except _KeyRefused as refused: + return _forbidden(refused.message) + except Exception: + self.logger.exception("Image category update failed") + return _coded(500, "category-failed", WRITE_FAILED) + try: + if not listed_keys: + return _coded( + 404, + "files-not-found", + "None of the selected images could be found.", + ) + counts = await self._run_in_executor( + self.db.set_image_categories, 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 + try: + keys = await self._run_in_executor( + lambda: list( + dict.fromkeys(_file_key(roots, roots[i], r) for i, _s, r in targets) + ) + ) + cleared = await self._run_in_executor(self.db.clear_image_categories, keys) + except _KeyRefused as refused: + return _forbidden(refused.message) + except Exception: + self.logger.exception("Image category removal failed") + return _coded(500, "category-failed", WRITE_FAILED) + 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/images.py b/py/api/images.py index 7d5987ef..c652dc59 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -802,6 +802,53 @@ def _favorite_root_key(root) -> str: return _favorite_relative_key(os.path.realpath(str(root))) +def _source_depth(source: str) -> int: + return len([part for part in re.split(r"[\\/]+", source) if part]) + + +def deepest_source(sources, absolute_path): + """``(source, POSIX path below it)`` of the deepest source holding the path. + + A source matches by its real path or by its spelling as configured (a + detected output folder can be a symlink), and the deepest real path wins. + The path is compared lexically after ``normpath`` and never resolved, so a + symlinked subfolder keeps its listed spelling. Containment folds case + where the OS does. *source* is the real path. None when no source + contains the path. + """ + path = os.path.normpath(str(absolute_path)) + folded = os.path.normcase(path) + best = None + for source in sources: + real = os.path.normpath(os.path.realpath(str(source))) + for spelling in dict.fromkeys((real, os.path.normpath(str(source)))): + prefix = os.path.join(os.path.normcase(spelling), "") + if not folded.startswith(prefix): + continue + if best is None or _source_depth(real) > _source_depth(best[0]): + best = (real, spelling) + break + if best is None: + return None + relative = os.path.relpath(path, best[1]) + return best[0], relative.replace(os.sep, "/") + + +def image_category_key(sources, absolute_path): + """The file key of an image's own category (#332), or None. + + ``(root key, relative key)``: the deepest output source that contains the + file and the path below it, keyed with the favorites helpers. Without + nested sources it equals the file's favorite key. A path no source + contains has no key (and so no own category). + """ + found = deepest_source(sources, absolute_path) + if found is None: + return None + source, relative = found + return _favorite_root_key(source), _favorite_relative_key(relative) + + def _name_suffix(name): """``Path(name).suffix`` as Python 3.10-3.13 compute it. 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/browser/server.py b/tests/browser/server.py index 71ea5669..160f4b0a 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -268,6 +268,7 @@ def create_app(root: Path) -> web.Application: "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_concept_fixture.py b/tests/browser/workspace_concept_fixture.py index 1c86dbbd..0c54e831 100644 --- a/tests/browser/workspace_concept_fixture.py +++ b/tests/browser/workspace_concept_fixture.py @@ -181,6 +181,7 @@ _TABLES = ( "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_library_editor_fixture.py b/tests/browser/workspace_library_editor_fixture.py index 8c0aed9e..ffc0534a 100644 --- a/tests/browser/workspace_library_editor_fixture.py +++ b/tests/browser/workspace_library_editor_fixture.py @@ -33,6 +33,7 @@ _TABLES = ( "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_review_fixture.py b/tests/browser/workspace_review_fixture.py index 45365dd7..16344eb2 100644 --- a/tests/browser/workspace_review_fixture.py +++ b/tests/browser/workspace_review_fixture.py @@ -85,6 +85,7 @@ _TABLES = ( "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_tagging_fixture.py b/tests/browser/workspace_tagging_fixture.py index aaf412fe..e1250763 100644 --- a/tests/browser/workspace_tagging_fixture.py +++ b/tests/browser/workspace_tagging_fixture.py @@ -51,6 +51,7 @@ _TABLES = ( "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_visual_fixture.py b/tests/browser/workspace_visual_fixture.py index f10cf457..6fbe1fcc 100644 --- a/tests/browser/workspace_visual_fixture.py +++ b/tests/browser/workspace_visual_fixture.py @@ -119,6 +119,7 @@ def seed_workspace_visual(api: PromptManagerAPI, root: Path) -> None: "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_vocabulary_fixture.py b/tests/browser/workspace_vocabulary_fixture.py index 17901872..95568762 100644 --- a/tests/browser/workspace_vocabulary_fixture.py +++ b/tests/browser/workspace_vocabulary_fixture.py @@ -86,6 +86,7 @@ _TABLES = ( "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", 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/fixtures/legacy_contract_snapshot.json b/tests/fixtures/legacy_contract_snapshot.json index 50a48f53..eb2085cd 100644 --- a/tests/fixtures/legacy_contract_snapshot.json +++ b/tests/fixtures/legacy_contract_snapshot.json @@ -2,6 +2,7 @@ "routes": [ "DELETE /prompt_manager/delete/{prompt_id}", "DELETE /prompt_manager/image-annotations/capture/{token}", + "DELETE /prompt_manager/images/category", "DELETE /prompt_manager/images/favorite", "DELETE /prompt_manager/images/{image_id}", "DELETE /prompt_manager/pairs/{pair_id}", @@ -26,6 +27,7 @@ "GET /prompt_manager/history/runs/{run_id}/outputs", "GET /prompt_manager/image-annotations/capture/{token}", "GET /prompt_manager/images/all", + "GET /prompt_manager/images/categories", "GET /prompt_manager/images/inspect/{image_path}", "GET /prompt_manager/images/output", "GET /prompt_manager/images/prompt/{image_path}", @@ -156,6 +158,7 @@ "POST /prompt_manager/vocabulary/tags", "POST /prompt_manager/vocabulary/tags/{tag_id}/delete", "PUT /prompt_manager/image-annotations/tags", + "PUT /prompt_manager/images/category", "PUT /prompt_manager/images/favorite", "PUT /prompt_manager/prompts/{prompt_id}", "PUT /prompt_manager/prompts/{prompt_id}/rating", diff --git a/tests/test_api_image_categories.py b/tests/test_api_image_categories.py new file mode 100644 index 00000000..02905267 --- /dev/null +++ b/tests/test_api_image_categories.py @@ -0,0 +1,375 @@ +"""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_a_symlinked_output_folder_keys_like_its_favorites(self): + """A detected output folder can be a symlink: keys use its real path.""" + real = self.comfy_dir / "real-output" + make_png(real / "a.png") + link = self.comfy_dir / "linked-output" + self._make_symlink(real, link, target_is_directory=True) + self.api._get_all_output_dirs = lambda: [link] + data = await self.put(["a.png"], "Scans") + self.assertEqual(data["updated"], 1) + self.assertEqual(self.rows(), [(_favorite_root_key(link), "a.png", "Scans")]) + self.assertEqual((await self.listed())["categories"], ["Scans"]) + self.assertEqual((await self.delete(["a.png"]))["cleared"], 1) + + async def test_lone_surrogates_are_400_not_500(self): + make_png(self.output_dir / "a.png") + await self.send( + "PUT", {"items": [self.target("a.png")], "category": "x\ud800"}, 400 + ) + await self.send("DELETE", {"items": [self.target("a\ud800.png")]}, 400) + self.assertEqual(self.rows(), []) + + async def test_delete_category_never_reaches_delete_image(self): + make_png(self.output_dir / "a.png") + calls = [] + + 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)) + + +class TestAllMatchingFilter(ImageCategoryTestCase): + """#392 "Select all N matching": the listing filter, resolved now.""" + + async def test_filter_sets_and_clears_every_match(self): + for name in ("bulk/a.png", "bulk/b.png", "other/c.png"): + make_png(self.output_dir / name) + data = await self.send( + "PUT", {"filter": {"subfolder": "bulk"}, "category": "Scans"} + ) + self.assertEqual((data["updated"], data["not_found"]), (2, 0)) + self.assertEqual( + sorted(row[1] for row in self.rows()), ["bulk/a.png", "bulk/b.png"] + ) + data = await self.send("DELETE", {"filter": {"subfolder": "bulk"}}) + self.assertEqual(data["cleared"], 2) + data = await self.send( + "PUT", {"filter": {"subfolder": "none"}, "category": "Scans"}, 404 + ) + self.assertEqual(data["code"], "files-not-found") + self.assertEqual( + (await self.send("DELETE", {"filter": {"q": "zz"}}))["cleared"], 0 + ) + + async def test_filter_refusals(self): + make_png(self.output_dir / "a.png") + item = self.target("a.png") + for body in ( + {"filter": {}, "items": [item], "category": "x"}, + {"filter": {"unknown": "1"}, "category": "x"}, + {"filter": "q=a", "category": "x"}, + {"filter": {"root": "0", "root_scope": "nothex"}, "category": "x"}, + ): + with self.subTest(body=body): + await self.send("PUT", body, 400) + data = await self.send( + "PUT", + {"filter": {"root": "0", "root_scope": "0" * 64}, "category": "x"}, + 409, + ) + self.assertEqual(data["code"], "source-changed") + self.assertEqual(self.rows(), []) + + async def test_more_than_500_matching_is_refused(self): + for number in range(501): + make_png(self.output_dir / f"many/{number:03}.png", size=(1, 1)) + data = await self.send("PUT", {"filter": {}, "category": "x"}, 400) + self.assertEqual(data["code"], "too-many-images") + self.assertEqual(data["error"], "Select at most 500 images.") + self.assertEqual(self.rows(), []) diff --git a/tests/test_duplicate_image_categories.py b/tests/test_duplicate_image_categories.py new file mode 100644 index 00000000..79452a68 --- /dev/null +++ b/tests/test_duplicate_image_categories.py @@ -0,0 +1,151 @@ +"""Duplicate cleanup keeps images' own categories (#332 MC11). + +The kept copy's own category wins; a kept copy without one receives the +removed copy's. The removed copy's row is never deleted, so a give-back +(the kept copy changed) and a failed delete lose nothing. +""" + +import asyncio +import os +from unittest.mock import patch + +try: + from tests.test_api_duplicate_files import WAIT, DuplicateFilesTestCase +except ImportError: # discovered with tests/ as the top-level directory + from test_api_duplicate_files import WAIT, DuplicateFilesTestCase + +from py.api import duplicate_files # noqa: E402 +from py.api.images import image_category_key # noqa: E402 +from py.config import GalleryConfig # noqa: E402 + + +class TestDuplicateCleanupKeepsCategories(DuplicateFilesTestCase): + def key(self, path): + return image_category_key(self.api._get_all_output_dirs(), path) + + def category(self, path): + return self.api.db.image_category(*self.key(path)) + + def set_category(self, path, value): + self.api.db.set_image_categories([self.key(path)], value) + + async def run_cleanup(self, keep="a.png"): + job, group = await self.found_group() + plan = await self.review(job.id, [self.keep_of(group, keep)]) + body = await self.apply(plan["token"]) + cleanup = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT) + return cleanup + + async def test_kept_copy_without_a_category_receives_the_removed_one(self): + self.trio() + b = self.output_dir / "b.png" + self.set_category(b, "Portraits") + cleanup = await self.run_cleanup() + self.assertFalse(b.exists()) + self.assertEqual(self.category(self.output_dir / "a.png"), "Portraits") + # The removed copy's row stays (hidden while its file is gone). + self.assertEqual(self.category(b), "Portraits") + self.assertIn( + "Kept 1 image category on the kept copies", "\n".join(cleanup.log) + ) + + async def test_kept_copy_keeps_its_own_category(self): + self.trio() + self.set_category(self.output_dir / "a.png", "Landscapes") + self.set_category(self.output_dir / "b.png", "Portraits") + self.set_category(self.output_dir / "sub" / "c.png", "Scans") + cleanup = await self.run_cleanup() + self.assertEqual(self.category(self.output_dir / "a.png"), "Landscapes") + self.assertNotIn("image categor", "\n".join(cleanup.log)) + + async def test_two_removed_copies_give_the_first_category(self): + self.trio() + self.set_category(self.output_dir / "b.png", "Portraits") + self.set_category(self.output_dir / "sub" / "c.png", "Scans") + cleanup = await self.run_cleanup() + self.assertIn(self.category(self.output_dir / "a.png"), {"Portraits", "Scans"}) + self.assertIn( + "Kept 1 image category on the kept copies", "\n".join(cleanup.log) + ) + + async def test_give_back_loses_no_category(self): + self.trio() + self.set_category(self.output_dir / "b.png", "Portraits") + with patch.object(duplicate_files, "still_kept", return_value=False): + cleanup = await self.run_cleanup() + self.assertEqual(cleanup.counters["removed"], 0) + rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]} + self.assertEqual( + rows["b.png"], duplicate_files.ROW_DETAILS[duplicate_files.KEEPER_GONE] + ) + self.assertTrue((self.output_dir / "b.png").is_file()) + self.assertEqual(self.category(self.output_dir / "b.png"), "Portraits") + + async def test_a_failed_delete_keeps_both_categories(self): + self.trio() + self.set_category(self.output_dir / "b.png", "Portraits") + real = duplicate_files._unlink + + def unlink(path): + if os.path.basename(path) == "b.png": + raise PermissionError(13, "denied") + real(path) + + with patch.object(duplicate_files, "_unlink", side_effect=unlink): + await self.run_cleanup() + self.assertTrue((self.output_dir / "b.png").is_file()) + self.assertEqual(self.category(self.output_dir / "b.png"), "Portraits") + self.assertEqual(self.category(self.output_dir / "a.png"), "Portraits") + + async def test_nested_sources_use_the_deepest_source_keys(self): + self.write("sub/a.png", b"same", mtime=1) + self.write("sub/b.png", b"same", mtime=2) + # The subfolder is listed first, so the cleanup's listed root is the + # file key's source; the outer source still holds the files. + GalleryConfig.MONITORING_DIRECTORIES = [ + self._root(self.output_dir / "sub"), + self._root(self.output_dir), + ] + sub_key = duplicate_files.favorite_root_key(self.output_dir / "sub") + self.api.db.set_image_categories([(sub_key, "b.png")], "Portraits") + await self.run_cleanup() + self.assertFalse((self.output_dir / "sub" / "b.png").exists()) + self.assertEqual(self.api.db.image_category(sub_key, "a.png"), "Portraits") + outer_key = duplicate_files.favorite_root_key(self.output_dir) + self.assertIsNone(self.api.db.image_category(outer_key, "sub/a.png")) + + async def test_outer_source_listed_first_still_keys_by_the_deepest(self): + self.write("sub/a.png", b"same", mtime=1) + self.write("sub/b.png", b"same", mtime=2) + GalleryConfig.MONITORING_DIRECTORIES = [ + self._root(self.output_dir), + self._root(self.output_dir / "sub"), + ] + sub_key = duplicate_files.favorite_root_key(self.output_dir / "sub") + self.api.db.set_image_categories([(sub_key, "b.png")], "Portraits") + job, group = await self.found_group() + [kept] = [f for f in group["files"] if f["path"].endswith("a.png")][:1] + plan = await self.review( + job.id, [self.keep_of(group, kept["path"], kept["root_index"])] + ) + body = await self.apply(plan["token"]) + cleanup = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT) + self.assertFalse((self.output_dir / "sub" / "b.png").exists()) + self.assertEqual(self.api.db.image_category(sub_key, "a.png"), "Portraits") + + async def test_a_symlinked_output_folder_still_moves_the_category(self): + link = self.output_dir.parent / "linked-output" + try: + os.symlink(self.output_dir, link, target_is_directory=True) + except (OSError, NotImplementedError): + self.skipTest("symlinks are not available") + self.trio() + with patch.object(self.api, "_get_all_output_dirs", return_value=[link]): + key = image_category_key([link], link / "b.png") + self.api.db.set_image_categories([key], "Portraits") + await self.run_cleanup() + kept = image_category_key([link], link / "a.png") + self.assertFalse((self.output_dir / "b.png").exists()) + self.assertEqual(self.api.db.image_category(*kept), "Portraits") diff --git a/tests/test_image_categories.py b/tests/test_image_categories.py new file mode 100644 index 00000000..9972c70e --- /dev/null +++ b/tests/test_image_categories.py @@ -0,0 +1,523 @@ +"""Image categories (#332): schema, migration, restore check, store and file key.""" + +import hashlib +import ntpath +import os +import sqlite3 +import tempfile +import unittest +from contextlib import closing +from pathlib import Path +from unittest.mock import patch + +try: + from open_handles import assert_closed +except ImportError: + from tests.open_handles import assert_closed + +from database.image_categories import ( + COPY_SQL, + INDEX, + STATEMENTS, + TABLE, + copy_image_category, + image_categories_schema_error, + migrate_image_categories, +) +from database.models import PromptModel +from database.operations import PromptDatabase +from py.api.images import image_category_key + +_DDL = ( + f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT NOT NULL, " + "category TEXT NOT NULL, updated_at TEXT NOT NULL, " + "PRIMARY KEY (root_key, relative_path))" +) +_INDEX_DDL = f"CREATE INDEX {INDEX} ON {TABLE}(category)" + + +def category_checksum(conn) -> tuple: + """sha256 of prompts and prompt_pairs ``(id, category, updated_at)``.""" + sums = [] + for table in ("prompts", "prompt_pairs"): + digest = hashlib.sha256() + for row in conn.execute( # nosec B608 - fixed table names + f"SELECT id, category, updated_at FROM {table} ORDER BY id" + ): + digest.update(repr(tuple(row)).encode()) + sums.append(digest.hexdigest()) + return tuple(sums) + + +def table_checksums(conn) -> dict: + """``{name: (schema sql, sha256 of rows)}`` of every table but ours.""" + found = {} + for name, sql in conn.execute( + "SELECT name, sql FROM sqlite_master WHERE type='table' AND name != ?", + (TABLE,), + ).fetchall(): + digest = hashlib.sha256() + rows = conn.execute(f'SELECT * FROM "{name}"') # nosec B608 + for text in sorted(repr(tuple(row)) for row in rows): + digest.update(text.encode()) + found[name] = (sql, digest.hexdigest()) + return found + + +class _TempDir(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + self.path = os.path.join(self.tmpdir, "categories.db") + + def tearDown(self): + PromptModel.close_all_instances() + assert_closed(self, self.tmpdir) + for name in os.listdir(self.tmpdir): + os.unlink(os.path.join(self.tmpdir, name)) + os.rmdir(self.tmpdir) + + +class TestImageCategoriesSchema(_TempDir): + def test_migration_creates_table_and_index(self): + db = PromptDatabase(self.path) + db.close_all() + PromptModel.close_all_instances() + with closing(sqlite3.connect(self.path)) as conn: + names = { + (kind, name) + for kind, name in conn.execute( + "SELECT type, name FROM sqlite_master WHERE tbl_name=?", (TABLE,) + ) + } + self.assertEqual( + names, + { + ("table", TABLE), + ("index", INDEX), + ("index", f"sqlite_autoindex_{TABLE}_1"), + }, + ) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {TABLE}").fetchone()[0], 0 + ) + self.assertIsNone(image_categories_schema_error(conn, allow_missing=False)) + + def test_migration_leaves_prompt_categories_unchanged(self): + """23,000 links; the migration reads and writes only its own objects.""" + db = PromptDatabase(self.path) + db.close_all() + PromptModel.close_all_instances() + with closing(sqlite3.connect(self.path)) as conn: + conn.execute(f"DROP TABLE {TABLE}") # a library from before #332 + categories = ["Landscapes", "Portraits", None, " ", "Scans"] + conn.executemany( + "INSERT INTO prompts (text, category, hash, updated_at) " + "VALUES (?, ?, ?, ?)", + [ + (f"p{n}", categories[n % 5], f"h{n}", f"2026-01-01 00:00:{n % 60}") + for n in range(200) + ], + ) + conn.executemany( + "INSERT INTO prompt_pairs (positive_prompt_id, negative_prompt_id, " + "category) VALUES (?, ?, ?)", + [(n, n + 100, categories[(n + 1) % 5]) for n in range(1, 50)], + ) + conn.executemany( + "INSERT INTO generated_images (prompt_id, image_path, filename, " + "file_path) VALUES (?, ?, ?, ?)", + [ + (n % 200 + 1, f"/out/{n}.png", f"{n}.png", f"/out/{n}.png") + for n in range(23000) + ], + ) + conn.commit() + before = table_checksums(conn) + prompts_before = category_checksum(conn) + seen = [] + + def authorize(action, arg1, arg2, _db, _trigger): + seen.append((action, arg1, arg2)) + return sqlite3.SQLITE_OK + + conn.set_authorizer(authorize) + migrate_image_categories(conn) + conn.commit() + conn.set_authorizer(None) + own = {TABLE, "sqlite_master", "sqlite_schema", INDEX} + touched = set() + for action, arg1, arg2 in seen: + if action in (sqlite3.SQLITE_READ, sqlite3.SQLITE_INSERT): + touched.add(arg1) + if action == sqlite3.SQLITE_READ: + self.assertIn(arg1, own, (action, arg1, arg2)) + elif action in ( + sqlite3.SQLITE_INSERT, + sqlite3.SQLITE_UPDATE, + sqlite3.SQLITE_DELETE, + ): + self.assertIn(arg1, own, (action, arg1, arg2)) + elif action == sqlite3.SQLITE_CREATE_TABLE: + self.assertEqual(arg1, TABLE) + elif action == sqlite3.SQLITE_CREATE_INDEX: + self.assertIn(arg1, (INDEX, f"sqlite_autoindex_{TABLE}_1")) + self.assertEqual(arg2, TABLE) + elif action == sqlite3.SQLITE_PRAGMA: + self.assertEqual((arg1, arg2), ("table_info", TABLE)) + self.assertIn(TABLE, touched) # the authorizer saw the work + self.assertEqual(table_checksums(conn), before) + self.assertEqual(category_checksum(conn), prompts_before) + schema = conn.execute( + "SELECT sql FROM sqlite_master ORDER BY name" + ).fetchall() + migrate_image_categories(conn) # a second run changes nothing + conn.commit() + self.assertEqual( + conn.execute("SELECT sql FROM sqlite_master ORDER BY name").fetchall(), + schema, + ) + # The whole startup sequence keeps them too. + PromptDatabase(self.path).close_all() + PromptModel.close_all_instances() + with closing(sqlite3.connect(self.path)) as conn: + self.assertEqual(category_checksum(conn), prompts_before) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {TABLE}").fetchone()[0], 0 + ) + + def test_structure_accepted_and_refused(self): + variants = { + "missing": None, + "extra column": ( + f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT " + "NOT NULL, category TEXT NOT NULL, updated_at TEXT NOT NULL, x TEXT, " + "PRIMARY KEY (root_key, relative_path))", + _INDEX_DDL, + ), + "wrong index": (_DDL, f"CREATE INDEX {INDEX} ON {TABLE}(updated_at)"), + "extra index": ( + _DDL, + _INDEX_DDL, + f"CREATE INDEX extra ON {TABLE}(root_key)", + ), + "trigger": ( + _DDL, + _INDEX_DDL, + f"CREATE TRIGGER t AFTER INSERT ON {TABLE} BEGIN SELECT 1; END", + ), + "wrong ddl": ( + f"CREATE TABLE {TABLE} (root_key TEXT, relative_path TEXT, " + "category TEXT, updated_at TEXT, " + "PRIMARY KEY (root_key, relative_path))", + _INDEX_DDL, + ), + "view": (f"CREATE VIEW {TABLE} AS SELECT 1",), + } + for label, statements in variants.items(): + with self.subTest(label), closing(sqlite3.connect(":memory:")) as conn: + for sql in statements or (): + conn.execute(sql) + error = image_categories_schema_error(conn) + if statements is None: + self.assertIsNone(error) + self.assertIsNotNone( + image_categories_schema_error(conn, allow_missing=False) + ) + continue + self.assertIsNotNone(error) + self.assertIsNotNone( + image_categories_schema_error(conn, check_values=False) + ) + with self.assertRaises(sqlite3.DatabaseError): + migrate_image_categories(conn) + # A table without its category index (never written by this code) is + # refused by the check but repaired by the install. + with closing(sqlite3.connect(":memory:")) as conn: + conn.execute(_DDL) + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', 'a.png', 'Scans', 'now')") + self.assertIsNotNone(image_categories_schema_error(conn)) + migrate_image_categories(conn) + self.assertIsNone(image_categories_schema_error(conn, allow_missing=False)) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {TABLE}").fetchone()[0], 1 + ) + with closing(sqlite3.connect(":memory:")) as conn: + migrate_image_categories(conn) + migrate_image_categories(conn) # idempotent + self.assertIsNone(image_categories_schema_error(conn, allow_missing=False)) + + def test_values_refused_but_startup_checks_structure_only(self): + bad = { + "blank": "", + "101 characters": "x" * 101, + "control character": "a\tb", + "nul": "a\x00b", + "leading space": " Landscapes", + "trailing space": "Landscapes ", + } + for label, value in bad.items(): + with self.subTest(label), closing(sqlite3.connect(":memory:")) as conn: + migrate_image_categories(conn) + conn.execute( + f"INSERT INTO {TABLE} VALUES ('r', 'a.png', ?, 'now')", (value,) + ) + self.assertIsNotNone(image_categories_schema_error(conn)) + self.assertIsNone( + image_categories_schema_error(conn, check_values=False) + ) + migrate_image_categories(conn) # startup still opens it + with closing(sqlite3.connect(":memory:")) as conn: + migrate_image_categories(conn) + conn.execute( + f"INSERT INTO {TABLE} VALUES ('r', ?, 'Scans', 'now')", ("x" * 5000,) + ) + self.assertIsNotNone(image_categories_schema_error(conn)) + conn.execute(f"DELETE FROM {TABLE}") + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', x'00ff', 'Scans', 'now')") + self.assertIsNotNone(image_categories_schema_error(conn)) + conn.execute(f"DELETE FROM {TABLE}") + for value in ("x" * 100, "Scans", "Café · 風景", "a b"): + conn.execute( + f"INSERT OR REPLACE INTO {TABLE} VALUES ('r', 'a.png', ?, 'now')", + (value,), + ) + self.assertIsNone(image_categories_schema_error(conn), value) + + def test_restore_check(self): + PromptDatabase(self.path).close_all() + PromptModel.close_all_instances() + ok, reason = PromptModel.verify_database_file( + self.path, require_library_state=True + ) + self.assertTrue(ok, reason) + with closing(sqlite3.connect(self.path)) as conn: + conn.execute(f"DROP TABLE {TABLE}") + conn.commit() + # A backup from before #332 is accepted ... + ok, reason = PromptModel.verify_database_file(self.path) + self.assertTrue(ok, reason) + # ... but a current database must have the table. + ok, reason = PromptModel.verify_database_file( + self.path, require_library_state=True + ) + self.assertFalse(ok) + self.assertIn("image categories", reason.lower()) + with closing(sqlite3.connect(self.path)) as conn: + conn.execute(_DDL) + conn.execute(_INDEX_DDL) + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', 'a.png', ' x', 'now')") + conn.commit() + ok, reason = PromptModel.verify_database_file(self.path) + self.assertFalse(ok) + self.assertIn("image categories", reason.lower()) + + +class TestImageCategoryStore(_TempDir): + def setUp(self): + super().setUp() + self.db = PromptDatabase(self.path) + + def tearDown(self): + self.db.close_all() + super().tearDown() + + def rows(self): + return ( + self.db.model.get_connection() + .execute( + f"SELECT root_key, relative_path, category, updated_at FROM {TABLE} " + "ORDER BY root_key, relative_path" + ) + .fetchall() + ) + + def test_set_counts_updated_and_unchanged(self): + keys = [("/out", "a.png"), ("/other", "a.png"), ("/out", "a.png")] + self.assertEqual( + self.db.set_image_categories(keys, "Portraits"), + {"updated": 2, "unchanged": 0}, + ) + stamps = {(r[0], r[1]): r[3] for r in self.rows()} + self.assertEqual( + self.db.set_image_categories( + [("/out", "a.png"), ("/out", "b.png")], "Portraits" + ), + {"updated": 1, "unchanged": 1}, + ) + # An unchanged file keeps its updated_at. + self.assertEqual( + {(r[0], r[1]): r[3] for r in self.rows()}[("/out", "a.png")], + stamps[("/out", "a.png")], + ) + self.assertEqual( + self.db.set_image_categories([("/out", "a.png")], "Scans"), + {"updated": 1, "unchanged": 0}, + ) + self.assertEqual(self.db.image_category("/out", "a.png"), "Scans") + self.assertEqual(self.db.image_category("/other", "a.png"), "Portraits") + self.assertIsNone(self.db.image_category("/out", "missing.png")) + self.assertEqual( + self.db.own_categories_of(["/out", "/out", "/none"]), + {("/out", "a.png"): "Scans", ("/out", "b.png"): "Portraits"}, + ) + + def test_clear_counts_removed_rows(self): + self.db.set_image_categories([("/out", "a.png"), ("/b", "c.png")], "Scans") + self.assertEqual( + self.db.clear_image_categories( + [("/out", "a.png"), ("/b", "c.png"), ("/out", "gone.png")] + ), + 2, + ) + self.assertEqual(self.rows(), []) + self.assertEqual(self.db.clear_image_categories([("/out", "a.png")]), 0) + + def test_a_failed_write_writes_nothing(self): + self.db.set_image_categories([("/out", "a.png")], "Scans") + keys = [("/out", "b.png"), ("/out", None)] # NOT NULL fails the second + with self.assertRaises(sqlite3.IntegrityError): + self.db.set_image_categories(keys, "Portraits") + self.assertEqual([r[:3] for r in self.rows()], [("/out", "a.png", "Scans")]) + + def test_store_never_touches_prompt_categories(self): + conn = self.db.model.get_connection() + with conn: + conn.execute( + "INSERT INTO prompts (text, category, hash) VALUES ('p', 'L', 'h')" + ) + before = category_checksum(conn) + seen = set() + conn.set_authorizer( + lambda action, arg1, *_: seen.add(arg1) or sqlite3.SQLITE_OK + ) + try: + self.db.set_image_categories([("/out", "a.png")], "Scans") + self.db.own_categories_of(["/out"]) + self.db.image_category("/out", "a.png") + self.db.clear_image_categories([("/out", "a.png")]) + finally: + conn.set_authorizer(None) + self.assertFalse(seen & {"prompts", "prompt_pairs", "generated_images"}, seen) + self.assertEqual(category_checksum(conn), before) + + def test_copy_keeps_kept_copy_category(self): + self.db.set_image_categories([("/out", "removed.png")], "Scans") + self.db.set_image_categories([("/out", "kept.png")], "Portraits") + with self.db._pair_transaction() as conn: + inserted = copy_image_category( + conn, ("/out", "removed.png"), ("/out", "kept.png") + ) + self.assertEqual(inserted, 0) + self.assertEqual(self.db.image_category("/out", "kept.png"), "Portraits") + with self.db._pair_transaction() as conn: + inserted = copy_image_category( + conn, ("/out", "removed.png"), ("/out", "new.png") + ) + self.assertEqual(inserted, 1) + self.assertEqual(self.db.image_category("/out", "new.png"), "Scans") + # The removed copy's row is never deleted. + self.assertEqual(self.db.image_category("/out", "removed.png"), "Scans") + with self.db._pair_transaction() as conn: + self.assertEqual( + copy_image_category(conn, ("/out", "none.png"), ("/out", "x.png")), 0 + ) + self.assertEqual( + copy_image_category( + conn, ("/out", "removed.png"), ("/out", "removed.png") + ), + 0, + ) + + def test_statements_never_scan_the_table(self): + conn = self.db.model.get_connection() + for sql in STATEMENTS: + with self.subTest(sql): + params = (None,) * sql.count("?") + plan = " ".join( + row[-1] for row in conn.execute(f"EXPLAIN QUERY PLAN {sql}", params) + ) + self.assertNotRegex(plan, rf"SCAN {TABLE}\b", plan) + plan = " ".join( + row[-1] + for row in conn.execute( + f"EXPLAIN QUERY PLAN SELECT rowid FROM {TABLE} WHERE category=? " + "ORDER BY rowid", + ("x",), + ) + ) + self.assertIn(INDEX, plan) + self.assertIn(COPY_SQL, STATEMENTS) + + +class TestImageCategoryKey(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.base = Path(os.path.realpath(self._tmp.name)) + self.outer = self.base / "output" + self.inner = self.outer / "nested" + (self.inner / "sub").mkdir(parents=True) + + def tearDown(self): + self._tmp.cleanup() + + def test_equals_the_favorite_key_without_nesting(self): + from py.api.images import _favorite_relative_key, _favorite_root_key + + key = image_category_key([self.outer], self.outer / "sub" / "a.png") + self.assertEqual( + key, (_favorite_root_key(self.outer), _favorite_relative_key("sub/a.png")) + ) + + def test_deepest_source_wins_in_any_order(self): + from py.api.images import _favorite_root_key + + path = self.inner / "sub" / "a.png" + for sources in ([self.outer, self.inner], [self.inner, self.outer]): + with self.subTest(sources=sources): + self.assertEqual( + image_category_key(sources, path), + (_favorite_root_key(self.inner), "sub/a.png"), + ) + # A file of the outer source only is keyed by it. + self.assertEqual( + image_category_key([self.outer, self.inner], self.outer / "b.png"), + (_favorite_root_key(self.outer), "b.png"), + ) + + def test_outside_every_source_has_no_key(self): + self.assertIsNone(image_category_key([self.inner], self.outer / "a.png")) + self.assertIsNone(image_category_key([], self.outer / "a.png")) + # A sibling sharing a name prefix is not inside. + self.assertIsNone( + image_category_key([self.outer], self.base / "output2" / "a.png") + ) + self.assertIsNone(image_category_key([self.outer], self.outer)) + + def test_symlinked_subfolder_is_not_resolved(self): + target = self.base / "elsewhere" + target.mkdir() + link = self.outer / "linked" + try: + os.symlink(target, link, target_is_directory=True) + except (OSError, NotImplementedError): + self.skipTest("symlinks are not available") + key = image_category_key([self.outer], link / "a.png") + self.assertEqual(key[1], "linked/a.png") + + def test_windows_folds_case_of_the_relative_key(self): + with patch("py.api.images._path_flavor", ntpath): + upper = image_category_key([self.outer], self.outer / "Sub" / "A.PNG") + lower = image_category_key([self.outer], self.outer / "sub" / "a.png") + self.assertEqual(upper[1], "sub/a.png") + self.assertEqual(upper, lower) + + def test_posix_keeps_case_distinct(self): + if os.path.normcase("A") == "a": + self.skipTest("this OS folds case") + upper = image_category_key([self.outer], self.outer / "A.png") + lower = image_category_key([self.outer], self.outer / "a.png") + self.assertNotEqual(upper, lower) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/vendor/comfyui-config.test.mjs b/tests/vendor/comfyui-config.test.mjs index c5b00677..61b10ef4 100644 --- a/tests/vendor/comfyui-config.test.mjs +++ b/tests/vendor/comfyui-config.test.mjs @@ -33,14 +33,15 @@ test('execution script discovers only the new matrix and preserves all 78 origin const executions = list('tools/vendor/playwright.execution.config.mjs'); assert.match(executions, /\[classic-start\]/); assert.match(executions, /\[nodes2-start\]/); - // + the fresh-host library location case, + #342 image attribution, + #343 linked text - assert.match(executions, /Total: 90 tests in 5 files/); + // + the fresh-host library location case, + #342 image attribution, + #343 linked text, + // + #332 image category routes on a generated output (image-attribution.spec.mjs, 2 projects) + assert.match(executions, /Total: 92 tests in 5 files/); for (const project of ['classic-start', 'nodes2-start']) { const cases = executions.split('\n').filter(line => line.includes(`[${project}]`)); assert.equal(cases.filter(line => line.includes('execution-history.spec.mjs')).length, 15); assert.equal(cases.filter(line => line.includes('execution-outputs.spec.mjs')).length, 8); assert.equal(cases.filter(line => line.includes('library-location.spec.mjs')).length, 1); - assert.equal(cases.filter(line => line.includes('image-attribution.spec.mjs')).length, 2); + assert.equal(cases.filter(line => line.includes('image-attribution.spec.mjs')).length, 3); assert.equal(cases.filter(line => line.includes('execution-text-link.spec.mjs')).length, 19); } assert.doesNotMatch(executions, /prompt-manager.spec/);