From 420c8bcf335763c2574701f3fbf20fda5d615785 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Thu, 29 Oct 2026 10:03:48 -0700 Subject: [PATCH] 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. --- database/annotations.py | 4 + database/image_categories.py | 2 + database/image_collections.py | 2 + database/image_favorites.py | 3 + database/image_removal.py | 305 ++++++++++++++++++++++++++++++ database/image_removal_schema.py | 191 +++++++++++++++++++ database/models.py | 8 + database/operations.py | 7 + tests/test_image_removal.py | 313 +++++++++++++++++++++++++++++++ 9 files changed, 835 insertions(+) create mode 100644 database/image_removal.py create mode 100644 database/image_removal_schema.py create mode 100644 tests/test_image_removal.py diff --git a/database/annotations.py b/database/annotations.py index e1bd394e..fbf755a6 100644 --- a/database/annotations.py +++ b/database/annotations.py @@ -40,6 +40,7 @@ from .annotation_schema import ( METADATA_TABLE, ) from .autotag_schema import AUTOTAG_TABLE +from .image_removal_schema import clear_marks from .library_commands import _UUID, LibraryCommandError, _validate_epoch from .models import utc_now_iso from .vocabulary import validate_vocabulary_tag_name @@ -236,6 +237,9 @@ class AnnotationOperations: @staticmethod def _annotation_owner(conn: sqlite3.Connection, locator: str, now: str) -> str: + # A file with image tags or Needs review state is in the library: a + # removal mark never hides it (#480). + clear_marks(conn, [locator]) row = conn.execute(_OWNER, (locator,)).fetchone() if row is not None: return row[0] diff --git a/database/image_categories.py b/database/image_categories.py index a1f5d5a1..944d1bab 100644 --- a/database/image_categories.py +++ b/database/image_categories.py @@ -221,6 +221,7 @@ class ImageCategoryOperations: sources. Returns ``{updated, unchanged}``; a file that already had exactly this value is ``unchanged`` and keeps its ``updated_at``. """ + from .image_removal_schema import clear_marks_of_keys from .models import utc_now_iso # models imports this module's schema keys = list(dict.fromkeys(tuple(key) for key in file_keys)) @@ -234,6 +235,7 @@ class ImageCategoryOperations: continue conn.execute(UPSERT_SQL, (*key, category, now)) counts["updated"] += 1 + clear_marks_of_keys(conn, keys) # #480: a categorized file is listed return counts def clear_image_categories(self, file_keys) -> int: diff --git a/database/image_collections.py b/database/image_collections.py index 9eeb4494..9669a08c 100644 --- a/database/image_collections.py +++ b/database/image_collections.py @@ -402,6 +402,7 @@ class ImageCollectionOperations: 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): @@ -438,6 +439,7 @@ class ImageCollectionOperations: ).rowcount for root, path in pairs ) + clear_marks_of_keys(conn, pairs) # #480: a collected file is listed return { "collection": collection, "created": created, diff --git a/database/image_favorites.py b/database/image_favorites.py index b6034986..fb5cde31 100644 --- a/database/image_favorites.py +++ b/database/image_favorites.py @@ -104,11 +104,14 @@ class ImageFavoriteOperations: """Mark one file as a favorite; repeating it is a no-op. Returns True.""" from .models import utc_now_iso # models imports this module's schema + from .image_removal_schema import clear_marks_of_keys + with self.model.get_connection() as conn: conn.execute( "INSERT OR IGNORE INTO image_favorites VALUES (?, ?, ?)", (root_key, relative_path, utc_now_iso()), ) + clear_marks_of_keys(conn, [(root_key, relative_path)]) # #480 return True def clear_image_favorite(self, root_key: str, relative_path: str) -> bool: diff --git a/database/image_removal.py b/database/image_removal.py new file mode 100644 index 00000000..3360433c --- /dev/null +++ b/database/image_removal.py @@ -0,0 +1,305 @@ +"""Remove a file's records from the library, and its removal mark (#480). + +A *removal* describes one physical file the way the API resolved it from the +walked listing (``py/api/image_delete.py``): + +- ``spellings``: every path string records may name it by + (``duplicate_spellings.Spellings.of``); +- ``locator``: its private canonical locator; +- ``favorites``: ``(root key, relative key)`` under every listed root holding + it, the keys of its favorite and its collection memberships; +- ``category``: its own category's file key, or None. + +``remove_image_records`` deletes, in one ``BEGIN IMMEDIATE`` transaction: + +- its image records (``generated_images``) under every spelling, plus legacy + rows matched by ``image_path``; their saved-pair links go by cascade and the + workflow documents they held are released; +- the image tags and auto-tag marks of its byte records. The annotation file + and byte rows go too, unless a Needs review item still cites them + (``media_review_items.byte_id`` is ``ON DELETE RESTRICT``): a decided or + revised item keeps its decision; +- its pending, undecided Needs review items (settled as duplicate removal + settles them); +- its own category, favorite and collection memberships. + +With ``mark`` it also writes the file's removal mark (library-only removal): +the file stays on disk, hidden until its bytes change. Prompts, recorded runs +and the media file itself are never touched here. +""" + +import os + +from .annotation_schema import BYTES_TABLE, FILES_TABLE +from .annotations import _BUMP_INDEX +from .autotag_schema import AUTOTAG_TABLE +from .duplicate_copies import ( + DELETE_FAVORITE_SQL, + DELETE_LINK_SQL, + DELETE_MEMBERSHIPS_SQL, + _file_link_rows, +) +from .image_categories import DELETE_SQL as DELETE_CATEGORY_SQL +from .image_removal_schema import ( + META_TABLE, + TABLE, + clear_marks, + has_marks_table, + mark_of, + write_mark as _write_mark, +) +from .library import display_title +from .models import utc_now_iso +from .review import _BUMP_REVIEW_INDEX, _PENDING_OF_LOCATOR, _SETTLE_ITEM +from .review_changes import notify_review_change +from .workflow_documents import documents_of_rows, release_documents + +BYTES_OF_LOCATOR_SQL = ( + f"SELECT b.id FROM {FILES_TABLE} f JOIN {BYTES_TABLE} b " # nosec B608 + "ON b.file_id=f.id WHERE f.locator=?" +) +TAG_COUNT_SQL = "SELECT COUNT(*) FROM media_image_tags WHERE byte_id=?" +AUTOTAG_COUNT_SQL = ( + f"SELECT COUNT(*) FROM {AUTOTAG_TABLE} WHERE byte_id=?" # nosec B608 +) +CLEAR_TAGS_SQL = "DELETE FROM media_image_tags WHERE byte_id=?" +CLEAR_AUTOTAG_SQL = f"DELETE FROM {AUTOTAG_TABLE} WHERE byte_id=?" # nosec B608 +BUMP_TAG_REVISION_SQL = ( + f"UPDATE {BYTES_TABLE} SET tag_revision=tag_revision+1 WHERE id=?" # nosec B608 +) +REVIEW_REFS_SQL = ( + f"SELECT COUNT(*) FROM {FILES_TABLE} f " # nosec B608 + f"JOIN {BYTES_TABLE} b ON b.file_id=f.id " + "JOIN media_review_items i ON i.byte_id=b.id WHERE f.locator=?" +) +DELETE_FILE_SQL = f"DELETE FROM {FILES_TABLE} WHERE locator=?" # nosec B608 +FILE_EXISTS_SQL = f"SELECT 1 FROM {FILES_TABLE} WHERE locator=?" # nosec B608 +FAVORITE_EXISTS_SQL = ( + "SELECT COUNT(*) FROM image_favorites WHERE root_key=? AND relative_path=?" +) +MEMBERSHIP_COUNT_SQL = ( + "SELECT COUNT(*) FROM image_collection_items WHERE root_key=? AND relative_path=?" +) +CATEGORY_EXISTS_SQL = ( + "SELECT COUNT(*) FROM image_categories WHERE root_key=? AND relative_path=?" +) +PAIR_LINK_COUNT_SQL = "SELECT COUNT(*) FROM image_prompt_pairs WHERE image_id=?" +PROMPT_LINKS_SQL = "SELECT id FROM generated_images WHERE prompt_id=?" +PROMPT_TITLE_SQL = ( + "SELECT p.text, s.title FROM prompts p " + "LEFT JOIN prompt_library_state s ON s.prompt_id=p.id WHERE p.id=?" +) +MARKS_SQL = f"SELECT locator, size, mtime_ns FROM {TABLE}" # nosec B608 +REVISION_SQL = f"SELECT revision FROM {META_TABLE} WHERE singleton_id=1" # nosec B608 +MAX_ORPHAN_TITLES = 3 +RECORD_KINDS = ( + "prompt_links", + "pair_links", + "image_tags", + "auto_tagged", + "categories", + "collections", + "favorites", + "review_items", +) + + +def _link_ids(conn, removal): + """``{image id: prompt id}`` of the records naming the file, any spelling.""" + found = {} + for spelled in dict.fromkeys(removal["spellings"]): + for image_id, prompt_id in _file_link_rows(conn, spelled): + found[image_id] = prompt_id + return found + + +def _byte_ids(conn, locator): + return [row[0] for row in conn.execute(BYTES_OF_LOCATOR_SQL, (locator,))] + + +def file_locator(path) -> str: + """The locator of an absolute path (``normcase(realpath)``).""" + return os.path.normcase(os.path.realpath(str(path))) + + +def removed_from_library(conn, path) -> bool: + """Whether *path* is hidden by a removal mark whose bytes still match. + + A mark whose file changed (other size or ``mtime_ns``) or is gone is + stale: it is deleted on *conn* (the caller's write transaction), and the + file is not hidden. + """ + locator = file_locator(path) + mark = mark_of(conn, locator) + if mark is None: + return False + try: + info = os.lstat(path) + except OSError: + clear_marks(conn, [locator]) + return False + if (info.st_size, info.st_mtime_ns) == tuple(mark): + return True + clear_marks(conn, [locator]) + return False + + +def write_mark(conn, locator, size, mtime_ns): + _write_mark(conn, locator, size, mtime_ns, utc_now_iso()) + + +class ImageRemovalOperations: + """PromptDatabase mixin: remove a file's records; removal marks.""" + + def count_removal_records(self, removals) -> dict: + """Counts of what removing *removals* deletes, in one read snapshot. + + ``{records: {kind: n}, orphaned_prompts: {count, titles}}``: prompts + whose every image record is among these files lose their last image + (they stay in the Library). + """ + from .maintenance import _read_snapshot + + def read(conn): + totals = dict.fromkeys(RECORD_KINDS, 0) + links = {} + for removal in removals: + found = _link_ids(conn, removal) + links.update(found) + for image_id in found: + totals["pair_links"] += conn.execute( + PAIR_LINK_COUNT_SQL, (image_id,) + ).fetchone()[0] + tagged = autotagged = False + for byte_id in _byte_ids(conn, removal["locator"]): + tagged |= bool( + conn.execute(TAG_COUNT_SQL, (byte_id,)).fetchone()[0] + ) + autotagged |= bool( + conn.execute(AUTOTAG_COUNT_SQL, (byte_id,)).fetchone()[0] + ) + totals["image_tags"] += int(tagged) + totals["auto_tagged"] += int(autotagged) + totals["review_items"] += len( + conn.execute(_PENDING_OF_LOCATOR, (removal["locator"],)).fetchall() + ) + totals["favorites"] += int( + any( + conn.execute(FAVORITE_EXISTS_SQL, key).fetchone()[0] + for key in removal["favorites"] + ) + ) + totals["collections"] += sum( + conn.execute(MEMBERSHIP_COUNT_SQL, key).fetchone()[0] + for key in removal["favorites"] + ) + if removal.get("category"): + totals["categories"] += conn.execute( + CATEGORY_EXISTS_SQL, removal["category"] + ).fetchone()[0] + totals["prompt_links"] = len(links) + return { + "records": totals, + "orphaned_prompts": self._orphaned_prompts(conn, links), + } + + return _read_snapshot(self.model.db_path, read) + + @staticmethod + def _orphaned_prompts(conn, links): + removed = set(links) + orphans = [] + for prompt_id in dict.fromkeys(links.values()): + remaining = [ + row[0] + for row in conn.execute(PROMPT_LINKS_SQL, (prompt_id,)) + if row[0] not in removed + ] + if not remaining: + orphans.append(prompt_id) + titles = [] + for prompt_id in orphans[:MAX_ORPHAN_TITLES]: + row = conn.execute(PROMPT_TITLE_SQL, (prompt_id,)).fetchone() + if row is not None: + titles.append(display_title((row[1],), row[0]) or "") + return {"count": len(orphans), "titles": titles} + + def remove_image_records(self, removal, *, mark=None) -> dict: + """Delete every record of one file (one transaction); see the module. + + ``mark``: ``(size, mtime_ns)`` to write the file's removal mark, or + None (the file was moved to the trash). Returns the counts deleted. + """ + counts = dict.fromkeys(RECORD_KINDS, 0) + locator = removal["locator"] + with self._pair_transaction() as conn: + links = _link_ids(conn, removal) + if links: + documents = documents_of_rows(conn, "id", list(links)) + for image_id in links: + counts["pair_links"] += conn.execute( + PAIR_LINK_COUNT_SQL, (image_id,) + ).fetchone()[0] + conn.execute(DELETE_LINK_SQL, (image_id,)) + release_documents(conn, documents) + counts["prompt_links"] = len(links) + counts.update(self._remove_annotations(conn, locator)) + for key in removal["favorites"]: + counts["favorites"] += conn.execute(DELETE_FAVORITE_SQL, key).rowcount + counts["collections"] += conn.execute( + DELETE_MEMBERSHIPS_SQL, key + ).rowcount + if removal.get("category"): + counts["categories"] += conn.execute( + DELETE_CATEGORY_SQL, tuple(removal["category"]) + ).rowcount + if mark is not None: + write_mark(conn, locator, *mark) + if counts["review_items"]: + notify_review_change(self.model.db_path) + return counts + + def _remove_annotations(self, conn, locator) -> dict: + """Tags, auto-tag marks, pending review items and, when free, the rows.""" + counts = {"image_tags": 0, "auto_tagged": 0, "review_items": 0} + bytes_ = _byte_ids(conn, locator) + for byte_id in bytes_: + tags = conn.execute(CLEAR_TAGS_SQL, (byte_id,)).rowcount + autotags = conn.execute(CLEAR_AUTOTAG_SQL, (byte_id,)).rowcount + if tags: + conn.execute(BUMP_TAG_REVISION_SQL, (byte_id,)) + counts["image_tags"] |= bool(tags) + counts["auto_tagged"] |= bool(autotags) + for (item_id,) in conn.execute(_PENDING_OF_LOCATOR, (locator,)).fetchall(): + counts["review_items"] += conn.execute(_SETTLE_ITEM, (item_id,)).rowcount + if counts["review_items"]: + conn.execute(_BUMP_REVIEW_INDEX) + if conn.execute(FILE_EXISTS_SQL, (locator,)).fetchone(): + if not conn.execute(REVIEW_REFS_SQL, (locator,)).fetchone()[0]: + conn.execute(DELETE_FILE_SQL, (locator,)) + conn.execute(_BUMP_INDEX) + counts["image_tags"] = int(counts["image_tags"]) + counts["auto_tagged"] = int(counts["auto_tagged"]) + return counts + + def removal_marks(self) -> dict: + """``{locator: (size, mtime_ns)}`` of every removal mark (one read).""" + conn = self.model.get_connection() + if not has_marks_table(conn): + return {} + return {row[0]: (row[1], row[2]) for row in conn.execute(MARKS_SQL)} + + def removal_revision(self) -> int: + conn = self.model.get_connection() + if not has_marks_table(conn): + return 0 + row = conn.execute(REVISION_SQL).fetchone() + return row[0] if row else 0 + + def clear_removal_marks(self, locators) -> int: + """Show these files again (their marks go); the count removed.""" + locators = list(locators) + if not locators: + return 0 + with self._pair_transaction() as conn: + return clear_marks(conn, locators) diff --git a/database/image_removal_schema.py b/database/image_removal_schema.py new file mode 100644 index 00000000..3fad49a6 --- /dev/null +++ b/database/image_removal_schema.py @@ -0,0 +1,191 @@ +"""Removal marks (#480): files removed from the library that stay on disk. + +``image_removals`` holds one row per hidden file: its private canonical +locator (the ``media_annotation_files.locator`` form) and the ``size`` and +``mtime_ns`` it had when it was removed. A file at that locator with other +bytes is a new file and shows again. ``image_removals_meta`` is one row, the +marks revision, bumped in the same transaction as every mark change so the +workspace listing caches can tell marks changed. + +The group installs like ``image_favorites``: optional in old backups, +required once a database is current. Older builds ignore it. +""" + +import os +import re +import sqlite3 + +from .annotation_schema import MAX_LOCATOR_LENGTH + +TABLE = "image_removals" +META_TABLE = "image_removals_meta" +TABLES = (TABLE, META_TABLE) +_TABLE_DDL = { + TABLE: ( + f"CREATE TABLE {TABLE} (locator TEXT PRIMARY KEY, " + "size INTEGER NOT NULL CHECK(size>=0), mtime_ns INTEGER NOT NULL, " + "removed_at TEXT NOT NULL) WITHOUT ROWID" + ), + META_TABLE: ( + f"CREATE TABLE {META_TABLE} (singleton_id INTEGER PRIMARY KEY " + "CHECK(singleton_id=1), revision INTEGER NOT NULL DEFAULT 0 " + "CHECK(revision>=0))" + ), +} +_SEED_META = ( + f"INSERT OR IGNORE INTO {META_TABLE} (singleton_id, revision) VALUES (1, 0)" +) +_COLUMNS = { + TABLE: ( + ("locator", "TEXT", 1, 1), + ("size", "INTEGER", 1, 0), + ("mtime_ns", "INTEGER", 1, 0), + ("removed_at", "TEXT", 1, 0), + ), + META_TABLE: ( + ("singleton_id", "INTEGER", 0, 1), + ("revision", "INTEGER", 1, 0), + ), +} +_AUTOINDEX = f"sqlite_autoindex_{TABLE}_1" + + +def _compact(sql: str) -> str: + return re.sub(r"\s+", "", sql or "").lower() + + +def _structure_error(conn: sqlite3.Connection, allow_missing: bool) -> str | None: + rows = conn.execute( + "SELECT type,name,tbl_name,sql FROM sqlite_master WHERE name IN (?,?) " + "OR tbl_name IN (?,?)", + TABLES + TABLES, + ).fetchall() + if not rows: + return None if allow_missing else "Missing image removal marks schema" + tables = {name: sql for kind, name, _tbl, sql in rows if kind == "table"} + others = {(kind, name) for kind, name, _tbl, _sql in rows if kind != "table"} + if set(tables) != set(TABLES) or others - {("index", _AUTOINDEX)}: + return "Image removal marks name is reserved by an unsupported schema object" + for table in TABLES: + if _compact(tables[table]) != _compact(_TABLE_DDL[table]): + return "Image removal marks have an invalid table definition" + columns = tuple( + (row[1], row[2].upper(), row[3], row[5]) + for row in conn.execute(f"PRAGMA table_info({table})") + ) + if columns != _COLUMNS[table]: + return "Image removal marks have invalid columns" + return None + + +def _values_error(conn: sqlite3.Connection) -> str | None: + limit = MAX_LOCATOR_LENGTH * 4 + if conn.execute( + f"SELECT 1 FROM {TABLE} WHERE typeof(locator) != 'text' " # nosec B608 + f"OR length(CAST(locator AS BLOB)) > {limit} OR locator = '' " + "OR typeof(size) != 'integer' OR typeof(mtime_ns) != 'integer' " + "OR typeof(removed_at) != 'text' OR length(CAST(removed_at AS BLOB)) > 64 " + "LIMIT 1" + ).fetchone(): + return "Image removal marks contain invalid values" + meta = conn.execute( + f"SELECT singleton_id, typeof(revision) FROM {META_TABLE}" # nosec B608 + ).fetchall() + if meta != [(1, "integer")]: + return "Image removal marks have an invalid revision" + return None + + +def image_removals_schema_error( + conn: sqlite3.Connection, *, allow_missing: bool = True +) -> str | None: + """Reject unsupported schemas and invalid persisted values (reads only).""" + error = _structure_error(conn, allow_missing) + if error: + return error + exists = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (TABLE,) + ).fetchone() + return _values_error(conn) if exists else None + + +def migrate_image_removals(conn: sqlite3.Connection) -> None: + """Install both tables atomically; validate before and after.""" + conn.execute("SAVEPOINT image_removals_schema") + try: + error = image_removals_schema_error(conn) + 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(_SEED_META) + error = image_removals_schema_error(conn, allow_missing=False) + if error: + raise sqlite3.DatabaseError(error) + conn.execute("RELEASE image_removals_schema") + except BaseException: + conn.execute("ROLLBACK TO image_removals_schema") + conn.execute("RELEASE image_removals_schema") + raise + + +_UNMARK_SQL = f"DELETE FROM {TABLE} WHERE locator=?" # nosec B608 +_MARK_SQL = ( + f"INSERT OR REPLACE INTO {TABLE} " # nosec B608 + "(locator, size, mtime_ns, removed_at) VALUES (?, ?, ?, ?)" +) +_MARK_OF_SQL = f"SELECT size, mtime_ns FROM {TABLE} WHERE locator=?" # nosec B608 +BUMP_MARKS_SQL = ( + f"UPDATE {META_TABLE} SET revision=revision+1 WHERE singleton_id=1" # nosec B608 +) + + +def has_marks_table(conn: sqlite3.Connection) -> bool: + return ( + conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (TABLE,) + ).fetchone() + is not None + ) + + +def locator_of_key(root_key: str, relative_key: str) -> str: + """The locator of a favorite-style file key (root key + relative key).""" + return os.path.normcase(os.path.join(root_key, *relative_key.split("/"))) + + +def mark_of(conn: sqlite3.Connection, locator: str): + """``(size, mtime_ns)`` of the file's removal mark, or None.""" + if not has_marks_table(conn): + return None + row = conn.execute(_MARK_OF_SQL, (locator,)).fetchone() + return (row[0], row[1]) if row else None + + +def write_mark( + conn: sqlite3.Connection, locator: str, size: int, mtime_ns: int, now: str +): + conn.execute(_MARK_SQL, (locator, int(size), int(mtime_ns), now)) + conn.execute(BUMP_MARKS_SQL) + + +def clear_marks(conn: sqlite3.Connection, locators) -> int: + """Delete the marks of *locators* on *conn* (no commit); the count removed. + + Every record write for a file calls this (#480): a file with records is + in the library, so it is never hidden. + """ + locators = list(dict.fromkeys(locators)) + if not locators or not has_marks_table(conn): + return 0 + removed = 0 + for locator in locators: + removed += conn.execute(_UNMARK_SQL, (locator,)).rowcount + if removed: + conn.execute(BUMP_MARKS_SQL) + return removed + + +def clear_marks_of_keys(conn: sqlite3.Connection, file_keys) -> int: + """``clear_marks`` for favorite-style ``(root key, relative key)`` pairs.""" + return clear_marks(conn, [locator_of_key(*key) for key in file_keys]) diff --git a/database/models.py b/database/models.py index 0ef8d35e..d64cbcba 100644 --- a/database/models.py +++ b/database/models.py @@ -44,6 +44,7 @@ try: except ImportError: from utils.hashing import register_prompt_functions from .image_categories import image_categories_schema_error, migrate_image_categories +from .image_removal_schema import image_removals_schema_error, migrate_image_removals from .workflow_documents import ( DOCUMENT_ID_COLUMNS, migrate_workflow_documents, @@ -366,6 +367,8 @@ class PromptModel: migrate_autotag_state(conn) # Image categories (#332): independent of every table above. migrate_image_categories(conn) + # Removal marks (#480): independent of every table above. + migrate_image_removals(conn) # Last: the older generated_images rebuilds run only before file_path. migrate_workflow_documents(conn) @@ -1242,6 +1245,11 @@ class PromptModel: ) if categories_error: return False, categories_error + removals_error = image_removals_schema_error( + conn, allow_missing=not require_library_state + ) + if removals_error: + return False, removals_error views_error = saved_views_schema_error( conn, allow_missing=not require_library_state ) diff --git a/database/operations.py b/database/operations.py index a40651b4..da5d0ae8 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_removal import ImageRemovalOperations, removed_from_library from .image_collections import ImageCollectionOperations from .annotations import AnnotationOperations from .review import ReviewOperations @@ -465,6 +466,7 @@ class PromptDatabase( WorkflowStorageOperations, CategoryImportOperations, ImageCategoryOperations, + ImageRemovalOperations, ): """Database operations class for managing prompts.""" @@ -1806,6 +1808,11 @@ class PromptDatabase( """ if not isinstance(metadata, dict): metadata = {} + if removed_from_library(conn, image_path): + # #480: the user removed this exact file from the library; its + # bytes are unchanged, so nothing links it again. + self.logger.debug("Not linking a file removed from the library") + return 0 filename = os.path.basename(image_path) file_path = normalize_image_path(image_path) file_size, width, height, image_format = _coerce_file_info(metadata) diff --git a/tests/test_image_removal.py b/tests/test_image_removal.py new file mode 100644 index 00000000..03046dcd --- /dev/null +++ b/tests/test_image_removal.py @@ -0,0 +1,313 @@ +"""Removing one file's records from the library, and removal marks (#480).""" + +import hashlib +import os +import sqlite3 +import tempfile +import unittest +from pathlib import Path + +try: + from tests.open_handles import assert_closed +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + +from database.image_removal_schema import ( + image_removals_schema_error, + migrate_image_removals, +) +from database.models import PromptModel +from database.operations import PromptDatabase + + +def _locator(path): + return os.path.normcase(os.path.realpath(str(path))) + + +class _RemovalCase(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) + base = Path(self.temp.name) + self.root = base / "output" + self.root.mkdir() + self.alias = base / "alias" + self.alias.symlink_to(self.root, target_is_directory=True) + self.db = PromptDatabase(str(base / "prompts.db")) + self.addCleanup(self.db.close_all) + self.epoch = self.db.get_execution_epoch() + + def file(self, name, data=b"image-bytes"): + path = self.root / name + path.write_bytes(data) + return path + + def removal(self, path): + root_key = os.path.normcase(os.path.realpath(str(self.root))) + return { + "spellings": [str(path), str(self.alias / path.name)], + "locator": _locator(path), + "favorites": [(root_key, path.name.lower())], + "category": (root_key, path.name.lower()), + } + + def link(self, path, text="a prompt", spelled=None): + prompt_id = self.db.save_prompt(text) + image_id = self.db.link_image_to_prompt(str(prompt_id), str(spelled or path)) + self.assertTrue(image_id) + return prompt_id, image_id + + def tag(self, path, tags): + data = path.read_bytes() + info = os.stat(path) + return self.db.replace_image_tags( + locator=_locator(path), + sha256=hashlib.sha256(data).hexdigest(), + size=info.st_size, + mtime_ns=info.st_mtime_ns, + expected_epoch=self.epoch, + expected_tag_revision=0, + tags=tags, + ) + + def queue(self, path): + data = path.read_bytes() + info = os.stat(path) + return self.db.record_review_item( + locator=_locator(path), + spelled_path=os.path.realpath(str(path)), + sha256=hashlib.sha256(data).hexdigest(), + size=info.st_size, + mtime_ns=info.st_mtime_ns, + epoch=self.epoch, + resolver_version=1, + evidence={"candidates": []}, + candidate_count=0, + limited=True, + ) + + def count(self, sql, *params): + conn = sqlite3.connect(self.db.model.db_path) + try: + return conn.execute(sql, params).fetchone()[0] + finally: + conn.close() + + def decorate(self, path): + """Give *path* one record of every kind.""" + removal = self.removal(path) + root_key, relative = removal["favorites"][0] + self.tag(path, ["portrait"]) + self.db.set_image_favorite(root_key, relative) + self.db.set_image_categories([removal["category"]], "keepers") + self.db.add_to_collection(items=[(root_key, relative)], name="Set") + return removal + + +class ImageRemovalRecordsTests(_RemovalCase): + def test_counts_and_removes_every_record_under_every_spelling(self): + path = self.file("a.png") + other = self.file("b.png", b"other") + orphan, _ = self.link(path, "only this image") + shared, _ = self.link(path, "two images") + self.db.link_image_to_prompt(str(shared), str(other)) + # A record under the symlinked folder spelling. + self.link(path, "aliased", spelled=self.alias / "a.png") + removal = self.decorate(path) + self.assertTrue(self.queue(path)) + + counted = self.db.count_removal_records([removal]) + self.assertEqual( + counted["records"], + { + "prompt_links": 3, + "pair_links": 0, + "image_tags": 1, + "auto_tagged": 0, + "categories": 1, + "collections": 1, + "favorites": 1, + "review_items": 1, + }, + ) + self.assertEqual(counted["orphaned_prompts"]["count"], 2) + + removed = self.db.remove_image_records(removal) + self.assertEqual(removed["prompt_links"], 3) + self.assertEqual(removed["review_items"], 1) + self.assertEqual(self.count("SELECT COUNT(*) FROM generated_images"), 1) + for table in ( + "media_image_tags", + "media_annotation_files", + "media_review_items", + "image_favorites", + "image_categories", + "image_collection_items", + ): + with self.subTest(table=table): + self.assertEqual(self.count(f"SELECT COUNT(*) FROM {table}"), 0) + # Prompts and the collection itself stay. + self.assertEqual(self.count("SELECT COUNT(*) FROM prompts"), 3) + self.assertEqual(self.count("SELECT COUNT(*) FROM image_collections"), 1) + self.assertTrue(path.exists()) + self.assertIsNotNone(self.db.get_prompt_by_id(orphan)) + + def test_a_kept_review_item_keeps_the_annotation_rows(self): + path = self.file("a.png") + removal = self.decorate(path) + self.assertTrue(self.queue(path)) + conn = sqlite3.connect(self.db.model.db_path) + try: + # A revised item is not settled; it still cites the bytes, whose + # foreign key is ON DELETE RESTRICT. + conn.execute("UPDATE media_review_items SET revision=1") + conn.commit() + finally: + conn.close() + removed = self.db.remove_image_records(removal) + self.assertEqual(removed["review_items"], 0) + self.assertEqual(removed["image_tags"], 1) + self.assertEqual(self.count("SELECT COUNT(*) FROM media_image_tags"), 0) + self.assertEqual(self.count("SELECT COUNT(*) FROM media_review_items"), 1) + self.assertEqual(self.count("SELECT COUNT(*) FROM media_annotation_files"), 1) + + def test_mark_is_written_with_the_records_and_bumps_the_revision(self): + path = self.file("a.png") + self.link(path) + before = self.db.removal_revision() + info = os.stat(path) + self.db.remove_image_records( + self.removal(path), mark=(info.st_size, info.st_mtime_ns) + ) + self.assertEqual( + self.db.removal_marks(), + {_locator(path): (info.st_size, info.st_mtime_ns)}, + ) + self.assertEqual(self.db.removal_revision(), before + 1) + self.assertEqual(self.db.clear_removal_marks([_locator(path)]), 1) + self.assertEqual(self.db.removal_marks(), {}) + self.assertEqual(self.db.removal_revision(), before + 2) + self.assertEqual(self.db.clear_removal_marks([_locator(path)]), 0) + + def test_a_failure_rolls_back_everything(self): + path = self.file("a.png") + self.link(path) + removal = self.decorate(path) + broken = {**removal, "category": ("only-one-value",)} + with self.assertRaises(sqlite3.Error): + self.db.remove_image_records(broken, mark=(1, 1)) + self.assertEqual(self.count("SELECT COUNT(*) FROM generated_images"), 1) + self.assertEqual(self.count("SELECT COUNT(*) FROM media_image_tags"), 1) + self.assertEqual(self.db.removal_marks(), {}) + + def test_parity_with_the_legacy_single_delete(self): + path = self.file("a.png") + _prompt, image_id = self.link(path) + twin = PromptDatabase(str(Path(self.temp.name) / "twin.db")) + try: + twin_prompt = twin.save_prompt("a prompt") + twin_image = twin.link_image_to_prompt(str(twin_prompt), str(path)) + self.assertTrue(twin.delete_image(twin_image)) + self.db.remove_image_records(self.removal(path)) + for table in ( + "generated_images", + "image_prompt_pairs", + "workflow_documents", + ): + with self.subTest(table=table): + conn = sqlite3.connect(twin.model.db_path) + try: + legacy = conn.execute( + f"SELECT COUNT(*) FROM {table}" + ).fetchone()[0] + finally: + conn.close() + self.assertEqual( + self.count(f"SELECT COUNT(*) FROM {table}"), legacy + ) + finally: + twin.close_all() + self.assertIsNotNone(image_id) + + +class RemovalMarkGateTests(_RemovalCase): + """A marked file is never linked again; any record write shows it again.""" + + def hide(self, path): + info = os.stat(path) + self.db.remove_image_records( + self.removal(path), mark=(info.st_size, info.st_mtime_ns) + ) + + def test_a_marked_file_is_not_linked_again(self): + path = self.file("a.png") + self.hide(path) + prompt_id = self.db.save_prompt("scanned again") + self.assertEqual(self.db.link_image_to_prompt(str(prompt_id), str(path)), 0) + self.assertEqual(self.count("SELECT COUNT(*) FROM generated_images"), 0) + self.assertIn(_locator(path), self.db.removal_marks()) + + def test_a_changed_file_drops_its_stale_mark_and_links(self): + path = self.file("a.png") + self.hide(path) + path.write_bytes(b"new bytes, new size") + prompt_id = self.db.save_prompt("a new image at the same path") + self.assertTrue(self.db.link_image_to_prompt(str(prompt_id), str(path))) + self.assertEqual(self.db.removal_marks(), {}) + + def test_explicit_record_writes_clear_the_mark(self): + writers = { + "favorite": lambda path, key: self.db.set_image_favorite(*key), + "category": lambda path, key: self.db.set_image_categories([key], "c"), + "collection": lambda path, key: self.db.add_to_collection( + items=[key], name="Set" + ), + "tags": lambda path, key: self.tag(path, ["t"]), + } + for kind, write in writers.items(): + with self.subTest(kind=kind): + path = self.file(f"{kind}.png", kind.encode()) + self.hide(path) + self.assertIn(_locator(path), self.db.removal_marks()) + write(path, self.removal(path)["favorites"][0]) + self.assertNotIn(_locator(path), self.db.removal_marks()) + + +class ImageRemovalSchemaTests(unittest.TestCase): + def setUp(self): + self.conn = sqlite3.connect(":memory:") + self.addCleanup(self.conn.close) + + def test_missing_is_allowed_only_for_old_backups(self): + self.assertIsNone(image_removals_schema_error(self.conn)) + self.assertIsNotNone( + image_removals_schema_error(self.conn, allow_missing=False) + ) + + def test_install_is_valid_and_idempotent(self): + migrate_image_removals(self.conn) + migrate_image_removals(self.conn) + self.assertIsNone(image_removals_schema_error(self.conn, allow_missing=False)) + self.assertEqual( + self.conn.execute("SELECT revision FROM image_removals_meta").fetchall(), + [(0,)], + ) + + def test_invalid_values_and_objects_are_refused(self): + migrate_image_removals(self.conn) + self.conn.execute("INSERT INTO image_removals VALUES ('', 1, 1, 'now')") + self.assertIn("invalid values", image_removals_schema_error(self.conn)) + self.conn.execute("DELETE FROM image_removals") + self.conn.execute("CREATE INDEX extra ON image_removals(size)") + self.assertIn("unsupported", image_removals_schema_error(self.conn)) + + def test_a_foreign_table_definition_is_refused(self): + self.conn.execute("CREATE TABLE image_removals (locator TEXT)") + with self.assertRaises(sqlite3.DatabaseError): + migrate_image_removals(self.conn) + + +if __name__ == "__main__": + unittest.main()