Files
ComfyAssets-ComfyUI_PromptM…/database/image_collections.py
T
Vito 66ca988d2a feat(images): delete selected images, from the library or to the trash (#480) (#488)
* feat(trash): a freedesktop trash backend that renames without replacing or refuses (#480)

* feat(images): remove a file's records and hide it with a removal mark (#480)

Deletes every record of one file under every recorded spelling (image links,
tags, auto-tag marks, pending review items, category, collection memberships,
favorite) in one transaction. A Needs review item that is kept keeps its
annotation rows (ON DELETE RESTRICT). The image_removals table hides an exact
file (locator, size, mtime_ns) until it changes; linking skips a marked file,
and an explicit tag, category, favorite or collection write shows it again.

* feat(images): review and apply Delete images as a reviewed plan job (#480)

POST /prompt_manager/images/delete/preview reviews both modes; the
maintenance apply takes {token, mode}. Library mode removes the records and
writes the removal mark; trash mode moves the file first and removes its
records only after a successful move. The workspace listing, the output scan,
the duplicate find and Tag output images skip marked files.

* feat(images): Delete… in the Images selection bar, inspector and viewer (#480)

The concept's Remove links… place, D159: one review per selection that asks
each time (nothing preselected) whether to remove the files from the library
only or also move them to the trash. The browser test server uses a fake
trash.

* test(images): portable delete tests, the native no-trash proof, D159 in TESTING.md (#480)

* perf(images): cache the removal marks so a listing page reads no extra table (#480)

The marks reload only when this process writes or clears one, or a restore
replaces the database. The route snapshot gains the preview route.

* test(images): the full inspector specs expect D159's Delete… in the panel and the viewer (#480)

* fix(images): the delete review starts nothing without a chosen mode, and dims its primary until then (#480)

An aria-disabled primary ignores every activation, a chosen one starts once,
and the apply helper refuses a missing mode. Tests cover a forced click, a
double click, Enter and Space on the primary, Enter in the dialog and a
submit event, and the server refusing an apply without an exact mode.
2026-10-29 19:40:26 +00:00

493 lines
20 KiB
Python

"""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 .image_removal_schema import clear_marks_of_keys
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
)
clear_marks_of_keys(conn, pairs) # #480: a collected file is listed
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