diff --git a/database/duplicate_copies.py b/database/duplicate_copies.py index 22129786..5399b7e1 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 own image category (#332), copied to the kept copy when that copy has @@ -54,6 +56,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=?" @@ -77,6 +88,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, @@ -140,14 +153,17 @@ 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``; ``categories``: ``(removed file key, kept file - key)`` pairs of the image categories (#332). Returns ``{moved, - merged, favorites, tags, review, categories}`` (``categories``: the - image categories copied to the kept copy). + 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``; ``categories``: ``(removed file + key, kept file key)`` pairs of the image categories (#332). Returns + ``{moved, merged, favorites, tags, review, collections, categories}`` + (``collections``: the memberships inserted on the kept copy; + ``categories``: the image categories copied to 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"]): @@ -180,6 +196,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) counts["categories"] = sum( copy_image_category(conn, removed, kept) @@ -224,3 +241,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..9eeb4494 --- /dev/null +++ b/database/image_collections.py @@ -0,0 +1,490 @@ +"""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, +but never recomputes ``name_key`` (case tables differ between Python +versions). +""" + +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 +# A name key's byte bound: case folding expands a character to at most +# three, each at most 4 UTF-8 bytes. +_MAX_KEY_BYTES = MAX_NAME_LENGTH * 3 * 4 +_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 _has_control(text: str) -> bool: + """Any Unicode control character (category Cc: C0, DEL and C1).""" + return any(unicodedata.category(char) == "Cc" for char in text) + + +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 _has_control(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} " + "OR length(CAST(name_key AS BLOB)) = 0 " + f"OR length(CAST(name_key AS BLOB)) > {_MAX_KEY_BYTES} " + f"OR length(CAST(created_at AS BLOB)) > {_TIMESTAMP_BYTES} LIMIT 1" + ).fetchone(): + return "Image collections contains invalid values" + # The key is not recomputed: Unicode case tables differ between Python + # versions, and a key written by another version must stay restorable. + # Its UNIQUE constraint keeps names distinct; writes compute it. + for (name,) in conn.execute(f"SELECT name FROM {COLLECTIONS_TABLE}"): # nosec B608 + if ( + not 1 <= len(name) <= MAX_NAME_LENGTH + or _has_control(name) + or name != name.strip() + ): + 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_count_snapshot(self, root_keys: list) -> tuple[dict, dict]: + """``({root key: {id: relative keys}}, {id: stored rows})`` read in + one read transaction, so a write between the reads cannot skew them. + """ + conn = self.model.get_connection() + if conn.in_transaction: + raise RuntimeError("Collection counts require their own transaction") + conn.execute("BEGIN") + try: + members = {} + for root_key in dict.fromkeys(root_keys): + found: dict[int, set[str]] = {} + for collection_id, path in conn.execute(ROOT_MEMBERS_SQL, (root_key,)): + found.setdefault(collection_id, set()).add(path) + members[root_key] = found + stored = {cid: count for cid, count in conn.execute(COUNTS_SQL)} + finally: + conn.rollback() # it only read + return members, stored + + 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 25e98e44..805d665f 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, @@ -338,6 +342,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). @@ -1159,8 +1164,8 @@ class PromptModel: path: File to inspect require_library_state: Also require the tables every current database has (Library state, image favorites, image - categories, saved views, image annotations and the review - queue), as a migrated + categories, saved views, 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 @@ -1234,6 +1239,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 3087d615..a40651b4 100644 --- a/database/operations.py +++ b/database/operations.py @@ -18,6 +18,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 @@ -453,6 +454,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 0d24ec0c..76fd0bb1 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 @@ -315,6 +316,7 @@ class PromptManagerAPI( ReviewRecordingMixin, ReviewEventsMixin, SavedViewRoutesMixin, + CollectionRoutesMixin, HistoryRoutesMixin, PromptRoutesMixin, ImageRoutesMixin, @@ -686,6 +688,7 @@ class PromptManagerAPI( self._register_review_routes(routes) self._register_review_event_routes(routes) self._register_saved_view_routes(routes) + self._register_collection_routes(routes) self._register_history_routes(routes) self._register_prompt_routes(routes) self._register_image_routes(routes) diff --git a/py/api/collections.py b/py/api/collections.py new file mode 100644 index 00000000..bc8046ec --- /dev/null +++ b/py/api/collections.py @@ -0,0 +1,350 @@ +"""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"[1-9][0-9]{0,8}") +_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: + 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) + stored_rows = None if stored is None else stored.get(collection["id"], 0) + rows.append( + { + **collection, + "count": count, + # Clamped: never below zero, whatever the snapshot saw. + "missing": None if count is None else max(0, stored_rows - 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. + + The scans come from the gallery cache; the root keys (realpath), the + listings and one read snapshot of every member and stored count run + in the executor, so a write between two reads cannot skew them. + """ + roots = list(self._get_all_output_dirs()) + root_keys = await self._run_in_executor( + lambda: [_favorite_root_key(root) for root in roots] + ) + distinct = {} # one folder configured twice counts once + for root, root_key in zip(roots, root_keys): + distinct.setdefault(root_key, root) + scans = { + root_key: await self._get_gallery_files(root) + for root_key, root in distinct.items() + } + return await self._run_in_executor(self._count_members, distinct, scans) + + def _count_members(self, roots: dict, scans: dict) -> tuple[dict, dict]: + """Blocking: match one snapshot's members against each listing.""" + members, stored = self.db.collection_count_snapshot(list(roots)) + listed: dict[int, int] = {} + for root_key, root in roots.items(): + listing = self._favorite_listing(root, scans[root_key]) + for collection_id, keys in members.get(root_key, {}).items(): + found = sum(1 for key in keys if key in listing) + listed[collection_id] = listed.get(collection_id, 0) + found + 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 len(query) != 3 or set(query) != {"root_index", "root_scope", "path"}: + # len() counts a repeated key (?path=a&path=b) twice. + return bad_request("Query must name root_index, root_scope and path once") + 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()) + target = { + "root_index": int(raw_index), + "root_scope": query["root_scope"], + "path": query["path"], + } + root_key, key = await self._run_in_executor( + self._collection_key, target, 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 = await self._run_in_executor(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 = await self._run_in_executor(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. + + Blocking (realpath per root, once per request): run in the executor. + 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, sources = {}, {} + for item in items: + index = item.get("root_index") if isinstance(item, dict) else None + root_key, key = self._collection_key(item, roots, sources) + 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, sources: dict + ) -> tuple[str, str]: + """The stored keys of one target, or raise its refusal (blocking). + + *sources* caches ``{index: (scope, root key)}`` for one request, so + each root's realpath runs once however many items name it. + """ + 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): + raise _Refused(_coded(409, "source-changed", SOURCES_CHANGED)) + if index not in sources: + sources[index] = ( + self._gallery_root_scope(roots[index], index), + _favorite_root_key(roots[index]), + ) + current, root_key = sources[index] + if not hmac.compare_digest(scope, current): + raise _Refused(_coded(409, "source-changed", SOURCES_CHANGED)) + 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 dc7f511c..34e15ed4 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 + ) totals["categories"] = 0 # image categories kept on the kept copies (#332) for done, target in enumerate(targets, start=1): if failure is not None: @@ -632,6 +634,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 0e197aa4..7157e994 100644 --- a/py/api/duplicate_spellings.py +++ b/py/api/duplicate_spellings.py @@ -104,7 +104,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 6baafd4c..61c79422 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, @@ -689,6 +691,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 e030c973..614287b9 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 @@ -81,6 +82,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 @@ -110,6 +112,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: @@ -1110,8 +1127,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() @@ -1125,8 +1143,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") @@ -1149,7 +1169,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, } @@ -1167,6 +1188,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 @@ -1223,7 +1246,7 @@ def _admitted_entries( entries: list, filters: dict, links: dict | None, - favorites: frozenset | None = None, + members: frozenset | None = None, image_tags: dict | None = None, own_categories: dict | None = None, ): @@ -1236,11 +1259,11 @@ def _admitted_entries( admits a file whose effective category equals it (#332 MC5): its own category (*own_categories*, ``{normalized path: category}``) when it has one, else the inherited one, so an unlinked file with its own category is - admitted. *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. + admitted. *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"] category = filters["category"] @@ -1258,7 +1281,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) @@ -1286,7 +1309,7 @@ def _filtered_page( size: int, filters: dict, links: dict | None, - favorites: frozenset | None = None, + members: frozenset | None = None, image_tags: dict | None = None, own_categories: dict | None = None, ) -> tuple[int, list, int | None]: @@ -1298,7 +1321,7 @@ def _filtered_page( """ total, page, positives = 0, [], set() for position, entry, match in _admitted_entries( - entries, filters, links, favorites, image_tags, own_categories + entries, filters, links, members, image_tags, own_categories ): total += 1 if match is not None: @@ -2002,6 +2025,7 @@ class ImageRoutesMixin: filters["category"], filters["tags"], filters["favorites"], + filters["collection"], ] ).encode() ).hexdigest() @@ -2012,17 +2036,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 @@ -2282,9 +2308,11 @@ class ImageRoutesMixin: ) -> tuple: """The ordered entries and what admits them (blocking; run in executor). - Returns ``(entries, key, links, image_tags, favorites, own)``: the + Returns ``(entries, key, links, image_tags, members, own)``: the listing in the filters' order, its keyset key, and the link, image-tag, - favorite and own-category maps _admitted_entries reads. One SELECT + member and own-category 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*; under a category it keeps only links whose raw recipe or positive category is that value (the inherited one is decided in Python), and @@ -2368,10 +2396,14 @@ 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), + own, ) - return entries, key, links, image_tags, favorites, own def _own_category_links(self, links: dict, own: dict, filters: dict) -> dict: """*links* plus the links of files whose own category is the filter's. @@ -2423,7 +2455,7 @@ class ImageRoutesMixin: or bool(filters["tags"]) or order == "rating" ) - entries, key, links, image_tags, favorites, own = self._workspace_admission( + entries, key, links, image_tags, members, own = self._workspace_admission( listings, filters, read_links=needed or after is None ) start = 0 @@ -2434,7 +2466,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, own + entries, start, limit + 1, filters, links, members, image_tags, own ) if not needed and after is not None: positive_count = _COUNT_OMITTED @@ -2456,12 +2488,12 @@ class ImageRoutesMixin: The links are read strictly: the positives must be the listing's own. """ - entries, _key, links, image_tags, favorites, own = self._workspace_admission( + entries, _key, links, image_tags, members, own = 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, own + entries, filters, links, members, image_tags, own ): if len(targets) == cap: return ListingSelection(targets, list(positives), unlinked, True) @@ -2483,9 +2515,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) @@ -2495,11 +2529,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 ) @@ -3684,25 +3718,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/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/full-inspector.spec.mjs b/tests/browser/full-inspector.spec.mjs index e4efc5a7..bd530bed 100644 --- a/tests/browser/full-inspector.spec.mjs +++ b/tests/browser/full-inspector.spec.mjs @@ -676,6 +676,9 @@ test('collection viewer uses real public transforms and resets on navigation',as async function slideshowFixture(page,request,preferences={videoAutoplay:false,videoMute:true,videoLoop:false}) { await expect(await request.post('/__test__/reset',{data:{full_inspector:true,include_video:true}})).toBeOK(); await page.goto('/prompt_manager/workspace#/images'); + // Install the fake clock before the fixture wraps setTimeout: the wrapped timer (which the viewer session captures) + // must be the fake one, so page.clock.runFor alone decides when the 3500 ms slideshow deadline fires. + await page.clock.install(); await page.evaluate(async preferences=>{ if(preferences) localStorage.setItem('gallerySettings',JSON.stringify(preferences)); const [{createMediaViewer},{createInspectionController},{createScope},{createPromptSender}]=await Promise.all([ @@ -695,7 +698,7 @@ async function slideshowFixture(page,request,preferences={videoAutoplay:false,vi },preferences); await expect(mediaDialog(page).getByRole('button',{name:'Zoom in',exact:true})).toHaveAttribute('aria-disabled','false'); await expect(mediaDialog(page).locator('#viewer-details > div > p.small')).toHaveText('complete.png'); - await page.clock.install();await page.clock.pauseAt(new Date(Date.now()+1000)); + await page.clock.pauseAt(new Date(Date.now()+1000)); } const slideshowButton=page=>mediaDialog(page).locator('#viewer-toolbar').getByRole('button',{name:/^(Play|Pause|Stop) slideshow$/}); const timerCount=page=>page.evaluate(()=>window.slideshowTimers.size); diff --git a/tests/browser/server.py b/tests/browser/server.py index 160f4b0a..dab18015 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/fixtures/legacy_contract_snapshot.json b/tests/fixtures/legacy_contract_snapshot.json index eb2085cd..27568173 100644 --- a/tests/fixtures/legacy_contract_snapshot.json +++ b/tests/fixtures/legacy_contract_snapshot.json @@ -1,5 +1,7 @@ { "routes": [ + "DELETE /prompt_manager/collections/items", + "DELETE /prompt_manager/collections/{collection_id}", "DELETE /prompt_manager/delete/{prompt_id}", "DELETE /prompt_manager/image-annotations/capture/{token}", "DELETE /prompt_manager/images/category", @@ -15,6 +17,8 @@ "GET /prompt_manager/backup", "GET /prompt_manager/backup/status", "GET /prompt_manager/categories", + "GET /prompt_manager/collections", + "GET /prompt_manager/collections/membership", "GET /prompt_manager/diagnostics", "GET /prompt_manager/export", "GET /prompt_manager/gallery", @@ -92,6 +96,7 @@ "GET /prompt_manager/vocabulary/tags/{tag_id}/impact", "GET /prompt_manager/web", "GET /prompt_manager/workspace", + "PATCH /prompt_manager/collections/{collection_id}", "PATCH /prompt_manager/library/{positive_id}", "PATCH /prompt_manager/pairs/{pair_id}", "PATCH /prompt_manager/saved-views/{view_id}", @@ -106,6 +111,7 @@ "POST /prompt_manager/bulk/delete", "POST /prompt_manager/bulk/tags", "POST /prompt_manager/cleanup", + "POST /prompt_manager/collections/items", "POST /prompt_manager/delete_duplicate_images", "POST /prompt_manager/diagnostics/test-link", "POST /prompt_manager/gallery/rescan", 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_collections.py b/tests/test_api_collections.py new file mode 100644 index 00000000..daa0e53b --- /dev/null +++ b/tests/test_api_collections.py @@ -0,0 +1,601 @@ +"""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") + + +class TestReviewFindings(CollectionRoutesTestCase): + """PR #400 review: loop work, counts snapshot, id and query rules.""" + + async def test_no_realpath_runs_on_the_event_loop(self): + import os + import threading + + names = [f"img{n}.png" for n in range(20)] + self.make(*names) + await self.call("GET", BASE) # warm the listing cache + loop_thread = threading.current_thread() + on_loop = [] + real = os.path.realpath + + def watched(*args, **kwargs): + if threading.current_thread() is loop_thread: + on_loop.append(args[0]) + return real(*args, **kwargs) + + items = [self.item(name) for name in names] + cid = None + with patch("os.path.realpath", side_effect=watched): + made = await self.add(name="Set", items=items) + cid = made["collection"]["id"] + await self.call("DELETE", ITEMS, {"collection_id": cid, "items": items[:5]}) + await self.call( + "GET", + MEMBERSHIP, + root_index="0", + root_scope=items[0]["root_scope"], + path="img9.png", + ) + await self.call("GET", BASE) + self.assertEqual(on_loop, []) + + async def test_missing_is_never_negative(self): + self.make("a.png", "b.png") + made = await self.add( + name="Set", items=[self.item("a.png"), self.item("b.png")] + ) + cid = made["collection"]["id"] + real = self.api.db.collection_count_snapshot + + def raced(root_keys): + members, stored = real(root_keys) + return members, {cid: 1} # a remove landed between two reads + + with patch.object(self.api.db, "collection_count_snapshot", side_effect=raced): + listed = await self.listed() + self.assertEqual((listed[0]["count"], listed[0]["missing"]), (2, 0)) + + async def test_path_ids_follow_the_listing_rule(self): + self.make("a.png") + made = await self.add(name="Set", items=[self.item("a.png")]) + cid = made["collection"]["id"] + for raw in ("0", "007", f"0{cid}"): + with self.subTest(raw=raw): + await self.call("PATCH", f"{BASE}/{raw}", {"name": "x"}, 400) + await self.call("DELETE", f"{BASE}/{raw}", None, 400) + self.assertEqual(len(await self.listed()), 1) + + async def test_membership_refuses_repeated_query_keys(self): + self.make("a.png", "b.png") + response = await self.client.get( + MEMBERSHIP + + f"?root_index=0&root_scope={self.scope()}&path=a.png&path=b.png" + ) + self.assertEqual(response.status, 400) + response = await self.client.get( + MEMBERSHIP + + f"?root_index=0&root_index=0&root_scope={self.scope()}&path=a.png" + ) + self.assertEqual(response.status, 400) + + async def test_deleting_a_collection_leaves_everything_else(self): + from py.api.images import _favorite_root_key + + try: + from tests.test_annotation_schema import insert_group_rows, table_rows + except ImportError: + from test_annotation_schema import insert_group_rows, table_rows + + self.make("a.png") + db = self.api.db + db.set_image_favorite(_favorite_root_key(self.output_dir), "a.png") + prompt = self._save_prompt("Linked") + db.link_image_to_prompt(prompt, str(self.output_dir / "a.png")) + with db.model.get_connection() as conn: + insert_group_rows(conn, locator=str(self.output_dir / "a.png")) + made = await self.add(name="Set", items=[self.item("a.png")]) + tables = ( + "image_favorites", + "generated_images", + "prompts", + "tags", + "media_annotation_files", + "media_annotation_bytes", + "media_image_tags", + ) + before = table_rows(db.model.get_connection(), tables) + file_bytes = (self.output_dir / "a.png").read_bytes() + await self.call("DELETE", f"{BASE}/{made['collection']['id']}") + self.assertEqual(table_rows(db.model.get_connection(), tables), before) + self.assertEqual((self.output_dir / "a.png").read_bytes(), file_bytes) + self.assertTrue(all(before[t] for t in tables)) diff --git a/tests/test_api_duplicate_files.py b/tests/test_api_duplicate_files.py index 9c78ba4a..9350a491 100644 --- a/tests/test_api_duplicate_files.py +++ b/tests/test_api_duplicate_files.py @@ -693,6 +693,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") @@ -849,8 +882,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 604d4ade..cd998722 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -2550,3 +2550,207 @@ 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_collection_filter_combines_with_an_images_own_category(self): + """#331 with #332: a member is admitted by its effective category.""" + from py.api.images import image_category_key + + night = self._save_prompt("Neon zebra", category="Night") + files = { + name: self.media(name, 100 + index) + for index, name in enumerate( + ("linked.png", "overridden.png", "own.png", "outside.png") + ) + } + for name in ("linked.png", "overridden.png"): + self.api.db.link_image_to_prompt(night, str(files[name])) + sources = self.api._get_all_output_dirs() + + def own(name, value): + key = image_category_key(sources, files[name]) + self.api.db.set_image_categories([key], value) + + own("overridden.png", "Portraits") + own("own.png", "Night") # unlinked: only its own category + own("outside.png", "Night") # not a member + members = [(self.output_dir, name) for name in files if name != "outside.png"] + grids = str(self.collect("Grids", *members)) + for category, expected in ( + ("Night", ["linked.png", "own.png"]), + ("Portraits", ["overridden.png"]), + ): + with self.subTest(category=category): + response, data = await self.get( + collection=grids, category=category, sort="name" + ) + 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_duplicate_image_categories.py b/tests/test_duplicate_image_categories.py index 79452a68..97492454 100644 --- a/tests/test_duplicate_image_categories.py +++ b/tests/test_duplicate_image_categories.py @@ -149,3 +149,36 @@ class TestDuplicateCleanupKeepsCategories(DuplicateFilesTestCase): kept = image_category_key([link], link / "a.png") self.assertFalse((self.output_dir / "b.png").exists()) self.assertEqual(self.api.db.image_category(*kept), "Portraits") + + async def test_a_copy_with_a_collection_and_a_category_moves_both(self): + # #331 + #332: one removed copy carries both records; one cleanup + # moves its membership and copies its category onto the kept copy. + self.trio() + db = self.api.db + root_key = duplicate_files.favorite_root_key(self.output_dir) + b = self.output_dir / "b.png" + grids = db.add_to_collection(name="LoRA grids", items=[(root_key, "b.png")]) + grids = grids["collection"]["id"] + self.set_category(b, "Portraits") + seen = {} + real_log = duplicate_files.DuplicateFilesMixin._log_moved_records + + def log_moved_records(job, totals): + seen.update(totals) + real_log(job, totals) + + with patch.object( + duplicate_files.DuplicateFilesMixin, + "_log_moved_records", + staticmethod(log_moved_records), + ): + cleanup = await self.run_cleanup() + self.assertFalse(b.exists()) + self.assertEqual(db.collection_members(grids, root_key), {"a.png"}) + self.assertEqual(db.file_collections(root_key, "b.png"), []) + self.assertEqual(self.category(self.output_dir / "a.png"), "Portraits") + self.assertEqual(seen["collections"], 1) + self.assertEqual(seen["categories"], 1) + log = "\n".join(cleanup.log) + self.assertIn("Moved 1 collection membership to the kept copies", log) + self.assertIn("Kept 1 image category on the kept copies", log) diff --git a/tests/test_image_collections.py b/tests/test_image_collections.py new file mode 100644 index 00000000..83307b02 --- /dev/null +++ b/tests/test_image_collections.py @@ -0,0 +1,268 @@ +"""Collections (#331): the database operations.""" + +import sqlite3 +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +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", "c1\x81", "del\x7f", 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_a_failure_after_the_membership_move_rolls_it_back(self): + self.require() + made = self.add([(ROOT, "b.png")], name="Grids")["collection"] + move = { + "removed": [], + "kept": [], + "targets": {}, + "kept_path": "/real/output/a.png", + "favorites": [(ROOT, "b.png", ROOT, "a.png")], + "removed_locator": "/real/output/b.png", + "kept_locator": "/real/output/a.png", + "sha256": "0" * 64, + "kept_size": 1, + "kept_mtime_ns": 1, + } + with patch.object( + type(self.db), "_move_copy_tags", side_effect=RuntimeError("boom") + ): + with self.assertRaises(RuntimeError): + self.db.move_copy_records(move) + self.assertEqual(self.db.file_collections(ROOT, "b.png"), [made["id"]]) + self.assertEqual(self.db.file_collections(ROOT, "a.png"), []) + # Without the failure the same move succeeds. + self.assertEqual(self.db.move_copy_records(move)["collections"], 1) + self.assertEqual(self.db.file_collections(ROOT, "a.png"), [made["id"]]) + + def test_count_snapshot_reads_in_one_transaction(self): + self.require() + made = self.add([(ROOT, "a.png"), (OTHER, "b.png")], name="Set")["collection"] + statements = [] + conn = self.db.model.get_connection() + conn.set_trace_callback(statements.append) + try: + members, stored = self.db.collection_count_snapshot([ROOT, OTHER, ROOT]) + finally: + conn.set_trace_callback(None) + self.assertEqual( + members, {ROOT: {made["id"]: {"a.png"}}, OTHER: {made["id"]: {"b.png"}}} + ) + self.assertEqual(stored, {made["id"]: 2}) + self.assertEqual(statements[0], "BEGIN") + self.assertEqual(statements[-1], "ROLLBACK") + self.assertFalse(conn.in_transaction) + + 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..df89faf7 --- /dev/null +++ b/tests/test_image_collections_schema.py @@ -0,0 +1,521 @@ +"""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): + # Column affinity turns a number into text; a blob stays a blob. + self.assert_refused("blob-name", f"UPDATE {COLLECTIONS} SET name = x'61'") + + 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_accepts_a_name_key_from_other_case_tables(self): + # Python versions fold some characters differently; a key written + # by another version must never make a backup unrestorable. + self.require_module() + path = self.closed_current( + "other-key.db", + ( + f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) " + "VALUES (1, 'Good', 'Good', ?)", + (NOW,), + ), + ) + self.assertEqual( + PromptModel.verify_database_file(path, require_library_state=True), + (True, "ok"), + ) + + def test_restore_refuses_an_empty_or_non_text_name_key(self): + self.assert_refused("empty-key", f"UPDATE {COLLECTIONS} SET name_key = ''") + self.assert_refused( + "blob-name-key", f"UPDATE {COLLECTIONS} SET name_key = x'61'" + ) + self.assert_refused( + "long-name-key", + (f"UPDATE {COLLECTIONS} SET name_key = ?", ("k" * 1201,)), + ) + + def test_restore_refuses_a_c1_control_character(self): + for name in ("a\x81b", "a\x7fb"): + with self.subTest(name=repr(name)): + self.assert_refused( + f"c1-{ord(name[1])}", + ( + f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", + (name, key(name)), + ), + ) + + 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_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."""