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.
This commit is contained in:
Vito Sansevero
@@ -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]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
@@ -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])
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user