diff --git a/database/duplicate_copies.py b/database/duplicate_copies.py index 2baea72b..ee4c07b4 100644 --- a/database/duplicate_copies.py +++ b/database/duplicate_copies.py @@ -9,6 +9,8 @@ IMMEDIATE`` transaction, so no row ever names a deleted file: already links the kept copy: then its saved-pair links join that record and it is deleted; - its favorite; +- its collection memberships (#331), joined onto the kept copy's as their + union under the same keys a favorite moves by; - its image tags (``media_image_tags`` of its byte record), joined onto the kept copy's byte record as their union, in order; - its pending Needs review items, settled: only an item still ``pending`` at @@ -50,6 +52,15 @@ INSERT_FAVORITE_SQL = "INSERT OR IGNORE INTO image_favorites VALUES (?, ?, ?)" DELETE_FAVORITE_SQL = ( "DELETE FROM image_favorites WHERE root_key = ? AND relative_path = ?" ) +COPY_MEMBERSHIPS_SQL = ( + "INSERT OR IGNORE INTO image_collection_items " + "(collection_id, root_key, relative_path, added_at) " + "SELECT collection_id, ?, ?, added_at FROM image_collection_items " + "WHERE root_key = ? AND relative_path = ?" +) +DELETE_MEMBERSHIPS_SQL = ( + "DELETE FROM image_collection_items WHERE root_key = ? AND relative_path = ?" +) COPY_BYTE_SQL = ( f"SELECT b.id FROM {FILES_TABLE} f JOIN {BYTES_TABLE} b " # nosec B608 "ON b.file_id=f.id WHERE f.locator=? AND b.sha256=?" @@ -73,6 +84,8 @@ FILE_LINK_SQL = ( DELETE_LINK_SQL, FAVORITE_SQL, DELETE_FAVORITE_SQL, + COPY_MEMBERSHIPS_SQL, + DELETE_MEMBERSHIPS_SQL, COPY_BYTE_SQL, COPY_TAG_COUNT_SQL, PENDING_COUNT_SQL, @@ -135,12 +148,15 @@ class DuplicateCopyOperations: by), ``targets`` (``{removed spelling: kept spelling}``, the spelling a moved record takes; else ``kept_path``), ``favorites`` (``(root key, 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}``. + file; they move its favorite and its collection memberships), + ``removed_locator``, ``kept_locator``, ``sha256`` and the kept copy's + ``kept_size`` and ``kept_mtime_ns``. Returns ``{moved, merged, + favorites, tags, review, collections}`` (``collections``: the + memberships inserted on the kept copy). """ - counts = dict.fromkeys(("moved", "merged", "favorites", "tags", "review"), 0) + counts = dict.fromkeys( + ("moved", "merged", "favorites", "tags", "review", "collections"), 0 + ) kept_keys = [normalize_image_path(path) for path in move["kept"]] with self._pair_transaction() as conn: for spelled in dict.fromkeys(move["removed"]): @@ -173,6 +189,7 @@ class DuplicateCopyOperations: counts["moved"] += 1 for favorite in move.get("favorites", ()): counts["favorites"] += _move_favorite(conn, *favorite) + counts["collections"] += _move_memberships(conn, *favorite) counts["tags"] = self._move_copy_tags(conn, move) for (item_id,) in conn.execute( _PENDING_OF_LOCATOR, (move["removed_locator"],) @@ -213,3 +230,12 @@ def _move_favorite(conn, root_key, removed_key, kept_root, kept_key): conn.execute(INSERT_FAVORITE_SQL, (kept_root, kept_key, row[0])) conn.execute(DELETE_FAVORITE_SQL, (root_key, removed_key)) return 1 + + +def _move_memberships(conn, root_key, removed_key, kept_root, kept_key): + """Copy a removed copy's memberships to the kept copy; the rows inserted.""" + inserted = conn.execute( + COPY_MEMBERSHIPS_SQL, (kept_root, kept_key, root_key, removed_key) + ).rowcount + conn.execute(DELETE_MEMBERSHIPS_SQL, (root_key, removed_key)) + return inserted diff --git a/database/image_collections.py b/database/image_collections.py new file mode 100644 index 00000000..43193b0b --- /dev/null +++ b/database/image_collections.py @@ -0,0 +1,459 @@ +"""Collections (#331): named, hand-picked groups of image files. + +A member names its file exactly as a favorite does (``image_favorites.py``): +the source folder's real, case-normalized path (``root_key``) and the path +below it (``relative_path``), both server-side only. A member whose file is +moved, renamed or deleted outside PromptManager is kept and simply not listed; +nothing follows a file to a new place. + +Names are trimmed with each run of whitespace collapsed to one space, 1 to 100 +characters without control characters; ``name_key`` (NFC of the case-folded +NFC name) is UNIQUE, so names are unique ignoring case and Unicode form. + +The two tables install like ``image_favorites`` and ``saved_views``: optional +in old backups, required once a database is current, created empty with no +backfill. Startup checks only the structure, so a value saved under older +rules never stops the library from opening; restore checks the values too. +""" + +import re +import sqlite3 +import unicodedata +from typing import Optional + +from .image_favorites import MAX_KEY_LENGTH + +COLLECTIONS_TABLE = "image_collections" +ITEMS_TABLE = "image_collection_items" +FILE_INDEX = "idx_image_collection_items_file" +TABLES = (COLLECTIONS_TABLE, ITEMS_TABLE) +MAX_NAME_LENGTH = 100 +_TIMESTAMP_BYTES = 64 +_CONTROL = re.compile(r"[\x00-\x1f\x7f]") +_WHITESPACE = re.compile(r"\s+") + +_TABLE_DDL = { + COLLECTIONS_TABLE: ( + f"CREATE TABLE {COLLECTIONS_TABLE} (id INTEGER PRIMARY KEY AUTOINCREMENT, " + "name TEXT NOT NULL, name_key TEXT NOT NULL UNIQUE, " + "created_at TEXT NOT NULL)" + ), + ITEMS_TABLE: ( + f"CREATE TABLE {ITEMS_TABLE} (collection_id INTEGER NOT NULL REFERENCES " + f"{COLLECTIONS_TABLE}(id) ON DELETE CASCADE, root_key TEXT NOT NULL, " + "relative_path TEXT NOT NULL, added_at TEXT NOT NULL, " + "PRIMARY KEY (collection_id, root_key, relative_path))" + ), +} +_INDEX_DDL = f"CREATE INDEX {FILE_INDEX} ON {ITEMS_TABLE}(root_key, relative_path)" +# (name, declared type, not null, default, primary key position) +_COLUMNS = { + COLLECTIONS_TABLE: ( + ("id", "INTEGER", 0, None, 1), + ("name", "TEXT", 1, None, 0), + ("name_key", "TEXT", 1, None, 0), + ("created_at", "TEXT", 1, None, 0), + ), + ITEMS_TABLE: ( + ("collection_id", "INTEGER", 1, None, 1), + ("root_key", "TEXT", 1, None, 2), + ("relative_path", "TEXT", 1, None, 3), + ("added_at", "TEXT", 1, None, 0), + ), +} +# (column, parent table, parent column, on delete) +_FOREIGN_KEYS = { + COLLECTIONS_TABLE: set(), + ITEMS_TABLE: {("collection_id", COLLECTIONS_TABLE, "id", "CASCADE")}, +} + + +class CollectionInvalid(ValueError): + """A name, id or target breaks the write rules (a 400).""" + + +def normalize_collection_name(name: object) -> str: + """The stored spelling of *name*: trimmed, whitespace collapsed, checked. + + Raises ``CollectionInvalid`` unless the result is 1 to 100 characters + without control characters. + """ + if not isinstance(name, str): + raise CollectionInvalid("Name must be a string") + collapsed = _WHITESPACE.sub(" ", name).strip() + if not 1 <= len(collapsed) <= MAX_NAME_LENGTH: + raise CollectionInvalid(f"Name must be 1 to {MAX_NAME_LENGTH} characters") + if _CONTROL.search(collapsed): + raise CollectionInvalid("Name contains control characters") + return collapsed + + +def collection_name_key(name: str) -> str: + """The uniqueness key of a name: NFC of the case-folded NFC spelling.""" + return unicodedata.normalize("NFC", unicodedata.normalize("NFC", name).casefold()) + + +def _compact(sql: str) -> str: + return re.sub(r"\s+", "", sql).lower() + + +def _expected_catalog() -> set: + objects = {("table", name, _compact(sql)) for name, sql in _TABLE_DDL.items()} + objects.add(("index", FILE_INDEX, _compact(_INDEX_DDL))) + objects |= {("index", f"sqlite_autoindex_{table}_1", "") for table in TABLES} + return objects + + +def _catalog(conn: sqlite3.Connection) -> list: + # sqlite_sequence (AUTOINCREMENT) is shared with other tables and ignored. + return conn.execute( + "SELECT type,name,sql FROM sqlite_master " + "WHERE name IN (?,?,?) OR tbl_name IN (?,?)", + (*TABLES, FILE_INDEX, *TABLES), + ).fetchall() + + +def _structure_error(conn: sqlite3.Connection) -> Optional[str]: + for table, expected in _COLUMNS.items(): + columns = tuple( + (row[1], row[2].upper(), row[3], row[4], row[5]) + for row in conn.execute(f"PRAGMA table_info({table})") + ) + if columns != expected: + return f"Image collections table {table} has invalid columns" + foreign_keys = { + (row[3], row[2], row[4], row[6]) + for row in conn.execute(f"PRAGMA foreign_key_list({table})") + } + if foreign_keys != _FOREIGN_KEYS[table]: + return f"Image collections table {table} has invalid foreign keys" + return None + + +def _decoded_too_long(conn, sql: str, limit: int) -> bool: + """Whether a byte-bounded candidate row decodes to over *limit* characters.""" + for row in conn.execute(sql): + for raw in row: + if len(bytes(raw).decode("utf-8", "replace")) > limit: + return True + return False + + +def _value_error(conn: sqlite3.Connection) -> Optional[str]: + # Only this module's constants are interpolated into the SQL below. Bytes + # bound characters from above (UTF-8 is at most 4 bytes each); the few + # candidates over the character bound are measured decoded in Python. + if conn.execute( + f"SELECT 1 FROM {COLLECTIONS_TABLE} WHERE " # nosec B608 + "typeof(id) != 'integer' OR typeof(name) != 'text' " + "OR typeof(name_key) != 'text' OR typeof(created_at) != 'text' " + f"OR length(CAST(name AS BLOB)) > {MAX_NAME_LENGTH * 4} " + f"OR length(CAST(created_at AS BLOB)) > {_TIMESTAMP_BYTES} LIMIT 1" + ).fetchone(): + return "Image collections contains invalid values" + for name, name_key in conn.execute( + f"SELECT name, name_key FROM {COLLECTIONS_TABLE}" # nosec B608 + ): + if ( + not 1 <= len(name) <= MAX_NAME_LENGTH + or _CONTROL.search(name) + or name != name.strip() + or name_key != collection_name_key(name) + ): + return "Image collections contains invalid names" + limit = MAX_KEY_LENGTH * 4 + if conn.execute( + f"SELECT 1 FROM {ITEMS_TABLE} WHERE " # nosec B608 + "typeof(collection_id) != 'integer' OR typeof(root_key) != 'text' " + "OR typeof(relative_path) != 'text' OR typeof(added_at) != 'text' " + f"OR length(CAST(root_key AS BLOB)) > {limit} " + f"OR length(CAST(relative_path AS BLOB)) > {limit} " + f"OR length(CAST(added_at AS BLOB)) > {_TIMESTAMP_BYTES} LIMIT 1" + ).fetchone() or _decoded_too_long( + conn, + f"SELECT CAST(root_key AS BLOB), CAST(relative_path AS BLOB) " # nosec B608 + f"FROM {ITEMS_TABLE} WHERE length(CAST(root_key AS BLOB)) > {MAX_KEY_LENGTH} " + f"OR length(CAST(relative_path AS BLOB)) > {MAX_KEY_LENGTH}", + MAX_KEY_LENGTH, + ): + return "Image collections contains invalid members" + if conn.execute(f"PRAGMA foreign_key_check({ITEMS_TABLE})").fetchone(): + return "Image collections contains members of a missing collection" + return None + + +def image_collections_schema_error( + conn: sqlite3.Connection, + *, + allow_missing: bool = True, + check_values: bool = True, +) -> Optional[str]: + """Reject partial or unsupported collections schemas and invalid rows. + + Reads only: nothing is ever repaired. A database without either table is + acceptable for backups that predate them unless ``allow_missing`` is + false. ``check_values=False`` (live startup) checks the structure only. + """ + objects = _catalog(conn) + if not objects: + return None if allow_missing else "Missing image collections schema" + present = {name for kind, name, _sql in objects if kind == "table"} + if present != set(TABLES): + return "Incomplete image collections schema" + actual = {(kind, name, _compact(sql or "")) for kind, name, sql in objects} + if actual != _expected_catalog(): + return "Image collections has an invalid schema definition" + error = _structure_error(conn) + if error or not check_values: + return error + return _value_error(conn) + + +def migrate_image_collections(conn: sqlite3.Connection) -> None: + """Install both tables atomically; validate before and after creating them.""" + conn.execute("SAVEPOINT image_collections_schema") + try: + error = image_collections_schema_error(conn, check_values=False) + if error: + raise sqlite3.DatabaseError(error) + for sql in _TABLE_DDL.values(): + conn.execute(sql.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1)) + conn.execute(_INDEX_DDL.replace("INDEX", "INDEX IF NOT EXISTS", 1)) + error = image_collections_schema_error( + conn, allow_missing=False, check_values=False + ) + if error: + raise sqlite3.DatabaseError(error) + conn.execute("RELEASE image_collections_schema") + except BaseException: + conn.execute("ROLLBACK TO image_collections_schema") + conn.execute("RELEASE image_collections_schema") + raise + + +class CollectionMissing(LookupError): + """No collection has the id (a 404 ``collection-not-found``).""" + + +class CollectionNameTaken(ValueError): + """Another collection already has the name's key (a 409 ``name-taken``).""" + + def __init__(self, collection_id: int): + super().__init__("Another collection has this name") + self.collection_id = collection_id + + +class CollectionFilesMissing(LookupError): + """A new name with nothing to insert (a 404 ``files-not-found``).""" + + +LIST_SQL = f"SELECT id, name FROM {COLLECTIONS_TABLE} ORDER BY name_key, id" +GET_SQL = f"SELECT id, name FROM {COLLECTIONS_TABLE} WHERE id = ?" +BY_KEY_SQL = f"SELECT id, name FROM {COLLECTIONS_TABLE} WHERE name_key = ?" +MEMBERS_SQL = ( + f"SELECT relative_path FROM {ITEMS_TABLE} " + "WHERE collection_id = ? AND root_key = ?" +) +ROOT_MEMBERS_SQL = ( + f"SELECT collection_id, relative_path FROM {ITEMS_TABLE} WHERE root_key = ?" +) +COUNTS_SQL = ( + f"SELECT c.id, (SELECT COUNT(*) FROM {ITEMS_TABLE} i " + f"WHERE i.collection_id = c.id) FROM {COLLECTIONS_TABLE} c" +) +FILE_SQL = ( + f"SELECT collection_id FROM {ITEMS_TABLE} " + "WHERE root_key = ? AND relative_path = ? ORDER BY collection_id" +) +# Every read, for the EXPLAIN QUERY PLAN check: none scans the items table and +# only the two whole-list reads (LIST_SQL, COUNTS_SQL) scan the collections. +READ_SQL = ( + LIST_SQL, + GET_SQL, + BY_KEY_SQL, + MEMBERS_SQL, + ROOT_MEMBERS_SQL, + COUNTS_SQL, + FILE_SQL, +) +_INSERT_ITEM_SQL = f"INSERT OR IGNORE INTO {ITEMS_TABLE} VALUES (?, ?, ?, ?)" +_DELETE_ITEM_SQL = ( + f"DELETE FROM {ITEMS_TABLE} " + "WHERE collection_id = ? AND root_key = ? AND relative_path = ?" +) + + +def _collection_id(value: object) -> int: + if type(value) is not int or not 1 <= value <= 999_999_999: + raise CollectionInvalid("Collection id must be a positive integer") + return value + + +def _items(items: object) -> list: + """Distinct ``(root key, relative key)`` pairs, in order.""" + if not isinstance(items, (list, tuple)): + raise CollectionInvalid("Items must be a list") + found = [] + for item in items: + if ( + not isinstance(item, (list, tuple)) + or len(item) != 2 + or not all( + isinstance(part, str) and 1 <= len(part) <= MAX_KEY_LENGTH + for part in item + ) + ): + raise CollectionInvalid("Each item is a root key and a relative key") + found.append((item[0], item[1])) + return list(dict.fromkeys(found)) + + +def _row(row) -> dict: + return {"id": row[0], "name": row[1]} + + +class ImageCollectionOperations: + """PromptDatabase mixin: collections and their members. + + Each write runs in its own ``BEGIN IMMEDIATE`` transaction + (``_pair_transaction``), so two tabs creating one name are serialized. + """ + + def list_collections(self) -> list[dict]: + """Every collection ``{id, name}`` in name-key order, then id.""" + rows = self.model.get_connection().execute(LIST_SQL) + return [_row(row) for row in rows] + + def get_collection(self, collection_id: int) -> Optional[dict]: + """``{id, name}`` of one collection, or None.""" + row = ( + self.model.get_connection() + .execute(GET_SQL, (_collection_id(collection_id),)) + .fetchone() + ) + return _row(row) if row else None + + def collection_exists(self, collection_id: int) -> bool: + return self.get_collection(collection_id) is not None + + def collection_members(self, collection_id: int, root_key: str) -> set[str]: + """The relative keys stored for one collection under one source.""" + rows = self.model.get_connection().execute( + MEMBERS_SQL, (_collection_id(collection_id), root_key) + ) + return {path for (path,) in rows} + + def root_collection_members(self, root_key: str) -> dict[int, set[str]]: + """``{collection id: relative keys}`` stored under one source.""" + found: dict[int, set[str]] = {} + for collection_id, path in self.model.get_connection().execute( + ROOT_MEMBERS_SQL, (root_key,) + ): + found.setdefault(collection_id, set()).add(path) + return found + + def collection_counts(self) -> dict[int, int]: + """``{collection id: stored rows}`` for every collection.""" + rows = self.model.get_connection().execute(COUNTS_SQL) + return {collection_id: count for collection_id, count in rows} + + def file_collections(self, root_key: str, relative_path: str) -> list[int]: + """Ids of the collections holding one file.""" + rows = self.model.get_connection().execute(FILE_SQL, (root_key, relative_path)) + return [collection_id for (collection_id,) in rows] + + def add_to_collection( + self, *, items: object, collection_id: object = None, name: object = None + ) -> dict: + """Add files to an existing id or a name; create the name when needed. + + A name whose key exists uses that collection; otherwise a collection + is created only when at least one item is inserted. Returns + ``{collection, created, added, already}``. + """ + from .models import utc_now_iso # models imports this module's schema + + if (collection_id is None) == (name is None): + raise CollectionInvalid("Give either a collection id or a name") + pairs = _items(items) + if name is not None: + spelled = normalize_collection_name(name) + else: + collection_id = _collection_id(collection_id) + now = utc_now_iso() + with self._pair_transaction() as conn: + created = False + if name is not None: + row = conn.execute( + BY_KEY_SQL, (collection_name_key(spelled),) + ).fetchone() + if row is None: + if not pairs: + raise CollectionFilesMissing("Nothing to add") + cursor = conn.execute( + f"INSERT INTO {COLLECTIONS_TABLE} " # nosec B608 + "(name, name_key, created_at) VALUES (?, ?, ?)", + (spelled, collection_name_key(spelled), now), + ) + row, created = (cursor.lastrowid, spelled), True + else: + row = conn.execute(GET_SQL, (collection_id,)).fetchone() + if row is None: + raise CollectionMissing("No such collection") + collection = _row(row) + added = sum( + conn.execute( + _INSERT_ITEM_SQL, (collection["id"], root, path, now) + ).rowcount + for root, path in pairs + ) + return { + "collection": collection, + "created": created, + "added": added, + "already": len(pairs) - added, + } + + def remove_from_collection(self, collection_id: object, items: object) -> dict: + """Remove files from one collection; ``{collection, removed}``.""" + collection_id = _collection_id(collection_id) + pairs = _items(items) + with self._pair_transaction() as conn: + row = conn.execute(GET_SQL, (collection_id,)).fetchone() + if row is None: + raise CollectionMissing("No such collection") + removed = sum( + conn.execute(_DELETE_ITEM_SQL, (collection_id, root, path)).rowcount + for root, path in pairs + ) + return {"collection": _row(row), "removed": removed} + + def rename_collection(self, collection_id: object, name: object) -> dict: + """The renamed ``{id, name}``; a case variant of its own name is fine.""" + collection_id = _collection_id(collection_id) + spelled = normalize_collection_name(name) + key = collection_name_key(spelled) + with self._pair_transaction() as conn: + if conn.execute(GET_SQL, (collection_id,)).fetchone() is None: + raise CollectionMissing("No such collection") + taken = conn.execute(BY_KEY_SQL, (key,)).fetchone() + if taken is not None and taken[0] != collection_id: + raise CollectionNameTaken(taken[0]) + conn.execute( + f"UPDATE {COLLECTIONS_TABLE} SET name = ?, name_key = ? " # nosec B608 + "WHERE id = ?", + (spelled, key, collection_id), + ) + return {"id": collection_id, "name": spelled} + + def delete_collection(self, collection_id: object) -> bool: + """Delete one collection and its members (never a file). True.""" + collection_id = _collection_id(collection_id) + with self._pair_transaction() as conn: + cursor = conn.execute( + f"DELETE FROM {COLLECTIONS_TABLE} WHERE id = ?", # nosec B608 + (collection_id,), + ) + if cursor.rowcount == 0: + raise CollectionMissing("No such collection") + return True diff --git a/database/models.py b/database/models.py index dad9cf07..fe30685e 100644 --- a/database/models.py +++ b/database/models.py @@ -22,6 +22,10 @@ from .managed_connection import ( from . import restore_stage from .library_schema import library_schema_error, migrate_library_state from .image_favorites import image_favorites_schema_error, migrate_image_favorites +from .image_collections import ( + image_collections_schema_error, + migrate_image_collections, +) from .annotation_schema import annotation_schema_error, migrate_annotation_state from .autotag_schema import ( autotag_schema_error, @@ -337,6 +341,7 @@ class PromptModel: migrate_library_state(conn) migrate_image_favorites(conn) migrate_saved_views(conn) + migrate_image_collections(conn) # A review group without its annotation group is refused before the # annotation install could make it look valid (catalog only; the rows # are validated once, in migrate_review_state). @@ -1156,7 +1161,7 @@ class PromptModel: 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 + collections, 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 @@ -1225,6 +1230,11 @@ class PromptModel: ) if views_error: return False, views_error + collections_error = image_collections_schema_error( + conn, allow_missing=not require_library_state + ) + if collections_error: + return False, collections_error annotation_error = annotation_schema_error( conn, allow_missing=not require_library_state ) diff --git a/database/operations.py b/database/operations.py index 0f45097f..def1adeb 100644 --- a/database/operations.py +++ b/database/operations.py @@ -13,6 +13,7 @@ from .models import PromptModel, normalize_image_path, utc_now_iso from .pairs import PairDependencyError, PairOperations # noqa: F401 (re-exported) from .library import LibraryOperations, titled from .image_favorites import ImageFavoriteOperations +from .image_collections import ImageCollectionOperations from .annotations import AnnotationOperations from .review import ReviewOperations from .review_decisions import ReviewDecisionOperations @@ -359,6 +360,7 @@ class PromptDatabase( ExecutionOutputOperations, HistoryReadOperations, ImageFavoriteOperations, + ImageCollectionOperations, AnnotationOperations, ReviewOperations, ReviewDecisionOperations, diff --git a/database/restore_stage.py b/database/restore_stage.py index 75d8883a..704435bf 100644 --- a/database/restore_stage.py +++ b/database/restore_stage.py @@ -59,6 +59,7 @@ SUMMARY_KEYS = ( "image_links", "runs", "saved_views", + "collections", ) @@ -144,7 +145,8 @@ def library_summary(conn: sqlite3.Connection) -> Dict[str, int]: ``positives`` are the Library's prompts (``library_counts``'s ``positive_prompts``); ``pairs`` the active saved pairs and ``archived_pairs`` the archived ones; ``image_links`` the indexed image - rows; ``runs`` the recorded executions. + rows; ``runs`` the recorded executions; ``collections`` the image + collections (#331). """ from .library import _ELIGIBLE @@ -169,6 +171,7 @@ def library_summary(conn: sqlite3.Connection) -> Dict[str, int]: "image_links": _count(conn, "generated_images"), "runs": _count(conn, "execution_runs"), "saved_views": _count(conn, "saved_views"), + "collections": _count(conn, "image_collections"), } @@ -180,6 +183,7 @@ _FINGERPRINT_SQL = ( ("prompt_pairs", "SELECT MAX(id), MAX(updated_at) FROM prompt_pairs"), ("execution_runs", "SELECT COUNT(*), MAX(rowid) FROM execution_runs"), ("saved_views", "SELECT MAX(rowid) FROM saved_views"), + ("image_collections", "SELECT MAX(rowid) FROM image_collections"), ) @@ -190,9 +194,9 @@ def live_fingerprint(conn: sqlite3.Connection): tables a restore replaces, or when a prompt, pair, run or saved view is the newest edited (counts plus the newest row and edit of each table); the apply refuses with ``plan-changed`` then. It does not see every - edit (a tag or favorite change on an older row); those are kept in the - restore's safety backup. A read transaction the caller already has is - used and left open; one opened here is ended here. + edit (a tag, favorite or collection membership change on an older row); + those are kept in the restore's safety backup. A read transaction the + caller already has is used and left open; one opened here is ended here. """ import json diff --git a/py/api/__init__.py b/py/api/__init__.py index 485578c7..ed94ecea 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -48,6 +48,7 @@ from .review import ReviewRoutesMixin from .review_recording import ReviewRecordingMixin from .review_events import ReviewEventsMixin from .saved_views import SavedViewRoutesMixin +from .collections import CollectionRoutesMixin from .history import HistoryRoutesMixin from .images import ImageRoutesMixin from .gallery_listing import WALK_STOPPED, RootListing, WalkOutcome @@ -314,6 +315,7 @@ class PromptManagerAPI( ReviewRecordingMixin, ReviewEventsMixin, SavedViewRoutesMixin, + CollectionRoutesMixin, HistoryRoutesMixin, PromptRoutesMixin, ImageRoutesMixin, @@ -684,6 +686,7 @@ class PromptManagerAPI( self._register_review_routes(routes) self._register_review_event_routes(routes) self._register_saved_view_routes(routes) + self._register_collection_routes(routes) self._register_history_routes(routes) self._register_prompt_routes(routes) self._register_image_routes(routes) @@ -1107,196 +1110,3 @@ class PromptManagerAPI( result["parameters"][field] = metadata[field] return result - - def _extract_readable_prompt(self, parsed_data): - """Extract human-readable prompt text from ComfyUI/A1111 data.""" - - def safe_to_string(value): - if isinstance(value, str): - return value - elif isinstance(value, list): - return " ".join(str(item) for item in value if item) - elif value is not None: - return str(value) - return None - - # First check if we already extracted a positive prompt (A1111 format) - if parsed_data.get("positive_prompt"): - return parsed_data["positive_prompt"] - - # Check if prompt is already a string - if isinstance(parsed_data.get("prompt"), str): - return parsed_data["prompt"] - - # Check if prompt is a simple value that can be converted - if parsed_data.get("prompt") and not isinstance( - parsed_data.get("prompt"), dict - ): - return safe_to_string(parsed_data["prompt"]) - - prompt_data = parsed_data.get("prompt") - if isinstance(prompt_data, dict): - # Use the enhanced logic from parse-metadata.py - positive_prompt = self._extract_positive_prompt_from_comfyui_data( - prompt_data - ) - if positive_prompt: - return positive_prompt - - # Check workflow data if available - workflow_data = parsed_data.get("workflow") - if isinstance(workflow_data, dict): - positive_prompt = self._extract_positive_prompt_from_comfyui_data( - workflow_data - ) - if positive_prompt: - return positive_prompt - - # Check parameters for positive prompt - if parsed_data.get("parameters", {}).get("positive"): - return safe_to_string(parsed_data["parameters"]["positive"]) - - return None - - def _get_node_inputs(self, node): - """Safely get inputs from a node, handling both dict and list formats. - - Old format: inputs is a dict with direct key-value pairs - inputs = {"text": "my prompt", "seed": 123} - - New format: inputs is a list of connection objects - inputs = [ - {"name": "text", "type": "STRING", "link": null, - "widget": {"name": "text"}}, - {"name": "clip", "type": "CLIP", "link": 11} - ] - """ - if not isinstance(node, dict): - return {} - - inputs = node.get("inputs", {}) - - # If inputs is already a dict, return it - if isinstance(inputs, dict): - return inputs - - # If inputs is a list, convert to dict format - if isinstance(inputs, list): - inputs_dict = {} - for input_item in inputs: - if isinstance(input_item, dict) and "name" in input_item: - name = input_item["name"] - inputs_dict[name] = input_item - return inputs_dict - - return {} - - def _find_text_in_node(self, node): - """Try to find text content in a node using various strategies. - - Handles both old and new workflow formats. - """ - if not isinstance(node, dict): - return None - - # Strategy 1: Check normalized inputs for 'text' field - inputs = self._get_node_inputs(node) - if "text" in inputs and isinstance(inputs["text"], str): - return inputs["text"] - - # Strategy 2: For text encoder nodes, check widgets_values - class_type = node.get("class_type", node.get("type", "")) - text_encoder_types = [ - "CLIPTextEncode", - "CLIPTextEncodeSDXL", - "CLIPTextEncodeSDXLRefiner", - "CLIPTextEncodeFlux", - "PromptManager", - "PromptManagerText", - "BNK_CLIPTextEncoder", - "Text Encoder", - "CLIP Text Encode", - ] - - if any( - encoder_type.lower() in class_type.lower() - for encoder_type in text_encoder_types - ): - widgets_values = node.get("widgets_values", []) - if widgets_values and len(widgets_values) > 0: - if isinstance(widgets_values[0], str) and widgets_values[0].strip(): - return widgets_values[0] - - return None - - def _extract_positive_prompt_from_comfyui_data(self, data): - """Extract positive prompt from ComfyUI data. - - Handles both old and new formats. - """ - if not isinstance(data, dict): - return None - - # Build nodes dictionary - nodes_by_id = {} - if "nodes" in data: - # Handle nodes array format - for node in data["nodes"]: - if isinstance(node, dict): - nid = node.get("id") - if nid is not None: - nodes_by_id[nid] = node - else: - # Handle flat dictionary format (node_id -> node_data) - for nid_str, node in data.items(): - try: - nid = int(nid_str) - except (ValueError, TypeError): - nid = nid_str - if isinstance(node, dict): - if "id" in node: - nid = node["id"] - nodes_by_id[nid] = node - - if not nodes_by_id: - return None - - # First, try to find positive/negative connection pattern - pos_id = None - for node in nodes_by_id.values(): - if isinstance(node, dict): - inputs = self._get_node_inputs(node) - if "positive" in inputs and "negative" in inputs: - try: - pos_input = inputs["positive"] - if isinstance(pos_input, list) and len(pos_input) > 0: - pos_id = int(pos_input[0]) - break - except (ValueError, TypeError, IndexError): - continue - - # Get text from the positive node - if pos_id is not None and pos_id in nodes_by_id: - text_val = self._find_text_in_node(nodes_by_id[pos_id]) - if text_val: - return text_val - - # Fallback: find any text encoder node with text content - text_nodes = [] - for node in nodes_by_id.values(): - if isinstance(node, dict): - text_val = self._find_text_in_node(node) - if text_val: - # Try to determine if this is positive or negative - node_title = node.get("title", "").lower() - if "neg" not in node_title and "negative" not in node_title: - # Prioritize non-negative prompts - text_nodes.insert(0, text_val) - else: - text_nodes.append(text_val) - - # Return the first positive-looking prompt - if text_nodes: - return text_nodes[0] - - return None diff --git a/py/api/admin.py b/py/api/admin.py index 4ebf8eea..ec3f6db7 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -47,7 +47,7 @@ try: extract_generation_params, parse_page_json, ) - from ...utils.prompt_pair_resolver import resolve_prompt_pair + from ...utils.prompt_pair_resolver import resolve_prompt_pair, sampler_less_prompt from ...utils.review_capture import reviewable except ImportError: from database.pair_ingestion import blank_positive, ingest_resolved_pair @@ -62,7 +62,7 @@ except ImportError: extract_generation_params, parse_page_json, ) - from utils.prompt_pair_resolver import resolve_prompt_pair + from utils.prompt_pair_resolver import resolve_prompt_pair, sampler_less_prompt from utils.review_capture import reviewable @@ -1754,13 +1754,24 @@ class AdminRoutesMixin: canonical.get("parameters"), str ): return resolution["state"] + elif canonical.get("prompt"): + # No node takes a positive or negative input (a guider or wrapper + # graph): the one wired text, never a guess (#252). + found = sampler_less_prompt(canonical["prompt"]) + prompt_text = found["text"] + if prompt_text is None and isinstance(canonical.get("parameters"), str): + # A metadata saver also wrote A1111 text: its positive prompt. + parsed = self._parse_comfyui_prompt(metadata)["positive_prompt"] + prompt_text = parsed.strip() if parsed and parsed.strip() else None + if prompt_text is None: + linked = self._link_recorded_run( + media_file, canonical, detailed, history + ) + if linked is not None: + return linked + return "unresolved" if found["state"] == "ambiguous" else "found" else: - parsed = self._parse_comfyui_prompt(metadata) - if not (parsed.get("prompt") or parsed.get("parameters")): - return None - prompt_text = self._extract_readable_prompt(parsed) - if prompt_text is not None and not isinstance(prompt_text, str): - prompt_text = str(prompt_text) + prompt_text = self._parse_comfyui_prompt(metadata)["positive_prompt"] if not (prompt_text and prompt_text.strip()): return "found" prompt_text = prompt_text.strip() diff --git a/py/api/annotations.py b/py/api/annotations.py index 47f905e6..d9a27a0f 100644 --- a/py/api/annotations.py +++ b/py/api/annotations.py @@ -27,7 +27,12 @@ from .annotation_suggest import ( validate_options, ) from .autotag_routes import model_busy_response -from .images import _forbidden, _has_traversal, _unavailable_link, _verified_title + +# The target check is shared with the collection routes (#331); both names +# stay importable from here. +from .media_target import MAX_TARGET_PATH # noqa: F401 +from .media_target import media_target as _target +from .images import _unavailable_link, _verified_title from .prompts import bad_request from .. import model_gate @@ -40,7 +45,7 @@ try: ) from ...database.execution_history import subscribe_history_restore from ...database.library_commands import LibraryCommandError - from ...utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore + from ...utils.media_capture import CaptureError, CaptureStore except ImportError: from database.annotations import ( MAX_IMAGE_TAGS, @@ -50,15 +55,13 @@ except ImportError: ) from database.execution_history import subscribe_history_restore from database.library_commands import LibraryCommandError - from utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore + from utils.media_capture import CaptureError, CaptureStore # Request bodies of the annotation routes are small (a target plus a tag set). MAX_BODY_BYTES = 64 * 1024 -MAX_TARGET_PATH = 1024 CAPTURE_URL = "/prompt_manager/image-annotations/capture/" _PREVIEW_CHUNK = 64 * 1024 -_TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"}) _OBSERVE_FIELDS = frozenset({"target"}) _SAVE_FIELDS = frozenset( { @@ -188,37 +191,6 @@ def _capture_error(exc: CaptureError) -> web.Response: return _coded(exc.code, _MESSAGES.get(exc.code) or str(exc)) -def _target(data: object): - """``((index, scope, relative), None)`` for a target, else ``(None, error)``. - - Only the listing identity is accepted: an int root index, the root's - 64-hex scope and a bounded root-relative POSIX media path, never an - absolute, drive-qualified, backslash or ``..`` path. Nothing here touches - the filesystem or the database. - """ - if not isinstance(data, dict) or set(data) != _TARGET_FIELDS: - return None, bad_request("target must name root_index, root_scope and path") - index, scope, relative = data["root_index"], data["root_scope"], data["path"] - if type(index) is not int or index < 0: - return None, bad_request("Invalid gallery root") - if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None: - return None, bad_request("Invalid gallery root scope") - if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH: - return None, bad_request("path must be a relative media path") - parts = relative.split("/") - if ( - "\0" in relative - or "\\" in relative - or relative.startswith("/") - or _has_traversal(relative) - or any(part in ("", ".", "..") or ":" in part for part in parts) - ): - return None, _forbidden() - if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS: - return None, _forbidden("Only media files can be tagged") - return (index, scope, relative), None - - def _expires_at(seconds: float | None) -> str: moment = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta( seconds=max(0.0, seconds or 0.0) diff --git a/py/api/collections.py b/py/api/collections.py new file mode 100644 index 00000000..51fab44d --- /dev/null +++ b/py/api/collections.py @@ -0,0 +1,325 @@ +"""Collections (#331): named, hand-picked groups of images. + +Routes under ``/prompt_manager/collections``. A file is named by its listing +identity ``{root_index, root_scope, path}`` (``media_target.py``) and stored +under the favorites' keys (``_favorite_root_key``, ``_favorite_relative_key``), +which never leave the server. A stale scope (the sources were reordered) is a +409 ``source-changed`` before anything is written. Responses are +``{"success": bool, ...}``; an unexpected failure is a logged 500 with a +generic message, never the exception text. +""" + +import hmac +import re +from functools import wraps +from typing import Callable + +from aiohttp import web + +from .images import ( + _favorite_relative_key, + _favorite_root_key, + _forbidden, + _json_body_or_error, + collection_not_found, +) +from .media_target import media_target +from .prompts import bad_request + +try: + from ...database.image_collections import ( + CollectionFilesMissing, + CollectionInvalid, + CollectionMissing, + CollectionNameTaken, + normalize_collection_name, + ) + from ...database.image_favorites import MAX_KEY_LENGTH +except ImportError: + from database.image_collections import ( + CollectionFilesMissing, + CollectionInvalid, + CollectionMissing, + CollectionNameTaken, + normalize_collection_name, + ) + from database.image_favorites import MAX_KEY_LENGTH + +BASE = "/prompt_manager/collections" +# The most files one add or remove names (the bulk tagging cap, MC12). +MAX_ITEMS = 500 +MAX_ID = 999_999_999 +_PATH_ID = re.compile(r"[0-9]{1,9}") +_ROOT_INDEX = re.compile(r"[0-9]{1,9}") +_ADD_TARGETS = (frozenset({"collection_id", "items"}), frozenset({"name", "items"})) +_REMOVE_FIELDS = frozenset({"collection_id", "items"}) +_RENAME_FIELDS = frozenset({"name"}) +_MEDIA_MESSAGE = "Only media files can be added to a collection" +READ_FAILED = "Collections could not be read." +CHANGE_FAILED = "The collection could not be changed." +SOURCES_CHANGED = "The sources changed. Refresh the images and try again." +FILES_NOT_FOUND = ( + "None of the selected images could be found. Refresh the images and try again." +) +NAME_TAKEN = "Another collection already has this name." + + +class _Refused(Exception): + """A request refused before any write; carries its response.""" + + def __init__(self, response: web.Response): + super().__init__(response.status) + self.response = response + + +def _coded(status: int, code: str, message: str, **extra) -> web.Response: + return web.json_response( + {"success": False, "code": code, "error": message, **extra}, + status=status, + headers={"Cache-Control": "no-store"}, + ) + + +def _collection_response(failure: str) -> Callable: + """Map the database's refusals to their status; anything else to a 500.""" + + def decorate(handler: Callable) -> Callable: + @wraps(handler) + async def wrapped(self, request: web.Request) -> web.Response: + try: + return await handler(self, request) + except _Refused as refused: + return refused.response + except CollectionInvalid as exc: + return bad_request(str(exc)) + except CollectionMissing: + return collection_not_found() + except CollectionFilesMissing: + return _coded(404, "files-not-found", FILES_NOT_FOUND) + except CollectionNameTaken as exc: + return _coded( + 409, "name-taken", NAME_TAKEN, collection_id=exc.collection_id + ) + except Exception: + self.logger.exception("Collection request failed") + return web.json_response( + {"success": False, "error": failure}, status=500 + ) + + return wrapped + + return decorate + + +def _path_id(request: web.Request) -> int: + raw = request.match_info["collection_id"] + if not raw.isascii() or _PATH_ID.fullmatch(raw) is None or int(raw) < 1: + raise _Refused(bad_request("Invalid collection id")) + return int(raw) + + +def _body_id(value: object) -> int: + if type(value) is not int or not 1 <= value <= MAX_ID: + raise _Refused(bad_request("Invalid collection id")) + return value + + +async def _body(request: web.Request, allowed: tuple) -> dict: + data, error = await _json_body_or_error(request) + if error is not None: + raise _Refused(error) + if frozenset(data) not in allowed: + names = " or ".join(", ".join(sorted(fields)) for fields in allowed) + raise _Refused(bad_request(f"Request must contain exactly {names}")) + return data + + +class CollectionRoutesMixin: + """Mixin providing the collection routes (#331).""" + + def _register_collection_routes(self, routes: web.RouteTableDef) -> None: + # The fixed paths go first: "/items" must never match "/{collection_id}". + routes.get(BASE)(self.list_collections) + routes.post(BASE + "/items")(self.add_collection_items) + routes.delete(BASE + "/items")(self.remove_collection_items) + routes.get(BASE + "/membership")(self.collection_membership) + routes.patch(BASE + "/{collection_id}")(self.rename_collection) + routes.delete(BASE + "/{collection_id}")(self.delete_collection) + + # -- reads ------------------------------------------------------------- + + @_collection_response(READ_FAILED) + async def list_collections(self, request: web.Request) -> web.Response: + """Every collection with its listed ``count`` and its ``missing`` rows. + + ``count`` is the members listed under every configured source (one + primary-key range read per source, matched against its cached + listing); ``missing`` the stored rows beyond it. When a scan or a read + fails both are null for every collection and the names still return. + """ + collections = await self._run_in_executor(self.db.list_collections) + counts = None + if collections: + try: + counts = await self._collection_counts() + except Exception: + self.logger.warning( + "Collection counts could not be read", exc_info=True + ) + rows = [] + for collection in collections: + listed, stored = (None, None) if counts is None else counts + count = None if listed is None else listed.get(collection["id"], 0) + missing = None if stored is None else stored.get(collection["id"], 0) + rows.append( + { + **collection, + "count": count, + "missing": None if count is None else missing - count, + } + ) + return web.json_response({"success": True, "collections": rows}) + + async def _collection_counts(self) -> tuple[dict, dict]: + """``({id: listed members}, {id: stored rows})`` over every source.""" + listed: dict[int, int] = {} + seen = set() + for root in self._get_all_output_dirs(): + root_key = _favorite_root_key(root) + if root_key in seen: # one folder configured twice counts once + continue + seen.add(root_key) + files = await self._get_gallery_files(root) + listing = await self._run_in_executor(self._favorite_listing, root, files) + members = await self._run_in_executor( + self.db.root_collection_members, root_key + ) + for collection_id, keys in members.items(): + found = sum(1 for key in keys if key in listing) + listed[collection_id] = listed.get(collection_id, 0) + found + stored = await self._run_in_executor(self.db.collection_counts) + return listed, stored + + @_collection_response(READ_FAILED) + async def collection_membership(self, request: web.Request) -> web.Response: + """GET ``?root_index&root_scope&path``: the ids holding that file.""" + query = request.query + if set(query) != {"root_index", "root_scope", "path"}: + return bad_request("Query must name root_index, root_scope and path") + raw_index = query["root_index"] + if not raw_index.isascii() or _ROOT_INDEX.fullmatch(raw_index) is None: + return bad_request("Invalid gallery root") + roots = list(self._get_all_output_dirs()) + root_key, key = self._collection_key( + { + "root_index": int(raw_index), + "root_scope": query["root_scope"], + "path": query["path"], + }, + roots, + ) + ids = await self._run_in_executor(self.db.file_collections, root_key, key) + return web.json_response({"success": True, "collection_ids": ids}) + + # -- writes ------------------------------------------------------------ + + @_collection_response(CHANGE_FAILED) + async def add_collection_items(self, request: web.Request) -> web.Response: + """POST ``{collection_id | name, items}``: add the listed files. + + Only files the source's listing holds are added; the rest count in + ``not_found``. A new name with nothing to add creates nothing (404 + ``files-not-found``). + """ + data = await _body(request, _ADD_TARGETS) + target = {} + if "name" in data: + target["name"] = normalize_collection_name(data["name"]) + else: + target["collection_id"] = _body_id(data["collection_id"]) + roots = list(self._get_all_output_dirs()) + keys = self._collection_items(data["items"], roots) + found, not_found = [], 0 + listings = {} + for index, root_key, key in keys: + if index not in listings: + root = roots[index] + files = await self._get_gallery_files(root) + listings[index] = await self._run_in_executor( + self._favorite_listing, root, files + ) + if key in listings[index]: + found.append((root_key, key)) + else: + not_found += 1 + result = await self._run_in_executor( + lambda: self.db.add_to_collection(items=found, **target) + ) + return web.json_response({"success": True, **result, "not_found": not_found}) + + @_collection_response(CHANGE_FAILED) + async def remove_collection_items(self, request: web.Request) -> web.Response: + """DELETE ``{collection_id, items}``: listed or not, every named file.""" + data = await _body(request, (_REMOVE_FIELDS,)) + collection_id = _body_id(data["collection_id"]) + roots = list(self._get_all_output_dirs()) + keys = self._collection_items(data["items"], roots) + result = await self._run_in_executor( + self.db.remove_from_collection, + collection_id, + [(root_key, key) for _index, root_key, key in keys], + ) + return web.json_response({"success": True, **result}) + + @_collection_response(CHANGE_FAILED) + async def rename_collection(self, request: web.Request) -> web.Response: + """PATCH ``{name}``: a case variant of its own name is allowed.""" + collection_id = _path_id(request) + data = await _body(request, (_RENAME_FIELDS,)) + name = normalize_collection_name(data["name"]) + collection = await self._run_in_executor( + self.db.rename_collection, collection_id, name + ) + return web.json_response({"success": True, "collection": collection}) + + @_collection_response(CHANGE_FAILED) + async def delete_collection(self, request: web.Request) -> web.Response: + """DELETE: the collection and its members; files are never touched.""" + collection_id = _path_id(request) + await self._run_in_executor(self.db.delete_collection, collection_id) + return web.json_response({"success": True, "deleted": True}) + + # -- request steps ----------------------------------------------------- + + def _collection_items(self, items: object, roots: list) -> list[tuple]: + """``(index, root key, relative key)`` per distinct file, in order. + + Every item is checked before anything is written: 400 for a malformed + shape, 403 for an unsafe path, 409 for a stale scope. Two spellings + that fold to one key (Windows) count once. + """ + if not isinstance(items, list) or not 1 <= len(items) <= MAX_ITEMS: + raise _Refused(bad_request(f"items must list 1 to {MAX_ITEMS} images")) + keys = {} + for item in items: + index = item.get("root_index") if isinstance(item, dict) else None + root_key, key = self._collection_key(item, roots) + keys.setdefault((root_key, key), index) + return [(index, root_key, key) for (root_key, key), index in keys.items()] + + def _collection_key(self, item: object, roots: list) -> tuple[str, str]: + """The stored keys of one target, or raise its refusal.""" + target, error = media_target(item, media_message=_MEDIA_MESSAGE) + if error is not None: + raise _Refused(error) + index, scope, relative = target + if index >= len(roots) or not hmac.compare_digest( + scope, self._gallery_root_scope(roots[index], index) + ): + raise _Refused(_coded(409, "source-changed", SOURCES_CHANGED)) + root_key = _favorite_root_key(roots[index]) + key = _favorite_relative_key(relative) + if max(len(root_key), len(key)) > MAX_KEY_LENGTH: + # A stored key past the schema bound would fail the next restore. + raise _Refused(_forbidden("Path is too long to add to a collection")) + return root_key, key diff --git a/py/api/duplicate_files.py b/py/api/duplicate_files.py index 6edf265d..9cf101f5 100644 --- a/py/api/duplicate_files.py +++ b/py/api/duplicate_files.py @@ -588,7 +588,9 @@ class DuplicateFilesMixin: job.update(total=len(targets), status_line="Removing duplicate files") spellings = await self._run_in_executor(self._source_spellings) removed, failure, keepers = 0, None, {} - totals = dict.fromkeys(("moved", "merged", "favorites", "tags", "review"), 0) + totals = dict.fromkeys( + ("moved", "merged", "favorites", "tags", "review", "collections"), 0 + ) for done, target in enumerate(targets, start=1): if failure is not None: self._record_target(job, "kept", target, NOT_ATTEMPTED) @@ -631,6 +633,11 @@ class DuplicateFilesMixin: job.log_line( f"Moved {_plural(totals['favorites'], 'favorite')} to the kept copies" ) + if totals["collections"]: + job.log_line( + f"Moved {_plural(totals['collections'], 'collection membership')} " + "to the kept copies" + ) if totals["tags"]: job.log_line( f"Added {_plural(totals['tags'], 'image tag')} to the kept copies" diff --git a/py/api/duplicate_spellings.py b/py/api/duplicate_spellings.py index 31e0f0b9..3653e5de 100644 --- a/py/api/duplicate_spellings.py +++ b/py/api/duplicate_spellings.py @@ -95,7 +95,8 @@ def copy_move(target, kept_stat, spellings): A record found under a folder spelling moves to the kept copy under the same spelling when that folder holds it too, else to its real - path. Its favorite moves under every listed root holding the file. + path. Its favorite and its collection memberships move under every + listed root holding the file. """ removed = spellings.of(target["abs"]) kept = spellings.of(target["keeper_abs"]) diff --git a/py/api/image_tag_jobs.py b/py/api/image_tag_jobs.py index 1d611745..66f148f5 100644 --- a/py/api/image_tag_jobs.py +++ b/py/api/image_tag_jobs.py @@ -58,7 +58,9 @@ from .annotations import ( from .autotag_routes import BATCH_CANCELLED, MODEL_BUSY_ERROR, model_busy_response from .images import ( MAX_LISTING_SELECTION, + CollectionMissing, ListingSourceChanged, + collection_not_found, _current_fingerprint, _listing_posix, _public_path, @@ -685,6 +687,8 @@ class ImageTagJobsMixin: selection = await self.resolve_listing_filter(spec.filter, cap) except ListingSourceChanged: return None, _coded("source-changed") + except CollectionMissing: + return None, collection_not_found() except ValueError as exc: return None, bad_request(str(exc)) except Exception: diff --git a/py/api/images.py b/py/api/images.py index 439a2800..10e70070 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -18,7 +18,7 @@ import urllib.parse from bisect import bisect_left, bisect_right, insort from collections.abc import Mapping from pathlib import Path -from typing import Callable, NamedTuple +from typing import Callable, Iterable, NamedTuple from aiohttp import web from PIL import Image, UnidentifiedImageError @@ -58,6 +58,7 @@ from .inspection_service import ( ) try: + from ...database.image_collections import CollectionMissing from ...database.image_favorites import MAX_KEY_LENGTH from ...database.library import display_title from ...database.models import normalize_image_path @@ -78,6 +79,7 @@ try: reveal_in_file_manager, ) except ImportError: + from database.image_collections import CollectionMissing from database.image_favorites import MAX_KEY_LENGTH from database.library import display_title from database.models import normalize_image_path @@ -104,6 +106,21 @@ IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS _CARD_RECORDS = ("prompt", "workflow") _PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" _LISTING_FAILED = "The image list could not be filtered. Check server logs for details." +_COLLECTION_ID = re.compile(r"[1-9][0-9]{0,8}") +COLLECTION_NOT_FOUND = "That collection no longer exists." + + +def collection_not_found() -> web.Response: + """The 404 a write or a listing answers for an unknown collection (#331).""" + return web.json_response( + { + "success": False, + "code": "collection-not-found", + "error": COLLECTION_NOT_FOUND, + }, + status=404, + headers={"Cache-Control": "no-store"}, + ) def _card_header(path: Path) -> tuple: @@ -1056,8 +1073,9 @@ def _workspace_filters(query: Mapping) -> dict: the linked positive's text, title and tag names. ``category`` is exact, ``tags`` is a comma list that must all be present, ``sort=rating`` ranks by the linked positive, and ``favorites=1`` keeps only favorites (the - Favorites page, #249 D51). Raises ValueError for any unknown or oversized - value. + Favorites page, #249 D51). ``collection=`` (1 to 9 digits) keeps only + that collection's members (#331); it wins over ``favorites`` (a page never + sets both). Raises ValueError for any unknown or oversized value. """ text = query.get("q", "").strip() folded = text.casefold() @@ -1071,8 +1089,10 @@ def _workspace_filters(query: Mapping) -> dict: {tag.strip() for tag in query.get("tags", "").split(",") if tag.strip()} ) favorites = query.get("favorites", "0") + collection = query.get("collection", "") if ( favorites not in ("0", "1") + or (collection != "" and _COLLECTION_ID.fullmatch(collection) is None) or len(text) > 500 or len(folded) > 500 or media not in ("all", "image", "video") @@ -1095,7 +1115,8 @@ def _workspace_filters(query: Mapping) -> dict: "path": exact.casefold(), "category": category, "tags": tags, - "favorites": favorites == "1", + "favorites": favorites == "1" and not collection, + "collection": int(collection) if collection else None, } @@ -1113,6 +1134,8 @@ LISTING_FILTER_KEYS = frozenset( "media", "sort", "favorites", + # A collection page's filter (#331): bulk actions act on its members. + "collection", } ) MAX_LISTING_FILTER_VALUE = 2048 @@ -1169,7 +1192,7 @@ def _admitted_entries( entries: list, filters: dict, links: dict | None, - favorites: frozenset | None = None, + members: frozenset | None = None, image_tags: dict | None = None, ): """Yield ``(position, entry, link match)`` for each entry the filters admit. @@ -1177,11 +1200,11 @@ def _admitted_entries( *links* maps normalized paths to (hit, rating, positive id) for every unambiguously linked image (see PromptDatabase.get_image_link_matches); a category or tag filter admits only those, and ``q`` admits a filename match - or a linked-positive hit. *favorites*, when given, admits only those listed - paths (``favorites=1``). *image_tags* maps normalized paths to their - _ImageTagMatch: a tag filter also admits a file whose current image tags - (with its linked positive's tags) carry every tag, and ``q`` also admits an - image tag name hit. + or a linked-positive hit. *members*, when given, admits only those listed + paths (the favorites of ``favorites=1`` or a ``collection``'s members). + *image_tags* maps normalized paths to their _ImageTagMatch: a tag filter + also admits a file whose current image tags (with its linked positive's + tags) carry every tag, and ``q`` also admits an image tag name hit. """ query, media, folder = filters["query"], filters["media"], filters["folder"] members_only = filters["category"] is not None or bool(filters["tags"]) @@ -1197,7 +1220,7 @@ def _admitted_entries( continue if media != "all" and (media == "video") != entry.video: continue - if favorites is not None and entry.path not in favorites: + if members is not None and entry.path not in members: continue match = None if links is None else links.get(entry.normalized) image = None if image_tags is None else image_tags.get(entry.normalized) @@ -1219,7 +1242,7 @@ def _filtered_page( size: int, filters: dict, links: dict | None, - favorites: frozenset | None = None, + members: frozenset | None = None, image_tags: dict | None = None, ) -> tuple[int, list, int | None]: """Count the *entries* the filters admit and take *size* from *start* on. @@ -1230,7 +1253,7 @@ def _filtered_page( """ total, page, positives = 0, [], set() for position, entry, match in _admitted_entries( - entries, filters, links, favorites, image_tags + entries, filters, links, members, image_tags ): total += 1 if match is not None: @@ -1934,6 +1957,7 @@ class ImageRoutesMixin: filters["category"], filters["tags"], filters["favorites"], + filters["collection"], ] ).encode() ).hexdigest() @@ -1944,17 +1968,19 @@ class ImageRoutesMixin: except ValueError as exc: return bad_request(str(exc)) - listings = [ - (index, root, await self._get_gallery_files(root)) - for index, root in enumerate(output_dirs) - if root_index is None or index == root_index - ] try: + if filters["collection"] is not None and not await self._run_in_executor( + self.db.collection_exists, filters["collection"] + ): + raise CollectionMissing("No such collection") + listings = await self._workspace_listings(output_dirs, root_index) page, total, has_more, next_cursor, positive_count = ( await self._run_in_executor( self._select_workspace_page, listings, filters, after, limit, scope ) ) + except CollectionMissing: + return collection_not_found() except Exception: # A filter that needs the prompt links or the image tags cannot be # served without them. Never echo the exception: it can carry SQL @@ -2215,9 +2241,11 @@ class ImageRoutesMixin: ) -> tuple: """The ordered entries and what admits them (blocking; run in executor). - Returns ``(entries, key, links, image_tags, favorites)``: the listing in - the filters' order, its keyset key, and the link, image-tag and favorite - maps _admitted_entries reads. One SELECT (get_image_link_matches) reads + Returns ``(entries, key, links, image_tags, members)``: the listing in + the filters' order, its keyset key, and the link, image-tag and member + maps _admitted_entries reads (members: a collection's, else the + favorites, else None). A collection deleted since the route checked it + raises CollectionMissing. One SELECT (get_image_link_matches) reads every linked positive when *read_links*. A failed read raises when the filters need it (category, tags, the rating sort) or *strict*; else the listing stays usable without links (a search then matches filenames @@ -2290,10 +2318,13 @@ class ImageRoutesMixin: def key(entry): return _workspace_key(entry, order) - favorites = ( - self._workspace_favorite_paths(listings) if filters["favorites"] else None + return ( + entries, + key, + links, + image_tags, + self._workspace_members(listings, filters), ) - return entries, key, links, image_tags, favorites def _select_workspace_page( self, @@ -2322,7 +2353,7 @@ class ImageRoutesMixin: or bool(filters["tags"]) or order == "rating" ) - entries, key, links, image_tags, favorites = self._workspace_admission( + entries, key, links, image_tags, members = self._workspace_admission( listings, filters, read_links=needed or after is None ) start = 0 @@ -2333,7 +2364,7 @@ class ImageRoutesMixin: else: start = bisect_right(entries, after, key=key) total, page, positive_count = _filtered_page( - entries, start, limit + 1, filters, links, favorites, image_tags + entries, start, limit + 1, filters, links, members, image_tags ) if not needed and after is not None: positive_count = _COUNT_OMITTED @@ -2355,12 +2386,12 @@ class ImageRoutesMixin: The links are read strictly: the positives must be the listing's own. """ - entries, _key, links, image_tags, favorites = self._workspace_admission( + entries, _key, links, image_tags, members = self._workspace_admission( listings, filters, read_links=True, strict=True ) targets, positives, unlinked, scopes = [], {}, 0, {} for _position, entry, match in _admitted_entries( - entries, filters, links, favorites, image_tags + entries, filters, links, members, image_tags ): if len(targets) == cap: return ListingSelection(targets, list(positives), unlinked, True) @@ -2382,9 +2413,11 @@ class ImageRoutesMixin: ) -> ListingSelection: """The images a parse_listing_filter result admits right now (#392). - Raises ValueError for a root outside the current roots, and + Raises ValueError for a root outside the current roots, ListingSourceChanged when the filter's ``root_scope`` no longer names - its root (the sources were renumbered since the listing). + its root (the sources were renumbered since the listing), and + CollectionMissing when its ``collection`` no longer exists (checked + before any scan and again after the members are read). """ output_dirs = list(self._get_all_output_dirs()) root_index = _workspace_root(params, output_dirs) @@ -2394,11 +2427,11 @@ class ImageRoutesMixin: raise ListingSourceChanged() raise ValueError("Invalid gallery root scope") filters = _workspace_filters(params) - listings = [ - (index, root, await self._get_gallery_files(root)) - for index, root in enumerate(output_dirs) - if root_index is None or index == root_index - ] + if filters["collection"] is not None and not await self._run_in_executor( + self.db.collection_exists, filters["collection"] + ): + raise CollectionMissing("No such collection") + listings = await self._workspace_listings(output_dirs, root_index) return await self._run_in_executor( self._select_workspace_matches, listings, filters, cap ) @@ -3583,25 +3616,58 @@ class ImageRoutesMixin: self._favorite_listing_cache[cache_key] = (files, listing) return listing - def _workspace_favorite_paths(self, listings: list[tuple]) -> frozenset: - """The listed paths of every favorite in the selected roots (blocking). + def _workspace_members(self, listings: list[tuple], filters: dict): + """The member set the filters ask for, or None (blocking).""" + collection = filters.get("collection") + if collection is not None: + members = self._workspace_member_paths( + listings, + lambda root_key: self.db.collection_members(collection, root_key), + ) + # Checked after the reads: a collection deleted before or while + # they ran is a 404, never an empty page passing for it. + if not self.db.collection_exists(collection): + raise CollectionMissing("No such collection") + return members + if filters["favorites"]: + return self._workspace_favorite_paths(listings) + return None - One primary-key read per root (get_root_favorites), matched against the - root's cached ``{relative key: listing spelling}``, so the work follows - the favorites rather than the files, and a key folds case exactly as - the toggle stored it (Windows). A favorite whose file is gone matches - nothing. A failed read raises: an unfiltered page must never pass for - the favorites. + def _workspace_favorite_paths(self, listings: list[tuple]) -> frozenset: + """The listed paths of every favorite in the selected roots (blocking).""" + return self._workspace_member_paths(listings, self.db.get_root_favorites) + + def _workspace_member_paths( + self, listings: list[tuple], keys_for_source: Callable[[str], Iterable[str]] + ) -> frozenset: + """The listed paths of the stored keys of each selected root (blocking). + + One read per root (``keys_for_source(root key)``: the favorites' or a + collection's primary-key range), matched against the root's cached + ``{relative key: listing spelling}``, so the work follows the members + rather than the files, and a key folds case exactly as it was stored + (Windows). A member whose file is gone matches nothing. A failed read + raises: an unfiltered page must never pass for the members. """ paths = set() for _index, root, files in listings: listing = self._favorite_listing(root, files) - for key in self.db.get_root_favorites(_favorite_root_key(root)): + for key in keys_for_source(_favorite_root_key(root)): listed = listing.get(key) if listed is not None: paths.add(root / listed) return frozenset(paths) + async def _workspace_listings( + self, output_dirs: list[Path], root_index: int | None + ) -> list[tuple]: + """``(index, root, scan)`` of every selected root (cached scans).""" + return [ + (index, root, await self._get_gallery_files(root)) + for index, root in enumerate(output_dirs) + if root_index is None or index == root_index + ] + def _workspace_favorites(self, page: list[tuple]) -> set[tuple[int, str]]: """``(root index, listing path)`` of the favorites among *page* entries.""" by_root: dict[tuple[int, Path], dict[str, str]] = {} diff --git a/py/api/media_target.py b/py/api/media_target.py new file mode 100644 index 00000000..ee1fcadf --- /dev/null +++ b/py/api/media_target.py @@ -0,0 +1,54 @@ +"""The listing identity of one media file, as a request names it. + +Shared by the image-annotation routes and the collection routes (#331): a +target is ``{root_index, root_scope, path}``, checked here without touching +the filesystem or the database. The scope is compared with the current roots +by the caller. +""" + +import os +import re + +from .images import _forbidden, _has_traversal +from .prompts import bad_request + +try: + from ...utils.media_capture import CAPTURE_EXTENSIONS +except ImportError: + from utils.media_capture import CAPTURE_EXTENSIONS + +MAX_TARGET_PATH = 1024 +TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"}) +_HEX64 = re.compile(r"[0-9a-f]{64}") +TAGGING_MEDIA_MESSAGE = "Only media files can be tagged" + + +def media_target(data: object, *, media_message: str = TAGGING_MEDIA_MESSAGE): + """``((index, scope, relative), None)`` for a target, else ``(None, error)``. + + Only the listing identity is accepted: an int root index, the root's + 64-hex scope and a bounded root-relative POSIX media path, never an + absolute, drive-qualified, backslash or ``..`` path (a 400 for a malformed + shape, a 403 for an unsafe or non-media path, carrying *media_message*). + """ + if not isinstance(data, dict) or set(data) != TARGET_FIELDS: + return None, bad_request("target must name root_index, root_scope and path") + index, scope, relative = data["root_index"], data["root_scope"], data["path"] + if type(index) is not int or index < 0: + return None, bad_request("Invalid gallery root") + if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None: + return None, bad_request("Invalid gallery root scope") + if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH: + return None, bad_request("path must be a relative media path") + parts = relative.split("/") + if ( + "\0" in relative + or "\\" in relative + or relative.startswith("/") + or _has_traversal(relative) + or any(part in ("", ".", "..") or ":" in part for part in parts) + ): + return None, _forbidden() + if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS: + return None, _forbidden(media_message) + return (index, scope, relative), None diff --git a/py/api/prompts.py b/py/api/prompts.py index e48296ea..8f5a1554 100644 --- a/py/api/prompts.py +++ b/py/api/prompts.py @@ -1227,7 +1227,9 @@ class PromptRoutesMixin: """ from .images import ( MAX_LISTING_SELECTION, + CollectionMissing, ListingSourceChanged, + collection_not_found, parse_listing_filter, ) @@ -1252,6 +1254,8 @@ class PromptRoutesMixin: }, status=409, ) + except CollectionMissing: + return collection_not_found() except ValueError as exc: return refused(str(exc)) if selection.truncated: diff --git a/tests/browser/server.py b/tests/browser/server.py index 71ea5669..95deff78 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -262,6 +262,8 @@ def create_app(root: Path) -> web.Application: with target.db.model.get_connection() as conn: for table in ( "saved_views", + "image_collection_items", + "image_collections", "image_favorites", "media_review_decisions", "media_review_items", diff --git a/tests/js/workspace-system-backups.test.js b/tests/js/workspace-system-backups.test.js index 05899351..3a1454ae 100644 --- a/tests/js/workspace-system-backups.test.js +++ b/tests/js/workspace-system-backups.test.js @@ -63,9 +63,9 @@ test('status line shows only server-reported backups (D103)', async () => { const PREVIEW = Object.freeze({ success: true, token: 'r'.repeat(43), expires_at: '2026-10-25T12:15:00+00:00', - summary: { prompts: 1234, positives: 900, pairs: 3, archived_pairs: 0, image_links: 56, runs: 1, saved_views: 2 }, + summary: { prompts: 1234, positives: 900, pairs: 3, archived_pairs: 0, image_links: 56, runs: 1, saved_views: 2, collections: 1 }, schema: { migrated: true, library_state: true }, integrity: 'ok', documents: { repaired: 0, cleared: 0, images: 0 }, - live: { prompts: 1, positives: 1, pairs: 0, archived_pairs: 0, image_links: 0, runs: 0, saved_views: 0 }, size: 4096, + live: { prompts: 1, positives: 1, pairs: 0, archived_pairs: 0, image_links: 0, runs: 0, saved_views: 0, collections: 0 }, size: 4096, safety_backup: { destination_label: 'prompts.db.backup_' }, }); diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index a5b85ba0..120c1e93 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -846,6 +846,14 @@ class TestSettingsReadMonitorThroughModule(AdminAPITestCase): self.assertEqual(data["settings"]["gallery_root_paths"], ["output"]) +def _wired_prompt(text): + """A sampler-less graph whose one encoder feeds a guider (#252).""" + return { + "1": {"class_type": "CLIPTextEncode", "inputs": {"text": text}}, + "2": {"class_type": "BasicGuider", "inputs": {"conditioning": ["1", 0]}}, + } + + def _raise(*args, **kwargs): raise RuntimeError("boom") @@ -2186,11 +2194,7 @@ class TestOutputScanInternals(AdminAPITestCase): def test_ingest_outcomes(self): f = Path("gen.png") - meta = { - "prompt": json.dumps( - {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}} - ) - } + meta = {"prompt": json.dumps(_wired_prompt("a cat"))} self.assertIsNone(self.api._ingest_scanned_file(f, {})) self.assertIsNone(self.api._ingest_scanned_file(f, {"prompt": "{}"})) self.assertEqual( @@ -2202,31 +2206,19 @@ class TestOutputScanInternals(AdminAPITestCase): def test_ingest_link_failures_are_tolerated(self): f = Path("gen.png") - meta = { - "prompt": json.dumps( - {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a dog"}}} - ) - } + meta = {"prompt": json.dumps(_wired_prompt("a dog"))} self.api.db.link_image_to_prompt = _raise self.assertEqual(self.api._ingest_scanned_file(f, meta), "added") self.assertEqual(self.api._ingest_scanned_file(f, meta), "found") def test_ingest_save_failure_counts_as_found(self): f = Path("gen.png") - meta = { - "prompt": json.dumps( - {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a bird"}}} - ) - } + meta = {"prompt": json.dumps(_wired_prompt("a bird"))} self.api.db.save_prompt = lambda *a, **k: None self.assertEqual(self.api._ingest_scanned_file(f, meta), "found") def test_batch_counts_and_per_file_errors(self): - meta = { - "prompt": json.dumps( - {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a fish"}}} - ) - } + meta = {"prompt": json.dumps(_wired_prompt("a fish"))} self.api.db.get_prompt_by_hash = _raise counts = self.api._ingest_scan_batch_sync( [(Path("a.png"), meta), (Path("b.png"), {})] @@ -2915,7 +2907,7 @@ class TestScanReadsVideos(AdminAPITestCase): shutil.which("ffmpeg") and shutil.which("ffprobe"), "ffmpeg not installed" ) async def test_scan_adds_prompts_found_in_videos(self): - graph = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}} + graph = _wired_prompt("a cat") _write_mp4( self.output_dir / "clip.mp4", json.dumps({"prompt": json.dumps(graph)}) ) diff --git a/tests/test_api_collections.py b/tests/test_api_collections.py new file mode 100644 index 00000000..c4dcdd86 --- /dev/null +++ b/tests/test_api_collections.py @@ -0,0 +1,494 @@ +"""Collections (#331): the routes under /prompt_manager/collections.""" + +import json +import ntpath +import sqlite3 +from unittest.mock import patch + +try: + from test_api_images import ImageAPITestCase, make_png + from open_handles import assert_closed +except ImportError: + from tests.test_api_images import ImageAPITestCase, make_png + from tests.open_handles import assert_closed + +BASE = "/prompt_manager/collections" +ITEMS = BASE + "/items" +MEMBERSHIP = BASE + "/membership" +SOURCES_CHANGED = "The sources changed. Refresh the images and try again." + + +class CollectionRoutesTestCase(ImageAPITestCase): + async def tearDownAsync(self): + self.api.db.close_all() + assert_closed(self, self._temp_db.name, self._output_tmp.name) + await super().tearDownAsync() + + def make(self, *names, root=None): + for name in names: + make_png((root or self.output_dir) / name) + + def scope(self, index=0): + roots = list(self.api._get_all_output_dirs()) + return self.api._gallery_root_scope(roots[index], index) + + def item(self, path, index=0, scope=None): + return { + "root_index": index, + "root_scope": scope or self.scope(index), + "path": path, + } + + def two_sources(self): + second = self.comfy_dir / "second" + second.mkdir() + self.api._get_all_output_dirs = lambda: [self.output_dir, second] + return second + + async def call(self, method, url, body=None, status=200, **params): + response = await self.client.request(method, url, json=body, params=params) + text = await response.text() + self.assertEqual(response.status, status, text) + data = json.loads(text) + self.assertNotIn(self.comfy_dir.as_posix(), text) + self.assertNotIn(str(self.comfy_dir).replace("\\", "\\\\"), text) + return data + + async def add(self, status=200, **body): + return await self.call("POST", ITEMS, body, status) + + async def listed(self): + return (await self.call("GET", BASE))["collections"] + + +class TestCollectionRoutes(CollectionRoutesTestCase): + async def test_each_route_succeeds(self): + self.make("a.png", "b.png", "c.png") + self.assertEqual( + await self.call("GET", BASE), {"success": True, "collections": []} + ) + made = await self.add( + name=" LoRa testing ", items=[self.item("a.png"), self.item("b.png")] + ) + grids = made["collection"] + self.assertEqual(grids["name"], "LoRa testing") + self.assertEqual( + made, + { + "success": True, + "collection": grids, + "created": True, + "added": 2, + "already": 0, + "not_found": 0, + }, + ) + again = await self.add( + collection_id=grids["id"], items=[self.item("b.png"), self.item("c.png")] + ) + self.assertEqual( + (again["created"], again["added"], again["already"]), (False, 1, 1) + ) + self.assertEqual( + await self.listed(), + [{"id": grids["id"], "name": "LoRa testing", "count": 3, "missing": 0}], + ) + membership = await self.call( + "GET", MEMBERSHIP, root_index="0", root_scope=self.scope(), path="a.png" + ) + self.assertEqual(membership, {"success": True, "collection_ids": [grids["id"]]}) + removed = await self.call( + "DELETE", + ITEMS, + {"collection_id": grids["id"], "items": [self.item("a.png")]}, + ) + self.assertEqual(removed, {"success": True, "collection": grids, "removed": 1}) + renamed = await self.call( + "PATCH", f"{BASE}/{grids['id']}", {"name": "LoRA grids"} + ) + self.assertEqual( + renamed, + {"success": True, "collection": {"id": grids["id"], "name": "LoRA grids"}}, + ) + deleted = await self.call("DELETE", f"{BASE}/{grids['id']}") + self.assertEqual(deleted, {"success": True, "deleted": True}) + self.assertEqual(await self.listed(), []) + # Files are never touched. + self.assertTrue((self.output_dir / "a.png").is_file()) + + async def test_names_are_unique_ignoring_case(self): + self.make("a.png", "b.png") + first = (await self.add(name="LoRa testing", items=[self.item("a.png")]))[ + "collection" + ] + second = await self.add(name="lora TESTING", items=[self.item("b.png")]) + self.assertEqual(second["collection"], first) + self.assertFalse(second["created"]) + self.assertEqual(len(await self.listed()), 1) + + async def test_bad_requests_are_400(self): + self.make("a.png") + made = await self.add(name="Set", items=[self.item("a.png")]) + cid = made["collection"]["id"] + good = [self.item("a.png")] + for body in ( + {"items": good}, + {"name": "x", "collection_id": cid, "items": good}, + {"name": "x"}, + {"name": "x", "items": good, "extra": 1}, + {"collection_id": "1", "items": good}, + {"collection_id": True, "items": good}, + {"collection_id": 0, "items": good}, + {"collection_id": 10**10, "items": good}, + {"name": "", "items": good}, + {"name": "x" * 101, "items": good}, + {"name": "bell\x07", "items": good}, + {"name": 5, "items": good}, + {"name": "x", "items": []}, + {"name": "x", "items": "a.png"}, + {"name": "x", "items": [self.item("a.png")] * 501}, + {"name": "x", "items": [{"root_index": 0, "path": "a.png"}]}, + {"name": "x", "items": [{**self.item("a.png"), "root_index": "0"}]}, + {"name": "x", "items": [{**self.item("a.png"), "root_scope": "xyz"}]}, + ): + with self.subTest(body=body): + data = await self.add(400, **body) + self.assertFalse(data["success"]) + for body in ({"items": good}, {"collection_id": cid}): + with self.subTest(remove=body): + await self.call("DELETE", ITEMS, body, 400) + for raw in ("abc", "0", "1234567890", "-1", "1.0"): + with self.subTest(path_id=raw): + await self.call("PATCH", f"{BASE}/{raw}", {"name": "x"}, 400) + await self.call("DELETE", f"{BASE}/{raw}", None, 400) + await self.call("PATCH", f"{BASE}/{cid}", {"name": ""}, 400) + await self.call("PATCH", f"{BASE}/{cid}", {"label": "x"}, 400) + for params in ( + {"root_index": "x", "root_scope": self.scope(), "path": "a.png"}, + {"root_index": "0", "path": "a.png"}, + {"root_index": "0", "root_scope": "nope", "path": "a.png"}, + ): + with self.subTest(membership=params): + await self.call("GET", MEMBERSHIP, status=400, **params) + self.assertEqual((await self.listed())[0]["count"], 1) + + async def test_unsafe_paths_are_403(self): + self.make("a.png") + for path in ( + "../a.png", + "sub/../../a.png", + "C:/a.png", + "a\\b.png", + "/a.png", + "notes.txt", + ): + with self.subTest(path=path): + await self.add(403, name="Set", items=[self.item(path)]) + await self.call( + "GET", + MEMBERSHIP, + status=403, + root_index="0", + root_scope=self.scope(), + path=path, + ) + self.assertEqual(await self.listed(), []) + + async def test_a_body_over_the_cap_is_413(self): + response = await self.client.post( + ITEMS, data=b'{"name":"x","items":[' + b" " * (1024 * 1024 + 1) + b"]}" + ) + self.assertEqual(response.status, 413) + + async def test_missing_collection_and_files_are_404(self): + self.make("a.png") + data = await self.add(404, collection_id=999, items=[self.item("a.png")]) + self.assertEqual(data["code"], "collection-not-found") + data = await self.call( + "DELETE", ITEMS, {"collection_id": 999, "items": [self.item("a.png")]}, 404 + ) + self.assertEqual(data["code"], "collection-not-found") + data = await self.call("PATCH", f"{BASE}/999", {"name": "x"}, 404) + self.assertEqual(data["code"], "collection-not-found") + data = await self.call("DELETE", f"{BASE}/999", None, 404) + self.assertEqual(data["code"], "collection-not-found") + data = await self.add(404, name="New", items=[self.item("never.png")]) + self.assertEqual(data["code"], "files-not-found") + self.assertEqual(await self.listed(), []) + + async def test_unlisted_files_count_as_not_found(self): + self.make("a.png") + made = await self.add( + name="Set", items=[self.item("a.png"), self.item("later.png")] + ) + self.assertEqual((made["added"], made["not_found"]), (1, 1)) + again = await self.add( + collection_id=made["collection"]["id"], + items=[self.item("a.png"), self.item("later.png")], + ) + self.assertEqual( + (again["added"], again["already"], again["not_found"]), (0, 1, 1) + ) + + async def test_name_taken_is_409(self): + self.make("a.png") + first = (await self.add(name="First", items=[self.item("a.png")]))["collection"] + second = (await self.add(name="Second", items=[self.item("a.png")]))[ + "collection" + ] + data = await self.call( + "PATCH", f"{BASE}/{second['id']}", {"name": "FIRST"}, 409 + ) + self.assertEqual( + (data["code"], data["collection_id"]), ("name-taken", first["id"]) + ) + renamed = await self.call("PATCH", f"{BASE}/{second['id']}", {"name": "SECOND"}) + self.assertEqual(renamed["collection"]["name"], "SECOND") + + async def test_stale_scope_is_409_after_reordering_sources(self): + second = self.two_sources() + self.make("a.png") + self.make("b.png", root=second) + stale = self.item("b.png", 1) + made = await self.add(name="Set", items=[self.item("a.png")]) + cid = made["collection"]["id"] + self.api._get_all_output_dirs = lambda: [second, self.output_dir] + for method, body in ( + ("POST", {"collection_id": cid, "items": [self.item("b.png", 0), stale]}), + ("DELETE", {"collection_id": cid, "items": [stale]}), + ): + with self.subTest(method=method): + data = await self.call(method, ITEMS, body, 409) + self.assertEqual(data["code"], "source-changed") + self.assertEqual(data["error"], SOURCES_CHANGED) + data = await self.call( + "GET", + MEMBERSHIP, + status=409, + root_index="1", + root_scope=stale["root_scope"], + path="b.png", + ) + self.assertEqual(data["code"], "source-changed") + # Nothing was written by the refused add; the member follows its folder. + listed = await self.listed() + self.assertEqual((listed[0]["count"], listed[0]["missing"]), (1, 0)) + membership = await self.call( + "GET", MEMBERSHIP, root_index="1", root_scope=self.scope(1), path="a.png" + ) + self.assertEqual(membership["collection_ids"], [cid]) + + async def test_counts_exclude_a_missing_member(self): + second = self.two_sources() + self.make("a.png", "b.png") + self.make("c.png", root=second) + made = await self.add( + name="Grids", + items=[self.item("a.png"), self.item("b.png"), self.item("c.png", 1)], + ) + other = await self.add(name="Other", items=[self.item("c.png", 1)]) + (self.output_dir / "b.png").unlink() + self.api.invalidate_gallery_cache() + listed = await self.listed() + self.assertEqual( + listed, + [ + { + "id": made["collection"]["id"], + "name": "Grids", + "count": 2, + "missing": 1, + }, + { + "id": other["collection"]["id"], + "name": "Other", + "count": 1, + "missing": 0, + }, + ], + ) + # The member is removable although it is no longer listed. + removed = await self.call( + "DELETE", + ITEMS, + {"collection_id": made["collection"]["id"], "items": [self.item("b.png")]}, + ) + self.assertEqual(removed["removed"], 1) + + async def test_counts_are_null_when_a_source_scan_fails(self): + self.make("a.png") + await self.add(name="Set", items=[self.item("a.png")]) + + async def broken(_root): + raise OSError("scan failed at /secret/place") + + with patch.object(self.api, "_get_gallery_files", side_effect=broken): + data = await self.call("GET", BASE) + self.assertEqual( + data["collections"][0] | {"id": 0}, + {"id": 0, "name": "Set", "count": None, "missing": None}, + ) + self.assertNotIn("secret", json.dumps(data)) + + async def test_windows_case_spellings_are_one_member(self): + self.make("Sub/A.png") + with patch("py.api.images._path_flavor", ntpath): + made = await self.add( + name="Set", items=[self.item("Sub/A.png"), self.item("SUB/a.PNG")] + ) + self.assertEqual((made["added"], made["already"]), (1, 0)) + (self.output_dir / "Sub" / "A.png").rename( + self.output_dir / "Sub" / "a.png" + ) + (self.output_dir / "Sub").rename(self.output_dir / "sub") + self.api.invalidate_gallery_cache() + listed = await self.listed() + self.assertEqual((listed[0]["count"], listed[0]["missing"]), (1, 0)) + membership = await self.call( + "GET", + MEMBERSHIP, + root_index="0", + root_scope=self.scope(), + path="sub/a.png", + ) + self.assertEqual(membership["collection_ids"], [made["collection"]["id"]]) + response = await self.client.get( + "/prompt_manager/images/output", + params={ + "mode": "workspace", + "collection": str(made["collection"]["id"]), + }, + ) + page = await response.json() + self.assertEqual([i["path"] for i in page["images"]], ["sub/a.png"]) + + async def test_posix_case_distinct_files_stay_two_members(self): + self.make("A.png", "a.png") + if len(list(self.output_dir.iterdir())) != 2: + self.skipTest("case-insensitive filesystem") + made = await self.add( + name="Set", items=[self.item("A.png"), self.item("a.png")] + ) + self.assertEqual(made["added"], 2) + + async def test_a_500_never_carries_the_exception_text(self): + self.make("a.png") + made = await self.add(name="Set", items=[self.item("a.png")]) + cid = made["collection"]["id"] + + def broken(*_args, **_kwargs): + raise sqlite3.OperationalError("disk I/O error at /secret/prompts.db") + + cases = ( + ( + "list_collections", + "GET", + BASE, + None, + {}, + "Collections could not be read.", + ), + ( + "file_collections", + "GET", + MEMBERSHIP, + None, + {"root_index": "0", "root_scope": self.scope(), "path": "a.png"}, + "Collections could not be read.", + ), + ( + "add_to_collection", + "POST", + ITEMS, + {"collection_id": cid, "items": [self.item("a.png")]}, + {}, + "The collection could not be changed.", + ), + ( + "remove_from_collection", + "DELETE", + ITEMS, + {"collection_id": cid, "items": [self.item("a.png")]}, + {}, + "The collection could not be changed.", + ), + ( + "rename_collection", + "PATCH", + f"{BASE}/{cid}", + {"name": "x"}, + {}, + "The collection could not be changed.", + ), + ( + "delete_collection", + "DELETE", + f"{BASE}/{cid}", + None, + {}, + "The collection could not be changed.", + ), + ) + for method_name, method, url, body, params, message in cases: + with self.subTest(method=method_name): + with patch.object(self.api.db, method_name, side_effect=broken): + data = await self.call(method, url, body, 500, **params) + self.assertEqual(data, {"success": False, "error": message}) + + +class TestCollectionListingFilter(CollectionRoutesTestCase): + """#392's "all N matching" filter carries ``collection`` like the listing.""" + + async def test_a_filter_with_a_collection_resolves_only_members(self): + second = self.two_sources() + self.make("a.png", "b.png") + self.make("c.png", root=second) + made = await self.add( + name="Set", items=[self.item("a.png"), self.item("c.png", 1)] + ) + cid = str(made["collection"]["id"]) + selection = await self.api.resolve_listing_filter({"collection": cid}) + self.assertEqual( + sorted((index, path) for index, _scope, path in selection.targets), + [(0, "a.png"), (1, "c.png")], + ) + selection = await self.api.resolve_listing_filter( + {"collection": cid, "root": "1", "root_scope": self.scope(1)} + ) + self.assertEqual([t[2] for t in selection.targets], ["c.png"]) + for bad in ("0", "x", "1234567890"): + with self.subTest(bad=bad): + with self.assertRaises(ValueError): + from py.api.images import parse_listing_filter + + parse_listing_filter({"collection": bad}) + + async def test_bulk_category_acts_on_collection_members_only(self): + self.make("a.png", "b.png") + inside = self._save_prompt("Inside", category="Old") + outside = self._save_prompt("Outside", category="Old") + self.api.db.link_image_to_prompt(inside, str(self.output_dir / "a.png")) + self.api.db.link_image_to_prompt(outside, str(self.output_dir / "b.png")) + made = await self.add(name="Set", items=[self.item("a.png")]) + cid = str(made["collection"]["id"]) + data = await self.call( + "POST", + "/prompt_manager/bulk/category", + {"images_filter": {"collection": cid}, "category": "New"}, + ) + self.assertEqual((data["updated_count"], data["prompt_count"]), (1, 1)) + self.assertEqual(self.api.db.get_prompt_by_id(inside)["category"], "New") + self.assertEqual(self.api.db.get_prompt_by_id(outside)["category"], "Old") + + async def test_a_deleted_collection_in_a_bulk_filter_is_a_404(self): + self.make("a.png") + data = await self.call( + "POST", + "/prompt_manager/bulk/category", + {"images_filter": {"collection": "77"}, "category": "New"}, + status=404, + ) + self.assertEqual(data["code"], "collection-not-found") + with self.assertRaises(Exception) as caught: + await self.api.resolve_listing_filter({"collection": "77"}) + self.assertEqual(type(caught.exception).__name__, "CollectionMissing") diff --git a/tests/test_api_core.py b/tests/test_api_core.py index 07e6ffa7..bf46c323 100644 --- a/tests/test_api_core.py +++ b/tests/test_api_core.py @@ -835,7 +835,6 @@ class TestMetadataParsing(unittest.TestCase): self.assertEqual(parsed["positive_prompt"], "a cat") self.assertEqual(parsed["negative_prompt"], "dog") self.assertEqual(parsed["parameters"]["parameters"], meta["parameters"]) - self.assertEqual(self.api._extract_readable_prompt(parsed), "a cat") def test_parse_comfyui_fields(self): meta = { @@ -856,97 +855,6 @@ class TestMetadataParsing(unittest.TestCase): self.assertEqual(parsed["prompt"], "plain text") self.assertEqual(parsed["workflow"], {"nodes": []}) - def test_extract_readable_prompt_variants(self): - extract = self.api._extract_readable_prompt - self.assertEqual(extract({"prompt": "plain"}), "plain") - self.assertEqual(extract({"prompt": ["a", "", "b"]}), "a b") - self.assertEqual(extract({"prompt": 42}), "42") - self.assertEqual( - extract({"prompt": None, "parameters": {"positive": ["x", "y"]}}), "x y" - ) - self.assertIsNone(extract({"prompt": {}, "parameters": {}})) - - graph = { - "1": {"class_type": "CLIPTextEncode", "inputs": {"text": "pos"}}, - "2": {"class_type": "CLIPTextEncode", "inputs": {"text": "neg"}}, - "3": { - "class_type": "KSampler", - "inputs": {"positive": ["1", 0], "negative": ["2", 0]}, - }, - } - self.assertEqual(extract({"prompt": graph}), "pos") - self.assertEqual(extract({"prompt": {"x": 1}, "workflow": graph}), "pos") - - def test_node_inputs_and_text_lookup(self): - self.assertEqual(self.api._get_node_inputs("nope"), {}) - self.assertEqual(self.api._get_node_inputs({"inputs": 5}), {}) - self.assertEqual(self.api._get_node_inputs({"inputs": {"a": 1}}), {"a": 1}) - listed = self.api._get_node_inputs( - {"inputs": [{"name": "text", "link": None}, {"nope": 1}, "junk"]} - ) - self.assertEqual(list(listed), ["text"]) - - self.assertIsNone(self.api._find_text_in_node(None)) - self.assertEqual(self.api._find_text_in_node({"inputs": {"text": "hi"}}), "hi") - self.assertEqual( - self.api._find_text_in_node( - {"type": "CLIPTextEncode", "widgets_values": ["w"]} - ), - "w", - ) - self.assertIsNone( - self.api._find_text_in_node( - {"type": "CLIPTextEncode", "widgets_values": [" "]} - ) - ) - self.assertIsNone(self.api._find_text_in_node({"type": "KSampler"})) - - def test_positive_prompt_from_nodes_array_and_fallbacks(self): - extract = self.api._extract_positive_prompt_from_comfyui_data - self.assertIsNone(extract("x")) - self.assertIsNone(extract({})) - self.assertIsNone(extract({"nodes": []})) - - workflow = { - "nodes": [ - {"id": 1, "type": "CLIPTextEncode", "widgets_values": ["pos"]}, - { - "id": 2, - "type": "CLIPTextEncode", - "widgets_values": ["neg"], - "title": "Negative", - }, - {"id": 3, "type": "KSampler", "inputs": [{"name": "seed", "link": 4}]}, - "junk", - ] - } - self.assertEqual(extract(workflow), "pos") - - only_negative = { - "nodes": [ - { - "id": 2, - "type": "CLIPTextEncode", - "widgets_values": ["neg"], - "title": "neg", - }, - ] - } - self.assertEqual(extract(only_negative), "neg") - - bad_link = { - "1": {"class_type": "CLIPTextEncode", "inputs": {"text": "pos"}}, - "3": { - "class_type": "KSampler", - "inputs": {"positive": "oops", "negative": []}, - }, - "4": { - "class_type": "KSampler", - "inputs": {"positive": ["zz", 0], "negative": []}, - }, - } - self.assertEqual(extract(bad_link), "pos") - if __name__ == "__main__": unittest.main() diff --git a/tests/test_api_duplicate_files.py b/tests/test_api_duplicate_files.py index 91bd0ea3..a708e229 100644 --- a/tests/test_api_duplicate_files.py +++ b/tests/test_api_duplicate_files.py @@ -687,6 +687,39 @@ class TestCleanupJob(DuplicateFilesTestCase): "Moved 2 database links to the kept copies", "\n".join(cleanup.log) ) + async def test_duplicate_removal_moves_collection_memberships(self): + self.trio() + db = self.api.db + root_key = duplicate_files.favorite_root_key(self.output_dir) + grids = db.add_to_collection(name="LoRA grids", items=[(root_key, "b.png")]) + grids = grids["collection"]["id"] + other = db.add_to_collection( + name="Other", items=[(root_key, "sub/c.png"), (root_key, "d.png")] + )["collection"]["id"] + db.add_to_collection(collection_id=grids, items=[(root_key, "sub/c.png")]) + _plan, cleanup = await self.remove_all_but("a.png") + # The union: a.png joins both; c.png's second grids row was a repeat. + self.assertEqual(db.collection_members(grids, root_key), {"a.png"}) + self.assertEqual(db.collection_members(other, root_key), {"a.png", "d.png"}) + self.assertEqual(sorted(db.file_collections(root_key, "a.png")), [grids, other]) + self.assertEqual(db.file_collections(root_key, "b.png"), []) + self.assertEqual(db.file_collections(root_key, "sub/c.png"), []) + self.assertIn( + "Moved 2 collection memberships to the kept copies", "\n".join(cleanup.log) + ) + + def test_one_collection_membership_is_singular(self): + job = type( + "Job", (), {"log": [], "log_line": lambda s, line: s.log.append(line)} + )() + totals = dict.fromkeys( + ("moved", "merged", "favorites", "tags", "review", "collections"), 0 + ) + duplicate_files.DuplicateFilesMixin._log_moved_records( + job, {**totals, "collections": 1} + ) + self.assertEqual(job.log, ["Moved 1 collection membership to the kept copies"]) + async def test_a_failed_removal_leaves_its_links_on_the_kept_copy(self): self.trio() prompt = self.save("Stays on c") @@ -843,8 +876,11 @@ class TestLinkQueries(DuplicateFilesTestCase): with self.subTest(sql=sql): self.assertNotIn("SCAN generated_images", plan) self.assertNotIn("SCAN image_favorites", plan) + self.assertNotIn("SCAN image_collection_items", plan) finally: conn.close() + statements = duplicate_files.query_plan_sql() + self.assertTrue(any("image_collection_items" in sql for sql in statements)) class TestReviewHelpers(DuplicateFilesTestCase): diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index f30f3947..26211a5b 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -2494,3 +2494,171 @@ class TestListingFilterSelection(ImageAPITestCase): self.assertEqual(response.status, 409) self.assertEqual((await response.json())["code"], "source-changed") + + +class TestWorkspaceCollectionFilter(ImageAPITestCase): + """#331: ``collection=`` keeps only that collection's listed members.""" + + media = TestWorkspaceImages.media + get = TestWorkspaceImages.get + _names = TestWorkspaceImages._names + + async def tearDownAsync(self): + self.api.db.close_all() + assert_closed(self, self._temp_db.name, self._output_tmp.name) + await super().tearDownAsync() + + def keys(self, root, relative): + from py.api.images import _favorite_relative_key, _favorite_root_key + + return (_favorite_root_key(root), _favorite_relative_key(relative)) + + def collect(self, name, *members): + """A collection holding ``(root, relative)`` members; its id.""" + result = self.api.db.add_to_collection( + name=name, items=[self.keys(root, rel) for root, rel in members] + ) + return result["collection"]["id"] + + def two_sources(self): + second = self.comfy_dir / "second" + second.mkdir() + self.api._get_all_output_dirs = lambda: [self.output_dir, second] + return second + + async def test_listing_keeps_only_collection_members(self): + second = self.two_sources() + self.media("a.png", 300) + self.media("b.png", 200) + self.media("sub/c.png", 100, root=second) + self.media("d.png", 50, root=second) + grids = self.collect( + "Grids", + (self.output_dir, "a.png"), + (second, "sub/c.png"), + (self.output_dir, "gone.png"), # stored, never listed + ) + self.collect("Other", (self.output_dir, "b.png")) + response, data = await self.get(collection=str(grids)) + self.assertEqual(response.status, 200, data) + self.assertEqual(self._names(data), ["a.png", "c.png"]) + self.assertEqual(data["total"], 2) + _, data = await self.get(collection=str(grids), root="1") + self.assertEqual(self._names(data), ["c.png"]) + # Without the filter the listing is unchanged. + _, data = await self.get() + self.assertEqual(data["total"], 4) + + async def test_collection_filter_combines_with_the_other_filters(self): + night = self._save_prompt("Neon zebra", category="Night", tags=["a"]) + day = self._save_prompt("Morning field", category="Day", tags=["a"]) + files = { + name: self.media(name, 100 + index) + for index, name in enumerate( + ("night.png", "day.png", "zebra_plain.png", "clip.mp4", "out.png") + ) + } + self.api.db.link_image_to_prompt(night, str(files["night.png"])) + self.api.db.link_image_to_prompt(day, str(files["day.png"])) + members = [(self.output_dir, name) for name in files if name != "out.png"] + grids = str(self.collect("Grids", *members)) + cases = ( + ({}, ["clip.mp4", "day.png", "night.png", "zebra_plain.png"]), + ({"q": "zebra"}, ["night.png", "zebra_plain.png"]), + ({"category": "Night"}, ["night.png"]), + ({"tags": "a"}, ["day.png", "night.png"]), + ({"media": "video"}, ["clip.mp4"]), + ({"media": "image", "q": "day"}, ["day.png"]), + ) + for params, expected in cases: + with self.subTest(**params): + response, data = await self.get(collection=grids, sort="name", **params) + self.assertEqual(response.status, 200, data) + self.assertEqual(self._names(data), expected) + self.assertEqual(data["total"], len(expected)) + + async def test_invalid_collection_is_a_400(self): + for value in ("x", "0", "-1", "1.5", "1234567890", " 1"): + with self.subTest(value=value): + response, data = await self.get(collection=value) + self.assertEqual(response.status, 400, data) + + async def test_unknown_collection_is_a_404(self): + self.media("a.png") + response, data = await self.get(collection="77") + self.assertEqual(response.status, 404) + self.assertEqual(data["code"], "collection-not-found") + self.assertNotIn("images", data) + self.api._get_all_output_dirs = lambda: [] + response, data = await self.get(collection="77") + self.assertEqual((response.status, data["code"]), (404, "collection-not-found")) + + async def test_a_collection_deleted_during_the_read_is_a_404(self): + self.media("a.png") + grids = self.collect("Grids", (self.output_dir, "a.png")) + real = self.api.db.collection_members + + def deleted_meanwhile(collection_id, root_key): + self.api.db.delete_collection(collection_id) + return real(collection_id, root_key) + + with patch.object( + self.api.db, "collection_members", side_effect=deleted_meanwhile + ): + response, data = await self.get(collection=str(grids)) + self.assertEqual((response.status, data["code"]), (404, "collection-not-found")) + + async def test_a_failed_members_read_is_a_500_never_an_unfiltered_page(self): + self.media("a.png") + self.media("b.png") + grids = self.collect("Grids", (self.output_dir, "a.png")) + + def broken(*_args, **_kwargs): + raise sqlite3.OperationalError("database is locked") + + with patch.object(self.api.db, "collection_members", side_effect=broken): + response, data = await self.get(collection=str(grids)) + self.assertEqual(response.status, 500) + self.assertNotIn("images", data) + self.assertNotIn("locked", json.dumps(data)) + + async def test_reordering_sources_keeps_members(self): + second = self.two_sources() + self.media("a.png", root=second) + self.media("b.png") + grids = str(self.collect("Grids", (second, "a.png"))) + _, data = await self.get(collection=grids) + self.assertEqual( + [(i["root_index"], i["path"]) for i in data["images"]], [(1, "a.png")] + ) + self.api._get_all_output_dirs = lambda: [second, self.output_dir] + _, data = await self.get(collection=grids) + self.assertEqual( + [(i["root_index"], i["path"]) for i in data["images"]], [(0, "a.png")] + ) + + async def test_cursor_from_another_collection_is_refused(self): + for index, name in enumerate(("a.png", "b.png", "c.png")): + self.media(name, 100 + index) + everything = [(self.output_dir, name) for name in ("a.png", "b.png", "c.png")] + first = str(self.collect("First", *everything)) + second = str(self.collect("Second", *everything)) + _, page = await self.get(collection=first, limit=1) + token = page["next_cursor"] + self.assertIsInstance(token, str) + response, _ = await self.get(collection=first, limit=1, cursor=token) + self.assertEqual(response.status, 200) + for params in ({"collection": second}, {}, {"favorites": "1"}): + with self.subTest(**params): + response, _ = await self.get(cursor=token, limit=1, **params) + self.assertEqual(response.status, 400) + + async def test_collection_wins_over_favorites(self): + self.media("a.png") + self.media("b.png") + grids = str(self.collect("Grids", (self.output_dir, "a.png"))) + from py.api.images import _favorite_root_key + + self.api.db.set_image_favorite(_favorite_root_key(self.output_dir), "b.png") + _, data = await self.get(collection=grids, favorites="1") + self.assertEqual(self._names(data), ["a.png"]) diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 2437220b..8e08b588 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -614,6 +614,50 @@ class TestRestoreFromFile(BackupTestCase): ) assert_closed(self, str(backups[0])) + def test_backup_and_restore_keep_collections(self): + with self.db.model.get_connection() as conn: + conn.execute( + "INSERT INTO image_collections (id, name, name_key, created_at) " + "VALUES (4, 'LoRa testing', 'lora testing', '2026-10-27T00:00:00Z')" + ) + conn.execute( + "INSERT INTO image_collection_items VALUES " + "(4, 'root', 'grid.png', '2026-10-27T00:00:00Z')" + ) + with tempfile.TemporaryDirectory() as directory: + # The rows still sit in the WAL; a backup must carry them. + backup = os.path.join(directory, "backup.db") + self.assertTrue(self.db.model.backup_database(backup)) + self.assertEqual( + PromptModel.verify_database_file(backup, require_library_state=True), + (True, "ok"), + ) + with self.db.model.get_connection() as conn: + conn.execute("DELETE FROM image_collections") + + self.db.model.restore_from_file(backup) + + conn = self.db.model.get_connection() + self.assertEqual( + [ + tuple(r) + for r in conn.execute("SELECT id, name FROM image_collections") + ], + [(4, "LoRa testing")], + ) + self.assertEqual( + [ + tuple(r) + for r in conn.execute( + "SELECT collection_id, root_key, relative_path " + "FROM image_collection_items" + ) + ], + [(4, "root", "grid.png")], + ) + PromptModel.close_all_instances() + assert_closed(self, directory) + def test_restore_without_a_live_database_returns_empty_backup_path(self): src = self._make_source_db(["fresh"]) self.db.close_all() diff --git a/tests/test_image_collections.py b/tests/test_image_collections.py new file mode 100644 index 00000000..9136c862 --- /dev/null +++ b/tests/test_image_collections.py @@ -0,0 +1,223 @@ +"""Collections (#331): the database operations.""" + +import sqlite3 +import tempfile +import unittest +from pathlib import Path + +try: + from tests.open_handles import assert_closed +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + +from database.models import PromptModel +from database.operations import PromptDatabase + +try: # guarded so a missing module fails assertions, not test loading + from database.image_collections import ( + CollectionFilesMissing, + CollectionInvalid, + CollectionMissing, + CollectionNameTaken, + READ_SQL, + ) +except ImportError: + CollectionFilesMissing = CollectionMissing = CollectionNameTaken = None + CollectionInvalid = ValueError + READ_SQL = None + +ROOT = "/real/output" +OTHER = "/real/other" + + +class ImageCollectionOperationsTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.addCleanup(assert_closed, self, self.temp.name) + self.addCleanup(PromptModel.reset_schema_cache) + self.db = PromptDatabase(str(Path(self.temp.name) / "prompts.db")) + self.addCleanup(self.db.close_all) + + def require(self): + self.assertTrue( + hasattr(self.db, "add_to_collection"), + "ImageCollectionOperations is not implemented", + ) + + def add(self, items, **target): + return self.db.add_to_collection(items=items, **target) + + def test_add_creates_a_collection_by_name(self): + self.require() + result = self.add( + [(ROOT, "a.png"), (ROOT, "b.png"), (OTHER, "a.png")], name="LoRa testing" + ) + collection = result["collection"] + self.assertEqual(collection["name"], "LoRa testing") + self.assertEqual( + result, + {"collection": collection, "created": True, "added": 3, "already": 0}, + ) + self.assertEqual(self.db.list_collections(), [collection]) + self.assertTrue(self.db.collection_exists(collection["id"])) + self.assertEqual(self.db.get_collection(collection["id"]), collection) + self.assertEqual( + self.db.collection_members(collection["id"], ROOT), {"a.png", "b.png"} + ) + self.assertEqual( + self.db.root_collection_members(OTHER), {collection["id"]: {"a.png"}} + ) + self.assertEqual(self.db.collection_counts(), {collection["id"]: 3}) + self.assertEqual(self.db.file_collections(ROOT, "a.png"), [collection["id"]]) + self.assertEqual(self.db.file_collections(ROOT, "zzz.png"), []) + + def test_add_by_name_ignores_case_and_unicode_form(self): + self.require() + first = self.add([(ROOT, "a.png")], name="Café grids")["collection"] + # NFD spelling (macOS) in another case joins the same collection. + second = self.add([(ROOT, "b.png")], name="CAFÉ GRIDS") + self.assertEqual(second["collection"], first) + self.assertFalse(second["created"]) + self.assertEqual(len(self.db.list_collections()), 1) + self.assertEqual( + self.db.collection_members(first["id"], ROOT), {"a.png", "b.png"} + ) + + def test_add_counts_already_present(self): + self.require() + made = self.add([(ROOT, "a.png"), (ROOT, "b.png")], name="Set") + result = self.add( + [(ROOT, "b.png"), (ROOT, "c.png"), (ROOT, "c.png")], + collection_id=made["collection"]["id"], + ) + self.assertEqual( + (result["created"], result["added"], result["already"]), (False, 1, 1) + ) + self.assertEqual(self.db.collection_counts(), {made["collection"]["id"]: 3}) + + def test_add_with_a_new_name_and_nothing_inserted_creates_nothing(self): + self.require() + with self.assertRaises(CollectionFilesMissing): + self.add([], name="Empty") + self.assertEqual(self.db.list_collections(), []) + # An existing name with nothing to add is not an error. + made = self.add([(ROOT, "a.png")], name="Kept")["collection"] + self.assertEqual( + self.add([], name="kept"), + {"collection": made, "created": False, "added": 0, "already": 0}, + ) + + def test_add_to_a_deleted_collection_raises_missing(self): + self.require() + made = self.add([(ROOT, "a.png")], name="Gone")["collection"] + self.db.delete_collection(made["id"]) + with self.assertRaises(CollectionMissing): + self.add([(ROOT, "b.png")], collection_id=made["id"]) + with self.assertRaises(CollectionMissing): + self.db.remove_from_collection(made["id"], [(ROOT, "a.png")]) + with self.assertRaises(CollectionMissing): + self.db.rename_collection(made["id"], "Back") + self.assertFalse(self.db.collection_exists(made["id"])) + self.assertIsNone(self.db.get_collection(made["id"])) + + def test_remove_returns_the_count(self): + self.require() + made = self.add([(ROOT, "a.png"), (ROOT, "b.png")], name="Set")["collection"] + result = self.db.remove_from_collection( + made["id"], [(ROOT, "a.png"), (ROOT, "a.png"), (ROOT, "never.png")] + ) + self.assertEqual(result, {"collection": made, "removed": 1}) + self.assertEqual(self.db.collection_members(made["id"], ROOT), {"b.png"}) + # Removing the last member keeps the (now empty) collection. + self.db.remove_from_collection(made["id"], [(ROOT, "b.png")]) + self.assertEqual(self.db.list_collections(), [made]) + self.assertEqual(self.db.collection_counts(), {made["id"]: 0}) + + def test_rename_refuses_a_taken_name(self): + self.require() + first = self.add([(ROOT, "a.png")], name="First")["collection"] + second = self.add([(ROOT, "a.png")], name="Second")["collection"] + with self.assertRaises(CollectionNameTaken) as caught: + self.db.rename_collection(second["id"], " first ") + self.assertEqual(caught.exception.collection_id, first["id"]) + self.assertEqual(self.db.get_collection(second["id"]), second) + + def test_rename_accepts_a_case_variant_of_its_own_name(self): + self.require() + made = self.add([(ROOT, "a.png")], name="lora grids")["collection"] + renamed = self.db.rename_collection(made["id"], "LoRA grids ") + self.assertEqual(renamed, {"id": made["id"], "name": "LoRA grids"}) + self.assertEqual(self.db.list_collections(), [renamed]) + other = self.db.rename_collection(made["id"], "Portraits") + self.assertEqual(other["name"], "Portraits") + + def test_delete_unknown_raises_missing(self): + self.require() + made = self.add([(ROOT, "a.png")], name="Doomed")["collection"] + self.assertTrue(self.db.delete_collection(made["id"])) + self.assertEqual(self.db.root_collection_members(ROOT), {}) + with self.assertRaises(CollectionMissing): + self.db.delete_collection(made["id"]) + with self.assertRaises(CollectionMissing): + self.db.delete_collection(999) + + def test_list_is_in_name_key_order(self): + self.require() + for name in ("beta", "Alpha", "été", "Zeta", "alpha two"): + self.add([(ROOT, "a.png")], name=name) + self.assertEqual( + [row["name"] for row in self.db.list_collections()], + ["Alpha", "alpha two", "beta", "Zeta", "été"], + ) + + def test_names_are_trimmed_and_collapsed(self): + self.require() + made = self.add([(ROOT, "a.png")], name=" LoRa \t testing\n ")["collection"] + self.assertEqual(made["name"], "LoRa testing") + for bad in ("", " ", "x" * 101, "bell\x07", 5, None): + with self.subTest(name=bad): + with self.assertRaises(CollectionInvalid): + self.add([(ROOT, "a.png")], name=bad) + self.assertEqual(len(self.db.list_collections()), 1) + + def test_targets_and_items_are_validated(self): + self.require() + made = self.add([(ROOT, "a.png")], name="Set")["collection"] + for target in ( + {}, + {"name": "x", "collection_id": made["id"]}, + {"collection_id": True}, + {"collection_id": 0}, + {"collection_id": "1"}, + ): + with self.subTest(target=target): + with self.assertRaises(CollectionInvalid): + self.add([(ROOT, "a.png")], **target) + for items in ([("root",)], [(ROOT, 5)], "a.png", [(ROOT, "x" * 4097)]): + with self.subTest(items=items): + with self.assertRaises(CollectionInvalid): + self.add(items, collection_id=made["id"]) + + def test_reads_use_an_index(self): + self.require() + self.add([(ROOT, "a.png")], name="Set") + conn = sqlite3.connect(self.db.model.db_path) + try: + for sql in READ_SQL: + plan = " ".join( + row[3] + for row in conn.execute( + "EXPLAIN QUERY PLAN " + sql, [None] * sql.count("?") + ) + ) + with self.subTest(sql=sql): + self.assertNotIn("SCAN image_collection_items", plan) + if "?" in sql: # a lookup, not the whole list + self.assertNotIn("SCAN image_collections", plan) + finally: + conn.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_image_collections_schema.py b/tests/test_image_collections_schema.py new file mode 100644 index 00000000..fb06ba2d --- /dev/null +++ b/tests/test_image_collections_schema.py @@ -0,0 +1,485 @@ +"""Collections schema (#331): install, validation and restore safety.""" + +import re +import sqlite3 +import tempfile +import unittest +import unicodedata +from contextlib import closing +from pathlib import Path +from unittest.mock import patch + +try: + from tests.open_handles import assert_closed + from tests.test_annotation_schema import LEGACY_TABLES, table_rows +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + from test_annotation_schema import LEGACY_TABLES, table_rows + +from database.models import PromptModel + +try: # guarded so a missing module fails assertions, not test loading + from database import image_collections + from database.image_collections import ( + image_collections_schema_error, + migrate_image_collections, + ) +except ImportError: + image_collections = None + image_collections_schema_error = None + migrate_image_collections = None + +COLLECTIONS = "image_collections" +ITEMS = "image_collection_items" +INDEX = "idx_image_collection_items_file" +TABLES = (COLLECTIONS, ITEMS) +DDL = { + COLLECTIONS: ( + f"CREATE TABLE {COLLECTIONS} (id INTEGER PRIMARY KEY AUTOINCREMENT, " + "name TEXT NOT NULL, name_key TEXT NOT NULL UNIQUE, " + "created_at TEXT NOT NULL)" + ), + ITEMS: ( + f"CREATE TABLE {ITEMS} (collection_id INTEGER NOT NULL REFERENCES " + f"{COLLECTIONS}(id) ON DELETE CASCADE, root_key TEXT NOT NULL, " + "relative_path TEXT NOT NULL, added_at TEXT NOT NULL, " + "PRIMARY KEY (collection_id, root_key, relative_path))" + ), +} +INDEX_DDL = f"CREATE INDEX {INDEX} ON {ITEMS}(root_key, relative_path)" +NOW = "2026-10-27T00:00:00.000000+00:00" +IMAGES = 23_000 + + +def _compact(sql): + return re.sub(r"\s+", "", sql or "").lower() + + +def expected_catalog(): + return { + ("table", COLLECTIONS, _compact(DDL[COLLECTIONS])), + ("table", ITEMS, _compact(DDL[ITEMS])), + ("index", INDEX, _compact(INDEX_DDL)), + ("index", f"sqlite_autoindex_{COLLECTIONS}_1", ""), + ("index", f"sqlite_autoindex_{ITEMS}_1", ""), + } + + +def catalog(conn): + return { + (kind, name, _compact(sql)) + for kind, name, sql in conn.execute( + "SELECT type,name,sql FROM sqlite_master " + "WHERE name IN (?,?,?) OR tbl_name IN (?,?)", + (*TABLES, INDEX, *TABLES), + ) + } + + +def key(name): + return unicodedata.normalize("NFC", unicodedata.normalize("NFC", name).casefold()) + + +def schema(conn, exclude=()): + return sorted( + tuple(row) + for row in conn.execute( + "SELECT type,name,tbl_name,sql FROM sqlite_master ORDER BY name" + ) + if row[1] not in exclude and row[2] not in exclude + ) + + +def row_counts(conn): + tables = [ + name + for (name,) in conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name" + ) + ] + return { + table: conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] # nosec + for table in tables + } + + +class ImageCollectionsSchemaTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.addCleanup(assert_closed, self, self.temp.name) + self.addCleanup(PromptModel.reset_schema_cache) + self.models = [] + self.addCleanup(self.close_models) + self.path = str(Path(self.temp.name) / "live.db") + + def close_models(self): + for model in self.models: + model.close_all() + + def model(self, path=None): + model = PromptModel(path or self.path) + self.models.append(model) + return model + + def require_module(self): + self.assertIsNotNone( + image_collections_schema_error, + "database.image_collections is not implemented", + ) + + def closed_current(self, name="current.db", *statements): + """A closed current database, altered afterwards by *statements*.""" + path = str(Path(self.temp.name) / name) + model = self.model(path) + model.close_all() + PromptModel.reset_schema_cache() + with closing(sqlite3.connect(path)) as raw, raw: + raw.execute("PRAGMA foreign_keys=OFF") + for statement in statements: + if isinstance(statement, tuple): + raw.execute(*statement) + else: + raw.execute(statement) + return path + + @staticmethod + def drop_group(path): + with closing(sqlite3.connect(path)) as raw, raw: + raw.execute(f"DROP TABLE {ITEMS}") + raw.execute(f"DROP TABLE {COLLECTIONS}") + + def test_migration_creates_tables_and_index(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + self.assertEqual(catalog(conn), expected_catalog()) + self.assertIsNone(image_collections_schema_error(conn, allow_missing=False)) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {COLLECTIONS}").fetchone()[0], 0 + ) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {ITEMS}").fetchone()[0], 0 + ) + self.assertEqual( + PromptModel.verify_database_file(self.path, require_library_state=True), + (True, "ok"), + ) + + def test_migration_is_idempotent(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + conn.execute( + f"INSERT INTO {COLLECTIONS} (name, name_key, created_at) " + "VALUES ('Kept', 'kept', ?)", + (NOW,), + ) + conn.execute(f"INSERT INTO {ITEMS} VALUES (1, 'root', 'a.png', ?)", (NOW,)) + before = schema(conn) + model.close_all() + for _ in range(2): + PromptModel.reset_schema_cache() + model = self.model() + with model.get_connection() as conn: + migrate_image_collections(conn) + self.assertEqual(schema(conn), before) + self.assertEqual( + [tuple(r) for r in conn.execute(f"SELECT * FROM {ITEMS}")], + [(1, "root", "a.png", NOW)], + ) + model.close_all() + + def test_migration_reads_only_its_own_tables(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + conn.execute("INSERT INTO prompts(id,text,hash) VALUES(1,'positive','p')") + conn.executemany( + "INSERT INTO generated_images(prompt_id,image_path,filename," + "file_path) VALUES(1,?,?,?)", + ((f"/out/{n}.png", f"{n}.png", f"/out/{n}.png") for n in range(IMAGES)), + ) + conn.execute("INSERT INTO image_favorites VALUES('root','a.png',?)", (NOW,)) + model.close_all() + self.drop_group(self.path) + with closing(sqlite3.connect(self.path)) as raw: + raw.execute("PRAGMA foreign_keys=ON") + counts_before = row_counts(raw) + schema_before = schema(raw) + legacy_before = table_rows(raw, LEGACY_TABLES) + self.assertEqual(counts_before["generated_images"], IMAGES) + touched = [] + + def record(action, arg1, arg2, _db, _trigger): + touched.append((action, arg1, arg2)) + return sqlite3.SQLITE_OK + + raw.set_authorizer(record) + try: + migrate_image_collections(raw) + finally: + raw.set_authorizer(None) + raw.commit() + reads = { + arg1 for action, arg1, _arg2 in touched if action == sqlite3.SQLITE_READ + } + self.assertTrue(reads) + self.assertLessEqual(reads, {"sqlite_master", COLLECTIONS, ITEMS}) + writes = { + arg1 + for action, arg1, _arg2 in touched + if action + in (sqlite3.SQLITE_INSERT, sqlite3.SQLITE_UPDATE, sqlite3.SQLITE_DELETE) + } + self.assertLessEqual(writes, {"sqlite_master"}) + self.assertEqual(catalog(raw), expected_catalog()) + counts_after = row_counts(raw) + self.assertEqual(counts_after.pop(COLLECTIONS), 0) + self.assertEqual(counts_after.pop(ITEMS), 0) + self.assertEqual(counts_after, counts_before) + self.assertEqual(schema(raw, exclude=TABLES), schema_before) + self.assertEqual(table_rows(raw, LEGACY_TABLES), legacy_before) + # A second run is a no-op. + after_first = schema(raw) + changes = raw.total_changes + migrate_image_collections(raw) + raw.commit() + self.assertEqual(schema(raw), after_first) + self.assertEqual(raw.total_changes, changes) + + def test_migration_refuses_a_foreign_object_of_the_same_name(self): + self.require_module() + for statements in ( + (f"CREATE TABLE {COLLECTIONS} (x)",), + (f"CREATE VIEW {COLLECTIONS} AS SELECT 1",), + (f"CREATE INDEX {INDEX} ON prompts(text)",), + ): + with self.subTest(statements=statements): + path = str(Path(self.temp.name) / f"foreign{len(self.models)}.db") + self.model(path).close_all() + PromptModel.reset_schema_cache() + self.drop_group(path) + with closing(sqlite3.connect(path)) as raw, raw: + for statement in statements: + raw.execute(statement) + with closing(sqlite3.connect(path)) as raw: + before = schema(raw) + self.assertIsNotNone(image_collections_schema_error(raw)) + with self.assertRaises(sqlite3.DatabaseError): + self.model(path).get_connection() + PromptModel.close_all_instances() + PromptModel.reset_schema_cache() + with closing(sqlite3.connect(path)) as raw: + self.assertEqual(schema(raw), before) + self.assertIsNone( + raw.execute( + "SELECT 1 FROM sqlite_master WHERE name=?", (ITEMS,) + ).fetchone() + ) + self.assertFalse(PromptModel.verify_database_file(path)[0]) + + def test_migration_rolls_back_on_error(self): + self.require_module() + self.model().close_all() + PromptModel.reset_schema_cache() + self.drop_group(self.path) + with closing(sqlite3.connect(self.path)) as raw: + with patch.object( + image_collections, + "image_collections_schema_error", + side_effect=[None, "Image collections went wrong"], + ): + with self.assertRaisesRegex(sqlite3.DatabaseError, "went wrong"): + migrate_image_collections(raw) + raw.commit() + self.assertEqual(catalog(raw), set()) + + def test_startup_accepts_a_value_saved_under_older_rules(self): + self.require_module() + path = self.closed_current( + "older.db", + ( + f"INSERT INTO {COLLECTIONS} (name, name_key, created_at) " + "VALUES (?, 'not the key', ?)", + (" " + "x" * 150 + " ", NOW), + ), + ) + model = self.model(path) + with model.get_connection() as conn: + self.assertIsNone(image_collections_schema_error(conn, check_values=False)) + self.assertIsNotNone(image_collections_schema_error(conn)) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {COLLECTIONS}").fetchone()[0], 1 + ) + model.close_all() + self.assertFalse( + PromptModel.verify_database_file(path, require_library_state=True)[0] + ) + + def test_restore_accepts_an_old_backup_without_collections(self): + self.require_module() + path = self.closed_current("old.db") + self.drop_group(path) + with closing(sqlite3.connect(path)) as raw: + self.assertIsNone(image_collections_schema_error(raw)) + self.assertEqual( + image_collections_schema_error(raw, allow_missing=False), + "Missing image collections schema", + ) + self.assertEqual(PromptModel.verify_database_file(path), (True, "ok")) + model = self.model(path) + with model.get_connection() as conn: + self.assertEqual(catalog(conn), expected_catalog()) + model.close_all() + self.assertEqual( + PromptModel.verify_database_file(path, require_library_state=True), + (True, "ok"), + ) + + def test_restore_requires_collections_on_a_current_stage(self): + self.require_module() + path = self.closed_current("stage.db") + self.drop_group(path) + self.assertEqual( + PromptModel.verify_database_file(path, require_library_state=True), + (False, "Missing image collections schema"), + ) + # One table without the other is never acceptable. + partial = self.closed_current("partial.db", f"DROP TABLE {ITEMS}") + self.assertFalse(PromptModel.verify_database_file(partial)[0]) + + def assert_refused(self, name, *statements): + self.require_module() + path = self.closed_current( + f"{name}.db", + ( + f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) " + "VALUES (1, 'Good', 'good', ?)", + (NOW,), + ), + (f"INSERT INTO {ITEMS} VALUES (1, 'root', 'a.png', ?)", (NOW,)), + *statements, + ) + with closing(sqlite3.connect(path)) as raw: + self.assertIsNone(image_collections_schema_error(raw, check_values=False)) + error = image_collections_schema_error(raw) + self.assertIsNotNone(error) + self.assertIn("collections", error.lower()) + ok, reason = PromptModel.verify_database_file(path) + self.assertFalse(ok) + self.assertEqual(reason, error) + self.assertFalse( + PromptModel.verify_database_file(path, require_library_state=True)[0] + ) + + def test_a_valid_current_database_is_accepted(self): + self.require_module() + path = self.closed_current( + "valid.db", + ( + f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) " + "VALUES (1, ?, ?, ?)", + ("LoRa testing", key("LoRa testing"), NOW), + ), + (f"INSERT INTO {ITEMS} VALUES (1, 'root', ?, ?)", ("é" * 4096, NOW)), + ) + self.assertEqual( + PromptModel.verify_database_file(path, require_library_state=True), + (True, "ok"), + ) + + def test_restore_refuses_a_non_text_name(self): + self.assert_refused("int-name", f"UPDATE {COLLECTIONS} SET name = 5") + + def test_restore_refuses_an_empty_name(self): + self.assert_refused("empty", f"UPDATE {COLLECTIONS} SET name = '', name_key=''") + + def test_restore_refuses_a_name_over_100_characters(self): + name = "é" * 101 + self.assert_refused( + "long", + (f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", (name, key(name))), + ) + + def test_restore_refuses_a_control_character(self): + name = "a\x07b" + self.assert_refused( + "control", + (f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", (name, key(name))), + ) + + def test_restore_refuses_a_leading_or_trailing_space(self): + for name in (" Good", "Good "): + with self.subTest(name=name): + self.assert_refused( + f"padded{len(name)}{name[0] == ' '}", + ( + f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", + (name, key(name)), + ), + ) + + def test_restore_refuses_a_wrong_name_key(self): + self.assert_refused("key", f"UPDATE {COLLECTIONS} SET name_key = 'Good'") + + def test_restore_refuses_a_non_text_key(self): + # Column affinity turns a number into text; a blob stays a blob. + self.assert_refused("blob-key", f"UPDATE {ITEMS} SET relative_path = x'61'") + + def test_restore_refuses_a_key_over_4096_characters(self): + self.assert_refused( + "long-key", (f"UPDATE {ITEMS} SET root_key = ?", ("é" * 4097,)) + ) + + def test_restore_refuses_a_timestamp_over_64_bytes(self): + self.assert_refused( + "stamp", (f"UPDATE {COLLECTIONS} SET created_at = ?", ("t" * 65,)) + ) + self.assert_refused( + "item-stamp", (f"UPDATE {ITEMS} SET added_at = ?", ("t" * 65,)) + ) + + def test_restore_refuses_an_item_of_a_missing_collection(self): + self.assert_refused("orphan", f"INSERT INTO {ITEMS} VALUES (9, 'r', 'b', 'n')") + + def test_restore_refuses_a_non_integer_collection_id(self): + self.assert_refused("text-id", f"UPDATE {ITEMS} SET collection_id = 'one'") + + def test_restore_refuses_a_trigger_or_extra_index(self): + self.require_module() + for name, statement in ( + ( + "trigger", + f"CREATE TRIGGER t AFTER INSERT ON {ITEMS} BEGIN SELECT 1; END", + ), + ("extra", f"CREATE INDEX extra ON {COLLECTIONS}(name)"), + ): + with self.subTest(name=name): + path = self.closed_current(f"{name}.db", statement) + with closing(sqlite3.connect(path)) as raw: + self.assertIsNotNone(image_collections_schema_error(raw)) + self.assertFalse(PromptModel.verify_database_file(path)[0]) + + def test_deleting_a_collection_deletes_its_members(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + conn.execute( + f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) " + "VALUES (1, 'A', 'a', ?), (2, 'B', 'b', ?)", + (NOW, NOW), + ) + conn.executemany( + f"INSERT INTO {ITEMS} VALUES (?, 'root', ?, ?)", + ((1, "a.png", NOW), (1, "b.png", NOW), (2, "a.png", NOW)), + ) + conn.execute(f"DELETE FROM {COLLECTIONS} WHERE id = 1") + self.assertEqual( + [tuple(r) for r in conn.execute(f"SELECT collection_id FROM {ITEMS}")], + [(2,)], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_image_tag_jobs.py b/tests/test_image_tag_jobs.py index 0af74981..f593aaf5 100644 --- a/tests/test_image_tag_jobs.py +++ b/tests/test_image_tag_jobs.py @@ -884,6 +884,16 @@ class TestFilterStart(ImageTagJobTestCase): self.assertEqual((status, data["code"]), (400, "no-images")) self.assertEqual(self.api._jobs.list(), []) + async def test_filter_with_a_deleted_collection_is_a_404_without_a_job(self): + self.png("a.png") + + status, data = await self.start( + {"mode": "add", "filter": {"collection": "77"}, "tags": ["x"]} + ) + + self.assertEqual((status, data["code"]), (404, "collection-not-found")) + self.assertEqual(self.api._jobs.list(), []) + async def test_filter_over_the_cap_is_refused_without_a_job(self): from py.api import image_tag_jobs diff --git a/tests/test_pair_scan.py b/tests/test_pair_scan.py index 1597b1f9..2ee70151 100644 --- a/tests/test_pair_scan.py +++ b/tests/test_pair_scan.py @@ -115,7 +115,12 @@ class TestPairScan(AdminAPITestCase): self.assertEqual(second["pairs_added"], 0) def test_legacy_single_prompt_second_pass_not_reported_as_new_link(self): - meta = {"prompt": json.dumps({"1": graph()["1"]})} + # A sampler-less graph: one encoder wired into a guider (#252). + single = { + "1": graph()["1"], + "2": {"class_type": "BasicGuider", "inputs": {"conditioning": ["1", 0]}}, + } + meta = {"prompt": json.dumps(single)} path = self.output_dir / "old.png" self.assertEqual(self.api._ingest_scanned_file(path, meta), "added") self.assertEqual(self.api._ingest_scanned_file(path, meta), "unchanged") @@ -292,3 +297,155 @@ class TestPairScan(AdminAPITestCase): self.api.db.get_image_prompt_info(str(path))["workflow_data"], workflow ) self.assertEqual(self.api._ingest_scanned_file(path, metadata), "unchanged") + + +def flux_graph(text="a red fox", encoder="CLIPTextEncode"): + """A sampler-less Flux graph: the prompt reaches a guider's conditioning (#252).""" + return { + "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "flux.safetensors"}}, + "2": {"class_type": "DualCLIPLoader", "inputs": {"clip_name1": "t5"}}, + "6": {"class_type": encoder, "inputs": {"text": text, "clip": ["2", 0]}}, + "22": { + "class_type": "BasicGuider", + "inputs": {"model": ["1", 0], "conditioning": ["6", 0]}, + }, + "13": { + "class_type": "SamplerCustomAdvanced", + "inputs": {"guider": ["22", 0], "sigmas": ["17", 0]}, + }, + "17": {"class_type": "BasicScheduler", "inputs": {"steps": 20}}, + "9": {"class_type": "SaveImage", "inputs": {"images": ["13", 0]}}, + } + + +class TestSamplerLessScan(AdminAPITestCase): + """Graphs where no node takes a positive or negative input (#252).""" + + def _scanned(self, name, prompt_graph, workflow=None): + meta = {"prompt": json.dumps(prompt_graph)} + if workflow is not None: + meta["workflow"] = json.dumps(workflow) + return self.api._ingest_scanned_file(self.output_dir / name, meta) + + def _texts(self): + return [p["text"] for p in self.api.db.get_recent_prompts()["prompts"]] + + def test_the_one_wired_prompt_is_linked(self): + self.assertEqual(self._scanned("flux.png", flux_graph()), "added") + self.assertEqual(self._texts(), ["a red fox"]) + + def test_a_text_node_feeding_the_encoder_is_one_prompt(self): + g = flux_graph(text=["30", 0]) + g["30"] = {"class_type": "TextInput", "inputs": {"text": "a red fox"}} + self.assertEqual(self._scanned("textinput.png", g), "added") + self.assertEqual(self._texts(), ["a red fox"]) + + def test_a_display_node_that_feeds_nothing_is_never_the_prompt(self): + # A real scan stored a "=== GIT CLONE REPORT ===" display text this way. + g = flux_graph() + g["40"] = { + "class_type": "easy showAnything", + "inputs": {"text": "=== GIT CLONE REPORT ===\nTotal repositories: 1"}, + } + self.assertEqual(self._scanned("report.png", g), "added") + self.assertEqual(self._texts(), ["a red fox"]) + + def test_only_unconsumed_texts_store_nothing(self): + g = { + "40": {"class_type": "easy showAnything", "inputs": {"text": "report"}}, + "41": {"class_type": "Textbox", "inputs": {"text": "install note"}}, + } + self.assertEqual(self._scanned("notes.png", g), "found") + self.assertEqual(self._texts(), []) + + def test_two_wired_texts_are_ambiguous_and_stay_unlinked(self): + # A LoRA loader's text input is wired too: the scan stored "". + g = flux_graph() + g["50"] = { + "class_type": "Lora Loader (LoraManager)", + "inputs": { + "model": ["1", 0], + "clip": ["2", 0], + "text": "", + }, + } + g["22"]["inputs"]["model"] = ["50", 0] + g["6"]["inputs"]["clip"] = ["50", 1] + self.assertEqual(self._scanned("lora.png", g), "unresolved") + self.assertEqual(self._texts(), []) + + def test_an_untraceable_wired_text_makes_the_graph_ambiguous(self): + # The encoder's text comes from a node no rule reads: never guess the other. + g = flux_graph(text=["60", 0]) + g["60"] = {"class_type": "SomeLLMNode", "inputs": {"seed": 1}} + g["50"] = { + "class_type": "Lora Loader (LoraManager)", + "inputs": {"model": ["1", 0], "text": ""}, + } + g["22"]["inputs"]["model"] = ["50", 0] + self.assertEqual(self._scanned("llm.png", g), "unresolved") + self.assertEqual(self._texts(), []) + + def test_this_repositorys_prompt_node_is_authoritative(self): + g = flux_graph(text="a red fox", encoder="PromptManager") + g["50"] = { + "class_type": "Lora Loader (LoraManager)", + "inputs": {"model": ["1", 0], "text": ""}, + } + g["22"]["inputs"]["model"] = ["50", 0] + self.assertEqual(self._scanned("own.png", g), "added") + self.assertEqual(self._texts(), ["a red fox"]) + + def test_a_numeric_text_is_never_stored_as_a_prompt(self): + self.assertEqual(self._scanned("number.png", flux_graph(text="2048")), "found") + self.assertEqual(self._texts(), []) + + def test_workflow_widget_values_are_never_read(self): + # The API text is linked from a node no rule reads; the UI workflow + # still holds a stale widget value, which the old scan stored. + g = flux_graph(text=["70", 0]) + g["70"] = {"class_type": "PromptBuilder", "inputs": {"prompt": "x", "n": 2}} + workflow = { + "nodes": [ + {"id": 6, "type": "CLIPTextEncode", "widgets_values": ["stale text"]} + ] + } + self.assertEqual(self._scanned("stale.png", g, workflow), "unresolved") + self.assertEqual(self._texts(), []) + + def test_a_lora_loaders_trigger_word_output_is_not_its_text(self): + # The encoder reads the loader's trigger-word output (slot 2), not its + # text input: the prompt is unknown, never the "" syntax. + g = flux_graph(text=["50", 2]) + g["50"] = { + "class_type": "Lora Loader (LoraManager)", + "inputs": {"model": ["1", 0], "text": ""}, + } + g["22"]["inputs"]["model"] = ["50", 0] + self.assertEqual(self._scanned("trigger.png", g), "unresolved") + self.assertEqual(self._texts(), []) + + def test_a_seeded_text_node_holds_a_template_not_the_prompt(self): + g = flux_graph(text=["80", 0]) + g["80"] = { + "class_type": "RandomPrompt", + "inputs": {"text": "a {red|blue} fox", "seed": 5}, + } + self.assertEqual(self._scanned("random.png", g), "unresolved") + self.assertEqual(self._texts(), []) + + def test_a1111_parameters_beside_an_undecided_graph_are_read(self): + # Metadata savers write both; the old scan read the A1111 positive. + g = flux_graph() + g["50"] = { + "class_type": "Lora Loader (LoraManager)", + "inputs": {"model": ["1", 0], "text": ""}, + } + g["22"]["inputs"]["model"] = ["50", 0] + meta = { + "prompt": json.dumps(g), + "parameters": "a red fox\nNegative prompt: blur\nSteps: 20", + } + path = self.output_dir / "saver.png" + self.assertEqual(self.api._ingest_scanned_file(path, meta), "added") + self.assertEqual(self._texts(), ["a red fox"]) diff --git a/tests/test_restore_plans.py b/tests/test_restore_plans.py index 5ad20eaa..f5825169 100644 --- a/tests/test_restore_plans.py +++ b/tests/test_restore_plans.py @@ -49,6 +49,7 @@ SUMMARY_KEYS = { "image_links", "runs", "saved_views", + "collections", } diff --git a/tests/test_restore_stage.py b/tests/test_restore_stage.py index 3c863799..72ae806e 100644 --- a/tests/test_restore_stage.py +++ b/tests/test_restore_stage.py @@ -226,6 +226,7 @@ class TestStageAndPublish(StageTestCase): for key in ("pairs", "archived_pairs", "image_links", "runs"): self.assertEqual(staged.summary[key], 0, key) self.assertEqual(staged.summary["saved_views"], 0) + self.assertEqual(staged.summary["collections"], 0) self.assertEqual(staged.schema, {"migrated": False, "library_state": True}) self.assertEqual( staged.digest, @@ -239,6 +240,48 @@ class TestStageAndPublish(StageTestCase): staged.discard() +class TestCollectionsSummary(StageTestCase): + """#331: collections join the review's summary and the live fingerprint.""" + + def add_collection(self, db, name): + with db.model.get_connection() as conn: + conn.execute( + "INSERT INTO image_collections (name, name_key, created_at) " + "VALUES (?, ?, '2026-10-27T00:00:00+00:00')", + (name, name.casefold()), + ) + + def test_summary_counts_collections(self): + src_path = os.path.join(self.uploads, "with-collections.db") + src = PromptDatabase(src_path) + try: + self.save(src, "kept") + self.add_collection(src, "First") + self.add_collection(src, "Second") + finally: + src.close_all() + staged = self.db.model.stage_restore(src_path) + try: + self.assertEqual(staged.summary["collections"], 2) + finally: + staged.discard() + self.assertIn("collections", restore_stage.SUMMARY_KEYS) + + def test_adding_a_collection_changes_the_fingerprint(self): + conn = self.db.model.get_connection() + summary, before = restore_stage.live_fingerprint(conn) + self.assertEqual(summary["collections"], 0) + self.add_collection(self.db, "After the review") + summary, after = restore_stage.live_fingerprint(conn) + self.assertEqual(summary["collections"], 1) + self.assertNotEqual(after, before) + # A delete and a re-add with the same count still changes it (rowid). + with self.db.model.get_connection() as write: + write.execute("DELETE FROM image_collections") + self.add_collection(self.db, "Replacement") + self.assertNotEqual(restore_stage.live_fingerprint(conn)[1], after) + + class TestOlderBackups(StageTestCase): """Rollback-journal and older-schema backups restore exactly as before.""" diff --git a/tests/test_sampler_less_prompt.py b/tests/test_sampler_less_prompt.py new file mode 100644 index 00000000..62ee6fbe --- /dev/null +++ b/tests/test_sampler_less_prompt.py @@ -0,0 +1,112 @@ +"""``sampler_less_prompt``: the one wired text of a graph with no sampler (#252).""" + +import time +import unittest + +from utils.prompt_pair_resolver import sampler_less_prompt + + +def _guided(**nodes): + """``nodes`` plus a guider consuming node ``"1"``.""" + graph = {"9": {"class_type": "BasicGuider", "inputs": {"conditioning": ["1", 0]}}} + graph.update(nodes) + return graph + + +class TestSamplerLessPrompt(unittest.TestCase): + def test_not_a_graph_or_too_large_is_absent(self): + for graph in (None, [], {}, "text"): + self.assertEqual(sampler_less_prompt(graph)["state"], "absent") + big = {str(i): {"class_type": "X", "inputs": {}} for i in range(3)} + self.assertEqual(sampler_less_prompt(big, max_visits=2)["state"], "absent") + + def test_one_literal_text(self): + graph = _guided(**{"1": {"class_type": "Enc", "inputs": {"text": " a fox "}}}) + self.assertEqual( + sampler_less_prompt(graph), + {"state": "resolved", "text": "a fox", "texts": ["a fox"]}, + ) + + def test_a_number_literal_is_never_a_prompt(self): + for value in (2048, "2048", "1.5", " 7 "): + graph = _guided(**{"1": {"class_type": "Enc", "inputs": {"text": value}}}) + self.assertEqual(sampler_less_prompt(graph)["state"], "absent", value) + + def test_a_long_digit_run_with_a_letter_is_text_and_fast(self): + text = "9" * 50_000 + "x" + graph = _guided(**{"1": {"class_type": "Enc", "inputs": {"text": text}}}) + started = time.monotonic() + self.assertEqual(sampler_less_prompt(graph)["text"], text) + self.assertLess(time.monotonic() - started, 1.0) + + def test_number_like_words_are_prompts(self): + for value in ("nan", "inf", "Infinity", "1_000"): + graph = _guided(**{"1": {"class_type": "Enc", "inputs": {"text": value}}}) + self.assertEqual(sampler_less_prompt(graph)["text"], value) + + def test_an_own_node_with_unreadable_text_beside_another_text_is_ambiguous(self): + graph = _guided( + **{ + "1": {"class_type": "PromptManager", "inputs": {"text": ["5", 0]}}, + "5": {"class_type": "LLM", "inputs": {"seed": 3}}, + "2": {"class_type": "Enc", "inputs": {"text": ""}}, + "3": {"class_type": "Sink", "inputs": {"c": ["2", 0]}}, + } + ) + self.assertEqual(sampler_less_prompt(graph)["state"], "ambiguous") + + def test_unhashable_class_type_is_an_unknown_node_not_a_crash(self): + graph = _guided(**{"1": {"class_type": ["x"], "inputs": {"text": "a fox"}}}) + self.assertEqual(sampler_less_prompt(graph)["text"], "a fox") + + def test_own_negative_node_is_not_authoritative(self): + graph = _guided( + **{ + "1": {"class_type": "Enc", "inputs": {"text": "a fox"}}, + "2": { + "class_type": "PromptManagerNegative", + "inputs": {"text": "blur"}, + }, + "3": {"class_type": "Sink", "inputs": {"c": ["2", 0]}}, + } + ) + found = sampler_less_prompt(graph) + self.assertEqual(found["state"], "ambiguous") + self.assertEqual(found["texts"], ["a fox", "blur"]) + + def test_own_node_text_wins_over_an_untraceable_text(self): + graph = _guided( + **{ + "1": {"class_type": "PromptManagerText", "inputs": {"text": "a fox"}}, + "2": {"class_type": "Enc", "inputs": {"text": ["5", 0]}}, + "5": {"class_type": "LLM", "inputs": {"seed": 3}}, + "6": {"class_type": "Sink", "inputs": {"c": ["2", 0]}}, + } + ) + self.assertEqual(sampler_less_prompt(graph)["text"], "a fox") + + def test_two_own_texts_are_ambiguous(self): + graph = _guided( + **{ + "1": {"class_type": "PromptManager", "inputs": {"text": "a fox"}}, + "2": {"class_type": "PromptManager", "inputs": {"text": "a cat"}}, + "3": {"class_type": "Sink", "inputs": {"c": ["2", 0]}}, + } + ) + self.assertEqual(sampler_less_prompt(graph)["state"], "ambiguous") + + def test_only_an_untraceable_text_is_ambiguous_with_nothing_to_show(self): + graph = _guided( + **{ + "1": {"class_type": "Enc", "inputs": {"text": ["5", 0]}}, + "5": {"class_type": "LLM", "inputs": {"seed": 3}}, + } + ) + self.assertEqual( + sampler_less_prompt(graph), + {"state": "ambiguous", "text": None, "texts": []}, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/prompt_pair_resolver.py b/utils/prompt_pair_resolver.py index b45c948c..15200481 100644 --- a/utils/prompt_pair_resolver.py +++ b/utils/prompt_pair_resolver.py @@ -2578,3 +2578,115 @@ def resolve_text_input( return resolver.text(inputs.get(input_name), "string")[0] except _Unresolved: return None + + +_NUMBER = re.compile(r"[+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?") + + +def _numeric(text: str) -> bool: + """Whether ``text`` is a plain number (a widget value such as ``2048``).""" + return _NUMBER.fullmatch(text) is not None + + +def sampler_less_prompt(graph: Any, *, max_visits: int = 500) -> Dict[str, Any]: + """The positive prompt of an API graph with no sampler pair, if one (#252). + + For graphs ``_no_sampler`` finds ``absent`` (no node takes a ``positive`` + or ``negative`` input, such as a Flux guider or a video wrapper). Only the + executed API graph is read, never the UI workflow. A candidate is a node + with a ``text`` input whose output another node consumes: a text that + feeds nothing did not make the image. Its text is traced as a STRING + (``resolve_text_input``); a linked text no rule reads is unknown unless + it is built from a candidate's text (``_covered``). Blank and + numeric texts are never prompts, and a node with a seed input holds a + template, not the prompt (unknown). This repository's own positive-text + nodes are authoritative: when they give a text, only they count. + + Returns ``{"state", "text", "texts"}``: ``resolved`` with the one distinct + text; ``ambiguous`` when several texts or an unknown one could be the + prompt (``texts`` lists the readable ones); ``absent`` when none. + """ + absent = {"state": "absent", "text": None, "texts": []} + if not isinstance(graph, dict) or not graph or len(graph) > max_visits: + return absent + nodes = {str(key): value for key, value in graph.items()} + consumed = { + str(value[0]) + for node in nodes.values() + if isinstance(node, dict) and isinstance(node.get("inputs"), dict) + for value in node["inputs"].values() + if _link(value) + } + own: list = [] + other: list = [] + unknown = False + for node_id, node in nodes.items(): + inputs = node.get("inputs") if isinstance(node, dict) else None + if node_id not in consumed or not _has_text_input(node): + continue + if _seeded(inputs): + # Generated at run time (a wildcard or a language model): its + # text input is a template, not the prompt that ran. + unknown = True + continue + value = inputs["text"] + text = resolve_text_input(nodes, node_id, "text", max_visits=max_visits) + if text is None: + # A literal that is not a string (a number) is never a prompt. + if _link(value) and not _covered(nodes, value): + unknown = True + continue + text = text.strip() + if not text or _numeric(text): + continue + kind = node.get("class_type") + if isinstance(kind, str) and kind in _OWN_CLASSES and "Negative" not in kind: + own.append(text) + else: + other.append(text) + texts = list(dict.fromkeys(own or other)) + if len(texts) == 1 and (own or not unknown): + return {"state": "resolved", "text": texts[0], "texts": texts} + if texts or unknown: + return {"state": "ambiguous", "text": None, "texts": texts} + return absent + + +def _has_text_input(node: Any) -> bool: + inputs = node.get("inputs") if isinstance(node, dict) else None + return isinstance(inputs, dict) and "text" in inputs + + +def _seeded(inputs: dict) -> bool: + """A seed input: the node's string output is a run-time value (the + resolver's ``runtime-value`` rule).""" + return any(isinstance(name, str) and name.endswith("seed") for name in inputs) + + +def _covered(nodes: dict, link: list) -> bool: + """Whether an unreadable text link is built from a candidate's text. + + True when a node upstream of ``link`` is read through its first output + (slot 0, the legacy encoder rule's text) and has an unseeded ``text`` + input: that candidate is read on its own, and whatever else is joined to + it at run time (a caption, a language model) is no stored prompt to + choose. Another output of such a node (a loader's trigger words) is not + its text. False when the text comes only from inputs no candidate holds + (a builder's ``prompt`` field): the prompt is unknown then. + """ + seen, stack = set(), [(str(link[0]), link[1])] + while stack: + node_id, slot = stack.pop() + if (node_id, slot) in seen: + continue + seen.add((node_id, slot)) + node = nodes.get(node_id) + inputs = node.get("inputs") if isinstance(node, dict) else None + if not isinstance(inputs, dict): + continue + if "text" in inputs and slot == 0 and not _seeded(inputs): + return True + stack.extend( + (str(value[0]), value[1]) for value in inputs.values() if _link(value) + ) + return False