diff --git a/TESTING.md b/TESTING.md index 936139c..14a10c0 100644 --- a/TESTING.md +++ b/TESTING.md @@ -1367,6 +1367,61 @@ concept fixture stars the concept's favorites (each other prompt's first image). restore. Old backups without the table restore, and backup plus restore keeps WAL-resident favorites and closes every handle (`tests/test_database_backup.py`). +Image annotations schema (#249, `database/annotation_schema.py`): the group +`media_annotation_metadata` (one row: schema version, `image_index_revision`), +`media_annotation_files` (one private locator per physical file, +`normcase(join(realpath(root), relative))`, server-side only), +`media_annotation_bytes` (one record per observed whole-file SHA-256 with the +`size`/`mtime_ns` verified when those bytes were last seen and the set's +`tag_revision`) and `media_image_tags` (the ordered tag set of one byte record, +joined into the global `tags` registry). It installs like favorites: optional in +old backups, required once current (`annotation_schema_error`, the single +`require_library_state` flag), validated read-only before and after a SAVEPOINT +install and on restore, never backfilled and never stat-ing media +(`tests/test_annotation_schema.py`). Writers (`database/annotations.py`) always +supply the UUID ids, run in `_pair_transaction` fenced by the storage epoch, and +replace or union one byte record's set (`tag-conflict` / `tag-limit` with the +current set, identical retries `replayed`); vocabulary rename, merge and delete +carry `media_image_tags` through the one `_move_tag_links` loop, and tag rows, +impact and command receipts report `image_count` / `images` / +`images_changed` as files (`tests/test_annotations.py`, +`tests/test_pair_tag_maintenance.py`, `tests/test_vocabulary.py`). Cheap reads +(spec amendment S1) use a byte record only while the file's current +`(size, mtime_ns)` equals its verified fingerprint; an in-place rewrite that +preserves both keeps the old set matching until the file is next observed. The +browser fixtures clear `media_image_tags`, `media_annotation_bytes` and +`media_annotation_files` on reset and never the metadata row. + +Capture store (#249, `utils/media_capture.py`, `tests/test_media_capture.py`): +one explicit action captures one rooted media file into an owner-only +`prompt-manager-captures-*` temp directory, streaming it in 1 MiB chunks while +hashing. Identity is compared like with like: the opened descriptor against +the path only on the API-stable facts (`size`, `mtime_ns`); the full identity +(`size`, `mtime_ns`, `ctime_ns`, `dev`/`ino` where meaningful) between the +`fstat` before and after the copy and between the path stats before and after +(native Windows CI with Python 3.12/3.13 showed `stat` and `fstat` disagreeing +on `st_ctime_ns`: change time versus creation time; 3.10/3.11 agree). Any +disagreement is `CaptureChanged` and the copy is removed. Error messages carry +`strerror` only, never a path. Aliases (a symlinked +file or directory, even inside the root), escapes, non-media names and +directories are `CaptureDenied`; the locator is +`normcase(join(realpath(root), relative))`, so a symlinked root shares the +listing's identity. Fixed limits: 256 MiB per file (stat and stream), 8 +completed captures, 512 MiB reserved, 2 capture/verify workers, 900 s idle, +1800 s hard; admission is nonblocking (`CaptureCapacity`) and lifetimes are +tested with an injected clock. An acquired capture outlives expiry and release +until its last use ends. Header-only dimensions (4096 px per side, +16,777,216 px) and video mark Generate unsupported; `verify` re-hashes the +current file without a copy. Every handle closes before unlink (a patched +`os.unlink` asserts it, plus `assert_closed`), and files are unlinked outside +the store lock. Known limit: on Windows before Python 3.13 `fstat` reports the +creation time, so a same-size rewrite that restores `mtime_ns` during the copy +window completes there with a digest of the copied bytes; POSIX and Windows on +3.13+ detect it through the descriptor's change time (native CI, 3.13 only). +Vocabulary rename, merge and delete bump the `tag_revision` of every image set +they change (and `image_index_revision`), so a Tag dialog opened before the +change gets `tag-conflict` on Save (`tests/test_pair_tag_maintenance.py`). + Budgets: `test_filter_work_is_bounded` always runs. It asserts that a plain first page makes 1 link SELECT and a plain cursor page none, that every filter shape makes exactly 1, the covering-index plan, and that the per-file values are built diff --git a/database/annotation_schema.py b/database/annotation_schema.py new file mode 100644 index 0000000..feaa1fa --- /dev/null +++ b/database/annotation_schema.py @@ -0,0 +1,290 @@ +"""Image annotation schema group (#249): file locators, byte records, image tags. + +The group installs like ``image_favorites`` and ``saved_views``: optional in old +backups, required once a database is current. It owns + +- ``media_annotation_metadata``: one row, the group's schema version and the + image index revision; +- ``media_annotation_files``: one private canonical locator per physical file + (``normcase(join(realpath(root), relative))``, server-side only); +- ``media_annotation_bytes``: one record per observed whole-file SHA-256 of a + file, with the size and ``mtime_ns`` verified when those bytes were last seen + and the tag set's revision; +- ``media_image_tags``: the ordered image tag set of one byte record, joined + into the global ``tags`` registry so vocabulary maintenance carries image + links like prompt and pair links. + +Installation never backfills anything from prompt or pair tags and never +touches the filesystem. Validation reads only: the catalog objects are compared +after whitespace compaction, columns, foreign keys and indexes through the +PRAGMAs, and persisted values in single bounded SQL passes. +""" + +import re +import sqlite3 +from typing import Optional + +from .image_favorites import MAX_KEY_LENGTH + +SCHEMA_VERSION = 1 +METADATA_TABLE = "media_annotation_metadata" +FILES_TABLE = "media_annotation_files" +BYTES_TABLE = "media_annotation_bytes" +IMAGE_TAGS_TABLE = "media_image_tags" +TABLES = (METADATA_TABLE, FILES_TABLE, BYTES_TABLE, IMAGE_TAGS_TABLE) +MAX_LOCATOR_LENGTH = MAX_KEY_LENGTH +_TIMESTAMP_BYTES = 64 + +_TABLE_DDL = { + METADATA_TABLE: ( + f"CREATE TABLE {METADATA_TABLE} (singleton_id INTEGER PRIMARY KEY " + "CHECK(singleton_id=1), schema_version INTEGER NOT NULL " + f"CHECK(schema_version={SCHEMA_VERSION}), " + "image_index_revision INTEGER NOT NULL DEFAULT 0 " + "CHECK(image_index_revision>=0))" + ), + FILES_TABLE: ( + f"CREATE TABLE {FILES_TABLE} (id TEXT PRIMARY KEY CHECK(length(id)=36), " + "locator TEXT NOT NULL UNIQUE, created_at TEXT NOT NULL)" + ), + BYTES_TABLE: ( + f"CREATE TABLE {BYTES_TABLE} (id TEXT PRIMARY KEY CHECK(length(id)=36), " + f"file_id TEXT NOT NULL REFERENCES {FILES_TABLE}(id) ON DELETE CASCADE, " + "sha256 TEXT NOT NULL CHECK(length(sha256)=64 AND sha256 NOT GLOB " + "'*[^0-9a-f]*'), size INTEGER NOT NULL CHECK(size>=0), " + "mtime_ns INTEGER NOT NULL, verified_at TEXT NOT NULL, " + "first_seen_at TEXT NOT NULL, tag_revision INTEGER NOT NULL DEFAULT 0 " + "CHECK(tag_revision>=0), UNIQUE(file_id, sha256))" + ), + IMAGE_TAGS_TABLE: ( + f"CREATE TABLE {IMAGE_TAGS_TABLE} (byte_id TEXT NOT NULL REFERENCES " + f"{BYTES_TABLE}(id) ON DELETE CASCADE, tag_id INTEGER NOT NULL " + "REFERENCES tags(id) ON DELETE CASCADE, position INTEGER NOT NULL " + "CHECK(position>=0), PRIMARY KEY(byte_id, tag_id))" + ), +} +_INDEX_DDL = { + "idx_media_annotation_bytes_fingerprint": ( + "CREATE INDEX idx_media_annotation_bytes_fingerprint ON " + f"{BYTES_TABLE}(file_id, size, mtime_ns)" + ), + "idx_media_image_tags_tag": ( + f"CREATE INDEX idx_media_image_tags_tag ON {IMAGE_TAGS_TABLE}(tag_id)" + ), + "idx_media_image_tags_position": ( + "CREATE UNIQUE INDEX idx_media_image_tags_position ON " + f"{IMAGE_TAGS_TABLE}(byte_id, position)" + ), +} +# SQLite names the indexes behind TEXT primary keys and UNIQUE constraints +# ``sqlite_autoindex__``; their catalog rows carry no SQL. +_AUTOINDEXES = {FILES_TABLE: 2, BYTES_TABLE: 2, IMAGE_TAGS_TABLE: 1} +# (name, declared type, not null, default, primary key position) +_COLUMNS = { + METADATA_TABLE: ( + ("singleton_id", "INTEGER", 0, None, 1), + ("schema_version", "INTEGER", 1, None, 0), + ("image_index_revision", "INTEGER", 1, "0", 0), + ), + FILES_TABLE: ( + ("id", "TEXT", 0, None, 1), + ("locator", "TEXT", 1, None, 0), + ("created_at", "TEXT", 1, None, 0), + ), + BYTES_TABLE: ( + ("id", "TEXT", 0, None, 1), + ("file_id", "TEXT", 1, None, 0), + ("sha256", "TEXT", 1, None, 0), + ("size", "INTEGER", 1, None, 0), + ("mtime_ns", "INTEGER", 1, None, 0), + ("verified_at", "TEXT", 1, None, 0), + ("first_seen_at", "TEXT", 1, None, 0), + ("tag_revision", "INTEGER", 1, "0", 0), + ), + IMAGE_TAGS_TABLE: ( + ("byte_id", "TEXT", 1, None, 1), + ("tag_id", "INTEGER", 1, None, 2), + ("position", "INTEGER", 1, None, 0), + ), +} +# (column, parent table, parent column, on delete) +_FOREIGN_KEYS = { + METADATA_TABLE: set(), + FILES_TABLE: set(), + BYTES_TABLE: {("file_id", FILES_TABLE, "id", "CASCADE")}, + IMAGE_TAGS_TABLE: { + ("byte_id", BYTES_TABLE, "id", "CASCADE"), + ("tag_id", "tags", "id", "CASCADE"), + }, +} + + +def _compact(sql: str) -> str: + return re.sub(r"\s+", "", sql).lower() + + +def _uuid(column: str) -> str: + """SQL: *column* is a lowercase hyphenated 36-character UUID string.""" + return ( + f"typeof({column}) = 'text' AND length({column}) = 36 " + f"AND substr({column}, 9, 1) = '-' AND substr({column}, 14, 1) = '-' " + f"AND substr({column}, 19, 1) = '-' AND substr({column}, 24, 1) = '-' " + f"AND length(replace({column}, '-', '')) = 32 " + f"AND replace({column}, '-', '') NOT GLOB '*[^0-9a-f]*'" + ) + + +def _timestamp(column: str) -> str: + return ( + f"typeof({column}) = 'text' " + f"AND length(CAST({column} AS BLOB)) <= {_TIMESTAMP_BYTES}" + ) + + +def _non_negative(column: str) -> str: + return f"typeof({column}) = 'integer' AND {column} >= 0" + + +# Every row of a table must satisfy its predicate; one pass each, LIMIT 1. +_ROW_PREDICATES = { + FILES_TABLE: " AND ".join( + ( + _uuid("id"), + "typeof(locator) = 'text'", + # UTF-8 is at most 4 bytes per character; the decoded length of + # the few candidates above the byte bound is checked in Python. + f"length(CAST(locator AS BLOB)) <= {MAX_LOCATOR_LENGTH * 4}", + _timestamp("created_at"), + ) + ), + BYTES_TABLE: " AND ".join( + ( + _uuid("id"), + _uuid("file_id"), + "typeof(sha256) = 'text' AND length(sha256) = 64 " + "AND sha256 NOT GLOB '*[^0-9a-f]*'", + _non_negative("size"), + "typeof(mtime_ns) = 'integer'", + _timestamp("verified_at"), + _timestamp("first_seen_at"), + _non_negative("tag_revision"), + ) + ), + IMAGE_TAGS_TABLE: " AND ".join( + (_uuid("byte_id"), "typeof(tag_id) = 'integer'", _non_negative("position")) + ), +} + + +def _expected_catalog() -> set: + objects = {("table", name, _compact(sql)) for name, sql in _TABLE_DDL.items()} + objects |= {("index", name, _compact(sql)) for name, sql in _INDEX_DDL.items()} + objects |= { + ("index", f"sqlite_autoindex_{table}_{n}", "") + for table, count in _AUTOINDEXES.items() + for n in range(1, count + 1) + } + return objects + + +def _catalog(conn: sqlite3.Connection) -> list: + names = TABLES + tuple(_INDEX_DDL) + marks = ",".join("?" * len(names)) + table_marks = ",".join("?" * len(TABLES)) + return conn.execute( # only "?" placeholders are interpolated + "SELECT type,name,sql FROM sqlite_master " + f"WHERE name IN ({marks}) OR tbl_name IN ({table_marks})", # nosec B608 + names + 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 annotations 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 annotations table {table} has invalid foreign keys" + if conn.execute(f"PRAGMA foreign_key_check({table})").fetchone(): + return f"Image annotations table {table} contains foreign key violations" + return None + + +def _value_error(conn: sqlite3.Connection) -> Optional[str]: + # Only this module's constants are interpolated into the SQL below. + rows = conn.execute( + "SELECT COUNT(*), MIN(typeof(singleton_id) = 'integer' " # nosec B608 + "AND singleton_id = 1 AND typeof(schema_version) = 'integer' " + f"AND schema_version = {SCHEMA_VERSION} " + f"AND {_non_negative('image_index_revision')}) FROM {METADATA_TABLE}" + ).fetchone() + if rows[0] != 1 or rows[1] != 1: + return "Image annotations metadata is invalid" + for table, predicate in _ROW_PREDICATES.items(): + if conn.execute( + f"SELECT 1 FROM {table} WHERE NOT ({predicate}) LIMIT 1" # nosec B608 + ).fetchone(): + return f"Image annotations table {table} contains invalid values" + # length() stops at an embedded NUL, so the SQL pass bounds bytes and the + # few candidates over the character bound are measured decoded here. + for (raw,) in conn.execute( + f"SELECT CAST(locator AS BLOB) FROM {FILES_TABLE} " # nosec B608 + f"WHERE length(CAST(locator AS BLOB)) > {MAX_LOCATOR_LENGTH}" + ): + if len(bytes(raw).decode("utf-8", "replace")) > MAX_LOCATOR_LENGTH: + return f"Image annotations table {FILES_TABLE} contains invalid values" + return None + + +def annotation_schema_error( + conn: sqlite3.Connection, *, allow_missing: bool = True +) -> Optional[str]: + """Reject partial or unsupported annotation schemas and invalid rows. + + Reads only: nothing is ever repaired. A database without any of the + group's objects is acceptable for backups that predate it unless + ``allow_missing`` is false. + """ + objects = _catalog(conn) + if not objects: + return None if allow_missing else "Missing image annotation schema" + present = {name for kind, name, _sql in objects if kind == "table"} + if present != set(TABLES): + return "Incomplete image annotation schema" + actual = {(kind, name, _compact(sql or "")) for kind, name, sql in objects} + if actual != _expected_catalog(): + return "Image annotations has an invalid schema definition" + return _structure_error(conn) or _value_error(conn) + + +def migrate_annotation_state(conn: sqlite3.Connection) -> None: + """Install the group atomically; validate before and after creating it.""" + conn.execute("SAVEPOINT media_annotation_schema") + try: + error = annotation_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)) + for sql in _INDEX_DDL.values(): + conn.execute(sql.replace("INDEX", "INDEX IF NOT EXISTS", 1)) + conn.execute( + f"INSERT OR IGNORE INTO {METADATA_TABLE} " + "(singleton_id, schema_version, image_index_revision) VALUES (1, ?, 0)", + (SCHEMA_VERSION,), + ) + error = annotation_schema_error(conn, allow_missing=False) + if error: + raise sqlite3.DatabaseError(error) + conn.execute("RELEASE media_annotation_schema") + except BaseException: + conn.execute("ROLLBACK TO media_annotation_schema") + conn.execute("RELEASE media_annotation_schema") + raise diff --git a/database/annotations.py b/database/annotations.py new file mode 100644 index 0000000..0a39be2 --- /dev/null +++ b/database/annotations.py @@ -0,0 +1,457 @@ +"""Per-image tag sets (#249): byte-exact storage in the shared vocabulary. + +A tag set belongs to one byte record (``media_annotation_bytes``): the +whole-file SHA-256 of one physical file, owned by the file's private locator +(``media_annotation_files``). Callers pass verified facts (digest, size, +``mtime_ns``); nothing here reads the filesystem or resolves an owner through +prompt links, so tagging one image never changes another image, a prompt or a +pair. + +Commands run in one ``_pair_transaction`` fenced by the storage epoch, like the +native Library and vocabulary commands. A save is a full replacement guarded +by the set's ``tag_revision``; the identical set sent again (a lost-response +retry, or the current set at its own revision) answers ``replayed: true`` and +writes nothing. ``add_image_tags`` unions into the existing set, keeping the +stored spelling. Names enter the global ``tags`` registry exactly as spelled, +so vocabulary rename, merge and delete carry image links with prompt and pair +links (``PromptDatabase._move_tag_links``). + +Cheap reads (spec amendment S1) treat a byte record as *current* only while +the file's present ``(size, mtime_ns)`` equals the record's verified +fingerprint; a file without a current record shows no image tags. +""" + +import re +import sqlite3 +from contextlib import contextmanager +from typing import Iterator +from uuid import uuid4 + +from .annotation_schema import ( + BYTES_TABLE, + FILES_TABLE, + IMAGE_TAGS_TABLE, + MAX_LOCATOR_LENGTH, + METADATA_TABLE, +) +from .library_commands import _UUID, LibraryCommandError, _validate_epoch +from .models import utc_now_iso +from .vocabulary import validate_vocabulary_tag_name + +MAX_IMAGE_TAGS = 100 +_LOOKUP_CHUNK = 500 +_DIGEST = re.compile(r"[0-9a-f]{64}") +_MAX_INT = 2**63 - 1 +_MIN_INT = -(2**63) +# Only this module's table constants are interpolated into the SQL below. +_OWNER = f"SELECT id FROM {FILES_TABLE} WHERE locator=?" # nosec B608 +_INSERT_OWNER = ( + f"INSERT INTO {FILES_TABLE}(id,locator,created_at) VALUES(?,?,?)" # nosec B608 +) +_BYTE = ( + f"SELECT id,tag_revision,size,mtime_ns FROM {BYTES_TABLE} " # nosec B608 + "WHERE file_id=? AND sha256=?" +) +_INSERT_BYTE = ( + f"INSERT INTO {BYTES_TABLE}(id,file_id,sha256,size,mtime_ns," # nosec B608 + "verified_at,first_seen_at,tag_revision) VALUES(?,?,?,?,?,?,?,0)" +) +_REFRESH_BYTE = ( + f"UPDATE {BYTES_TABLE} SET size=?,mtime_ns=?,verified_at=? WHERE id=?" # nosec B608 +) +_REFRESH_BY_LOCATOR = ( + f"UPDATE {BYTES_TABLE} SET size=?,mtime_ns=?,verified_at=? " # nosec B608 + f"WHERE sha256=? AND file_id=({_OWNER})" +) +_BYTE_TAGS = ( + f"SELECT t.name FROM {IMAGE_TAGS_TABLE} mit " # nosec B608 + "JOIN tags t ON t.id=mit.tag_id WHERE mit.byte_id=? ORDER BY mit.position" +) +_CLEAR_TAGS = f"DELETE FROM {IMAGE_TAGS_TABLE} WHERE byte_id=?" # nosec B608 +_INSERT_TAG = ( + f"INSERT INTO {IMAGE_TAGS_TABLE}(byte_id,tag_id,position) " # nosec B608 + "VALUES(?,?,?)" +) +_BUMP_REVISION = ( + f"UPDATE {BYTES_TABLE} SET tag_revision=tag_revision+1 WHERE id=?" # nosec B608 +) +_REVISION = f"SELECT tag_revision FROM {BYTES_TABLE} WHERE id=?" # nosec B608 +_BUMP_INDEX = ( + f"UPDATE {METADATA_TABLE} SET " # nosec B608 + "image_index_revision=image_index_revision+1 WHERE singleton_id=1" +) +_READ = ( + "SELECT b.id,b.tag_revision," # nosec B608 + f"(SELECT COUNT(*) FROM {BYTES_TABLE} o WHERE o.file_id=f.id AND o.sha256!=?) " + f"FROM {FILES_TABLE} f LEFT JOIN {BYTES_TABLE} b " + "ON b.file_id=f.id AND b.sha256=? WHERE f.locator=?" +) +_FINGERPRINTS = ( + "SELECT f.locator,b.id,b.size,b.mtime_ns,b.verified_at " # nosec B608 + f"FROM {FILES_TABLE} f JOIN {BYTES_TABLE} b ON b.file_id=f.id " + "WHERE f.locator IN ({marks})" +) +_TAGS_OF = ( + f"SELECT mit.byte_id,t.name FROM {IMAGE_TAGS_TABLE} mit " # nosec B608 + "JOIN tags t ON t.id=mit.tag_id WHERE mit.byte_id IN ({marks}) " + "ORDER BY mit.byte_id,mit.position" +) +_RECORDS = ( + "SELECT f.locator,b.id,b.size,b.mtime_ns,t.name " # nosec B608 + f"FROM {FILES_TABLE} f JOIN {BYTES_TABLE} b ON b.file_id=f.id " + f"LEFT JOIN {IMAGE_TAGS_TABLE} mit ON mit.byte_id=b.id " + "LEFT JOIN tags t ON t.id=mit.tag_id " + "ORDER BY f.locator,b.first_seen_at,b.id,mit.position" +) + + +class AnnotationCommandError(LibraryCommandError): + """An image tag command lost a race or hit the set's bound. + + ``code`` is ``tag-conflict`` (the set changed since it was read) or + ``tag-limit`` (the union would exceed ``MAX_IMAGE_TAGS``); ``current`` is + the stored set ``{sha256, tags, tag_revision}`` read inside the + transaction, so the caller can show it. + """ + + def __init__(self, code: str, epoch: str, current: dict | None = None): + super().__init__(code, epoch) + self.current = current + + +def normalize_image_tags(names: object) -> list[str]: + """Trim and validate each name, fold case-insensitive repeats, cap at 100. + + The first occurrence keeps its spelling and position. Raises ``ValueError`` + for anything but a list of valid vocabulary names or for more than + ``MAX_IMAGE_TAGS`` distinct names. + """ + if not isinstance(names, list): + raise ValueError("tags must be a list of strings") + result: list[str] = [] + seen: set[str] = set() + for value in names: + name = validate_vocabulary_tag_name(value) + key = name.casefold() + if key in seen: + continue + seen.add(key) + result.append(name) + if len(result) > MAX_IMAGE_TAGS: + raise ValueError(f"An image can carry at most {MAX_IMAGE_TAGS} tags") + return result + + +def _validate_locator(locator: object) -> None: + if not isinstance(locator, str) or not 0 < len(locator) <= MAX_LOCATOR_LENGTH: + raise ValueError( + f"locator must be a nonblank path of at most {MAX_LOCATOR_LENGTH} chars" + ) + + +def _validate_digest(sha256: object) -> None: + if not isinstance(sha256, str) or not _DIGEST.fullmatch(sha256): + raise ValueError("sha256 must be 64 lowercase hexadecimal characters") + + +def _validate_int(value: object, name: str, minimum: int) -> None: + if type(value) is not int or not minimum <= value <= _MAX_INT: + raise ValueError(f"{name} must be an integer not below {minimum}") + + +def _validate_target(locator: object, sha256: object, size: object, mtime_ns: object): + _validate_locator(locator) + _validate_digest(sha256) + _validate_int(size, "size", 0) + _validate_int(mtime_ns, "mtime_ns", _MIN_INT) + + +def _validate_fingerprint(locator: object, fingerprint: object) -> None: + _validate_locator(locator) + if ( + not isinstance(fingerprint, tuple) + or len(fingerprint) != 2 + or any(type(value) is not int for value in fingerprint) + ): + raise ValueError("stats must map locators to (size, mtime_ns) integer pairs") + + +def _new_id() -> str: + """A canonical UUID; the schema's CHECK lets NULL through, this does not.""" + value = str(uuid4()) + if not _UUID.fullmatch(value): + raise ValueError("Annotation ids must be canonical UUIDs") + return value + + +def _receipt(epoch: str, sha256: str, tags: list, revision: int, replayed: bool): + return { + "library_epoch": epoch, + "sha256": sha256, + "tags": list(tags), + "tag_revision": revision, + "replayed": replayed, + } + + +class AnnotationOperations: + """PromptDatabase mixin: image tag sets, fingerprints and cheap reads. + + ``_`` helpers take the open ``_pair_transaction`` connection. + """ + + @contextmanager + def _annotation_read(self) -> Iterator[sqlite3.Connection]: + """One deferred read transaction, so paired SELECTs see one snapshot.""" + conn = self.model.get_connection() + if conn.in_transaction: + raise sqlite3.ProgrammingError( + "Annotation reads require their own snapshot" + ) + with conn: + conn.execute("BEGIN") + yield conn + + @staticmethod + def _annotation_owner(conn: sqlite3.Connection, locator: str, now: str) -> str: + row = conn.execute(_OWNER, (locator,)).fetchone() + if row is not None: + return row[0] + file_id = _new_id() + conn.execute(_INSERT_OWNER, (file_id, locator, now)) + return file_id + + @staticmethod + def _annotation_byte( + conn: sqlite3.Connection, + file_id: str, + sha256: str, + size: int, + mtime_ns: int, + now: str, + ) -> tuple[str, int, bool, bool]: + """The byte record for these bytes, refreshed or created. + + Returns ``(id, tag_revision, created, refingerprinted)``; the last is + True when an existing record's verified ``(size, mtime_ns)`` changed, + which moves its currency for cheap reads (the caller bumps the index + revision when nothing else does). + """ + row = conn.execute(_BYTE, (file_id, sha256)).fetchone() + if row is None: + byte_id = _new_id() + conn.execute( + _INSERT_BYTE, (byte_id, file_id, sha256, size, mtime_ns, now, now) + ) + return byte_id, 0, True, False + conn.execute(_REFRESH_BYTE, (size, mtime_ns, now, row[0])) + return row[0], row[1], False, (row[2], row[3]) != (size, mtime_ns) + + @staticmethod + def _byte_tags(conn: sqlite3.Connection, byte_id: str) -> list[str]: + return [row[0] for row in conn.execute(_BYTE_TAGS, (byte_id,))] + + def _write_image_tags( + self, conn: sqlite3.Connection, byte_id: str, tags: list[str] + ) -> int: + """Rewrite the junction in order, bump both revisions; return the new one.""" + tag_ids = self._ensure_tags(conn, tags) + conn.execute(_CLEAR_TAGS, (byte_id,)) + conn.executemany( + _INSERT_TAG, + [(byte_id, tag_ids[name], position) for position, name in enumerate(tags)], + ) + conn.execute(_BUMP_REVISION, (byte_id,)) + conn.execute(_BUMP_INDEX) + return conn.execute(_REVISION, (byte_id,)).fetchone()[0] + + def _current_set(self, conn, sha256, byte_id, revision) -> dict: + return { + "sha256": sha256, + "tags": self._byte_tags(conn, byte_id), + "tag_revision": revision, + } + + def read_image_tags(self, locator: str, sha256: str) -> dict: + """The stored set of these exact bytes, or an empty set at revision 0. + + ``prior_count`` is the number of other byte records at the locator + (earlier or later bytes of the same file). + """ + _validate_locator(locator) + _validate_digest(sha256) + with self._annotation_read() as conn: + row = conn.execute(_READ, (sha256, sha256, locator)).fetchone() + if row is None: + return {"tags": [], "tag_revision": 0, "prior_count": 0} + tags = self._byte_tags(conn, row[0]) if row[0] is not None else [] + return {"tags": tags, "tag_revision": row[1] or 0, "prior_count": row[2]} + + def replace_image_tags( + self, + *, + locator: str, + sha256: str, + size: int, + mtime_ns: int, + expected_epoch: str, + expected_tag_revision: int, + tags: list, + ) -> dict: + """Replace the set of these exact bytes, guarded by epoch and revision. + + Creates the owner and byte record when the bytes are new (an explicit + empty set is a real save). The identical set at the current revision, + or one behind it (a lost-response retry), is ``replayed`` and writes + nothing; any other revision mismatch is ``tag-conflict`` with the + current set and no writes. + """ + _validate_target(locator, sha256, size, mtime_ns) + _validate_epoch(expected_epoch) + _validate_int(expected_tag_revision, "expected_tag_revision", 0) + tags = normalize_image_tags(tags) + with self._pair_transaction() as conn: + epoch = self._library_epoch(conn, expected_epoch) + now = utc_now_iso() + file_id = self._annotation_owner(conn, locator, now) + byte_id, revision, created, refingerprinted = self._annotation_byte( + conn, file_id, sha256, size, mtime_ns, now + ) + current = self._byte_tags(conn, byte_id) + if ( + not created + and current == tags + and expected_tag_revision in (revision, revision - 1) + ): + if refingerprinted: + conn.execute(_BUMP_INDEX) + return _receipt(epoch, sha256, current, revision, True) + if expected_tag_revision != revision: + raise AnnotationCommandError( + "tag-conflict", + epoch, + {"sha256": sha256, "tags": current, "tag_revision": revision}, + ) + revision = self._write_image_tags(conn, byte_id, tags) + return _receipt(epoch, sha256, tags, revision, False) + + def add_image_tags( + self, + *, + locator: str, + sha256: str, + size: int, + mtime_ns: int, + expected_epoch: str, + tags: list, + ) -> dict: + """Union ``tags`` into the set of these exact bytes, keeping stored spelling. + + ``added`` counts the names that were new (case-insensitively). A union + over ``MAX_IMAGE_TAGS`` is ``tag-limit`` with the current set and no + writes; a union that adds nothing to an existing record is ``replayed``. + """ + _validate_target(locator, sha256, size, mtime_ns) + _validate_epoch(expected_epoch) + tags = normalize_image_tags(tags) + with self._pair_transaction() as conn: + epoch = self._library_epoch(conn, expected_epoch) + now = utc_now_iso() + file_id = self._annotation_owner(conn, locator, now) + byte_id, revision, created, refingerprinted = self._annotation_byte( + conn, file_id, sha256, size, mtime_ns, now + ) + current = self._byte_tags(conn, byte_id) + present = {name.casefold() for name in current} + added = [name for name in tags if name.casefold() not in present] + merged = current + added + if len(merged) > MAX_IMAGE_TAGS: + raise AnnotationCommandError( + "tag-limit", + epoch, + {"sha256": sha256, "tags": current, "tag_revision": revision}, + ) + if not created and not added: + if refingerprinted: + conn.execute(_BUMP_INDEX) + return {**_receipt(epoch, sha256, current, revision, True), "added": 0} + revision = self._write_image_tags(conn, byte_id, merged) + return { + **_receipt(epoch, sha256, merged, revision, False), + "added": len(added), + } + + def refresh_image_fingerprint( + self, locator: str, sha256: str, size: int, mtime_ns: int + ) -> bool: + """Record that these exact bytes were just seen with this stat. + + True when a byte record was refreshed; nothing is created and no tag + or tag revision changes. The image index revision moves because the + record's currency for cheap reads may have changed. + """ + _validate_target(locator, sha256, size, mtime_ns) + with self.model.get_connection() as conn: + cursor = conn.execute( + _REFRESH_BY_LOCATOR, (size, mtime_ns, utc_now_iso(), sha256, locator) + ) + if cursor.rowcount > 0: + conn.execute(_BUMP_INDEX) + return cursor.rowcount > 0 + + def current_image_tags(self, stats: dict) -> dict: + """Image tags of the files whose current ``(size, mtime_ns)`` is verified. + + ``stats`` maps locators to their present fingerprint. A locator is + answered (possibly with an empty set) only when a byte record carries + exactly that verified fingerprint; the most recently verified record + wins when several do. Lookups go in chunks of 500 locators. + """ + if not isinstance(stats, dict): + raise ValueError( + "stats must map locators to (size, mtime_ns) integer pairs" + ) + for locator, fingerprint in stats.items(): + _validate_fingerprint(locator, fingerprint) + locators = list(stats) + chosen: dict[str, tuple[tuple[str, str], str]] = {} + result: dict[str, list[str]] = {} + with self._annotation_read() as conn: + for start in range(0, len(locators), _LOOKUP_CHUNK): + chunk = locators[start : start + _LOOKUP_CHUNK] + marks = ",".join("?" * len(chunk)) + for locator, byte_id, size, mtime_ns, verified_at in conn.execute( + _FINGERPRINTS.format(marks=marks), chunk + ): + if (size, mtime_ns) != stats[locator]: + continue + key = (verified_at, byte_id) + if locator not in chosen or key > chosen[locator][0]: + chosen[locator] = (key, byte_id) + owners = {byte_id: locator for locator, (_key, byte_id) in chosen.items()} + result = {locator: [] for locator in chosen} + byte_ids = list(owners) + for start in range(0, len(byte_ids), _LOOKUP_CHUNK): + chunk = byte_ids[start : start + _LOOKUP_CHUNK] + marks = ",".join("?" * len(chunk)) + for byte_id, name in conn.execute(_TAGS_OF.format(marks=marks), chunk): + result[owners[byte_id]].append(name) + return result + + def iter_image_tag_records(self) -> Iterator[tuple[str, int, int, list[str]]]: + """Stream every byte record as ``(locator, size, mtime_ns, tags)``. + + One SELECT ordered by locator, then by the bytes' first sighting, with + each set's names in position order. + """ + conn = self.model.get_connection() + current: tuple | None = None + for locator, byte_id, size, mtime_ns, name in conn.execute(_RECORDS): + if current is None or current[0] != byte_id: + if current is not None: + yield current[1] + current = (byte_id, (locator, size, mtime_ns, [])) + if name is not None: + current[1][3].append(name) + if current is not None: + yield current[1] diff --git a/database/models.py b/database/models.py index 8e9e82f..5268577 100644 --- a/database/models.py +++ b/database/models.py @@ -17,6 +17,7 @@ from typing import Optional, Set, Tuple from .library_schema import library_schema_error, migrate_library_state from .image_favorites import image_favorites_schema_error, migrate_image_favorites +from .annotation_schema import annotation_schema_error, migrate_annotation_state from .saved_views import migrate_saved_views, saved_views_schema_error from .execution_history_schema import ( history_schema_error, @@ -302,6 +303,7 @@ class PromptModel: migrate_library_state(conn) migrate_image_favorites(conn) migrate_saved_views(conn) + migrate_annotation_state(conn) def _migrate_prompt_pairs(self, conn: sqlite3.Connection) -> None: conn.execute("SAVEPOINT prompt_pair_schema") @@ -1021,9 +1023,9 @@ class PromptModel: Args: path: File to inspect require_library_state: Also require the tables every current - database has (Library state, image favorites and saved - views), as a migrated restore stage must; old backups may - lack them. + database has (Library state, image favorites, saved views + and image annotations), as a migrated restore stage must; + old backups may lack them. Returns: (ok, reason): reason is "ok" when the file is acceptable, otherwise @@ -1088,6 +1090,11 @@ class PromptModel: ) if views_error: return False, views_error + annotation_error = annotation_schema_error( + conn, allow_missing=not require_library_state + ) + if annotation_error: + return False, annotation_error return True, "ok" except sqlite3.Error as e: return False, f"Invalid SQLite database: {e}" diff --git a/database/operations.py b/database/operations.py index 26b28d1..8fa24dd 100644 --- a/database/operations.py +++ b/database/operations.py @@ -12,6 +12,7 @@ from .models import PromptModel, normalize_image_path, utc_now_iso from .pairs import PairDependencyError, PairOperations from .library import LibraryOperations, titled from .image_favorites import ImageFavoriteOperations +from .annotations import AnnotationOperations from .saved_views import SavedViewOperations from .library_commands import LibraryCommandOperations from .vocabulary import VocabularyOperations @@ -322,6 +323,7 @@ class PromptDatabase( ExecutionOutputOperations, HistoryReadOperations, ImageFavoriteOperations, + AnnotationOperations, SavedViewOperations, ): """Database operations class for managing prompts.""" @@ -1084,16 +1086,39 @@ class PromptDatabase( } @staticmethod + def _bump_image_tag_bytes(conn: sqlite3.Connection, tag_id: int) -> int: + """Move the tag revision of every image set that carries *tag_id*. + + Vocabulary rename, merge and delete change what an image's set reads + as, so a Tag dialog opened before the change must see ``tag-conflict`` + on Save instead of recreating the old name. Bumps the image index + revision once when any set was affected; returns that count. + """ + cursor = conn.execute( + "UPDATE media_annotation_bytes SET tag_revision = tag_revision + 1 " + "WHERE id IN (SELECT byte_id FROM media_image_tags WHERE tag_id = ?)", + (tag_id,), + ) + if cursor.rowcount > 0: + conn.execute( + "UPDATE media_annotation_metadata SET " + "image_index_revision = image_index_revision + 1 WHERE singleton_id = 1" + ) + return cursor.rowcount + + @classmethod def _move_tag_links( - conn: sqlite3.Connection, source_id: int, target_id: int + cls, conn: sqlite3.Connection, source_id: int, target_id: int ) -> None: - """Repoint component and pair links from source to target. + """Repoint component, pair and image links from source to target. Owners that already carry the target keep that one link; the source's - leftover duplicates are removed. The source registry row is left for the - caller to delete. + leftover duplicates are removed (an image set keeps the target's + position). The source registry row is left for the caller to delete. + Image sets that carried the source get a new tag revision first. """ - for junction in ("prompt_tags", "prompt_pair_tags"): + cls._bump_image_tag_bytes(conn, source_id) + for junction in ("prompt_tags", "prompt_pair_tags", "media_image_tags"): conn.execute( f"UPDATE OR IGNORE {junction} SET tag_id = ? WHERE tag_id = ?", (target_id, source_id), @@ -1146,6 +1171,7 @@ class PromptDatabase( conn.execute("DELETE FROM tags WHERE id = ?", (old_tag_id,)) else: # Simple rename + self._bump_image_tag_bytes(conn, old_tag_id) conn.execute( "UPDATE tags SET name = ? WHERE id = ?", (new_name, old_tag_id) ) @@ -1187,6 +1213,7 @@ class PromptDatabase( affected = cursor.fetchone()["c"] # CASCADE delete handles prompt_tags entries + self._bump_image_tag_bytes(conn, tag_id) conn.execute("DELETE FROM tags WHERE id = ?", (tag_id,)) conn.commit() @@ -1239,7 +1266,10 @@ class PromptDatabase( pair_count = conn.execute( "SELECT COUNT(*) FROM prompt_pair_tags WHERE tag_id = ?", (src_id,) ).fetchone()[0] - if src_count or pair_count: + image_count = conn.execute( + "SELECT COUNT(*) FROM media_image_tags WHERE tag_id = ?", (src_id,) + ).fetchone()[0] + if src_count or pair_count or image_count: self._move_tag_links(conn, src_id, target_id) tags_merged += 1 affected += src_count diff --git a/database/vocabulary.py b/database/vocabulary.py index a3152d9..5bca3fb 100644 --- a/database/vocabulary.py +++ b/database/vocabulary.py @@ -3,6 +3,9 @@ A count is the number of eligible Library parents that carry the value directly or through one active saved pair, which is exactly what ``list_library(tags=[name])`` / ``list_library(category=value)`` return. +A tag row's ``image_count`` is the number of files (private locators) with any +byte record carrying the tag (#249), so a file counts once however many of its +byte versions carry it, and equal bytes at two paths count twice. Tag commands run in one ``_pair_transaction``: the storage epoch is checked first, then the exact ``{id, name}`` identity, and the receipt is read on the @@ -26,6 +29,11 @@ except ImportError: from utils.validators import validate_category, validate_tags _CATEGORY_DOMAIN = "prompt-manager-vocabulary-category:v1\n" +# Files carrying a tag through any of their byte records (#249 image tags). +_IMAGE_FILES = ( + "SELECT mit.tag_id,COUNT(DISTINCT b.file_id) FROM media_image_tags mit " + "JOIN media_annotation_bytes b ON b.id=mit.byte_id" +) def _tag_rows(page: str) -> str: @@ -40,9 +48,13 @@ def _tag_rows(page: str) -> str: "WHERE pt.tag_id IN (SELECT id FROM page) AND " + _ELIGIBLE + " " "UNION SELECT ppt.tag_id,pp.positive_prompt_id FROM prompt_pairs pp " "JOIN prompt_pair_tags ppt ON ppt.pair_id=pp.id " - "WHERE pp.archived=0 AND ppt.tag_id IN (SELECT id FROM page)) " - "SELECT page.id,page.name,COUNT(DISTINCT owners.pid) AS prompt_count " + "WHERE pp.archived=0 AND ppt.tag_id IN (SELECT id FROM page))," + "images(tag_id,n) AS (" + _IMAGE_FILES + " " + "WHERE mit.tag_id IN (SELECT id FROM page) GROUP BY mit.tag_id) " + "SELECT page.id,page.name,COUNT(DISTINCT owners.pid) AS prompt_count," + "COALESCE(MAX(images.n),0) AS image_count " "FROM page LEFT JOIN owners ON owners.tag_id=page.id " + "LEFT JOIN images ON images.tag_id=page.id " "GROUP BY page.id ORDER BY page.id" ) @@ -61,7 +73,10 @@ _TAG_IMPACT = ( "AS negative_variants," "(SELECT COUNT(*) FROM prompt_pairs pp JOIN prompt_pair_tags ppt " "ON ppt.pair_id=pp.id WHERE ppt.tag_id=:tag AND pp.archived!=0) " - "AS archived_negative_variants" + "AS archived_negative_variants," + "(SELECT COUNT(DISTINCT b.file_id) FROM media_image_tags mit " + "JOIN media_annotation_bytes b ON b.id=mit.byte_id WHERE mit.tag_id=:tag) " + "AS images" ) _MAX_ID = 2**63 - 1 @@ -205,7 +220,16 @@ class VocabularyOperations: "SELECT COUNT(DISTINCT pair_id) FROM prompt_pair_tags WHERE tag_id=?", (tag_id,), ).fetchone()[0] - return {"components_changed": components, "pairs_changed": pairs} + images = conn.execute( + "SELECT COUNT(DISTINCT b.file_id) FROM media_image_tags mit " + "JOIN media_annotation_bytes b ON b.id=mit.byte_id WHERE mit.tag_id=?", + (tag_id,), + ).fetchone()[0] + return { + "components_changed": components, + "pairs_changed": pairs, + "images_changed": images, + } def create_vocabulary_tag(self, name: str, *, expected_epoch: str) -> dict: """Insert an unused registry row, or reuse the exact existing name.""" @@ -250,6 +274,7 @@ class VocabularyOperations: "SELECT id FROM tags WHERE name=?", (new_name,) ).fetchone() if target is None: + self._bump_image_tag_bytes(conn, tag_id) conn.execute("UPDATE tags SET name=? WHERE id=?", (new_name, tag_id)) result_id = tag_id else: @@ -275,6 +300,7 @@ class VocabularyOperations: epoch = self._library_epoch(conn, expected_epoch) self._vocabulary_identity(conn, epoch, tag_id, name) counts = self._tag_owner_counts(conn, tag_id) + self._bump_image_tag_bytes(conn, tag_id) conn.execute("DELETE FROM tags WHERE id=?", (tag_id,)) return {"library_epoch": epoch, "deleted": True, **counts} diff --git a/tests/browser/full-inspector.spec.mjs b/tests/browser/full-inspector.spec.mjs index 5c82d3e..a22fb3a 100644 --- a/tests/browser/full-inspector.spec.mjs +++ b/tests/browser/full-inspector.spec.mjs @@ -804,6 +804,7 @@ test('slideshow preference failures leave native playback usable',async({page,re }); test('full inspector attaches matched theme and mobile evidence with reachable content',async({page,request},testInfo)=>{ + test.setTimeout(120_000); // 32 evidence captures across four viewports, like the concept comparisons await expect(await request.post('/__test__/reset',{data:{full_inspector:true,include_video:true}})).toBeOK(); await page.goto('/prompt_manager/workspace#/images'); const measurements=[]; diff --git a/tests/browser/server.py b/tests/browser/server.py index 77dc0f4..c5ae192 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -189,6 +189,9 @@ def create_app(root: Path) -> web.Application: for table in ( "saved_views", "image_favorites", + "media_image_tags", + "media_annotation_bytes", + "media_annotation_files", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace-images.spec.mjs b/tests/browser/workspace-images.spec.mjs index c600759..63d04bf 100644 --- a/tests/browser/workspace-images.spec.mjs +++ b/tests/browser/workspace-images.spec.mjs @@ -1512,6 +1512,25 @@ test('Refine opened before discovery settles lists the discovered sources and fo }); // #249 final review m1: as in the old field row, a Source offers only its own folders. +test('a search reload keeps the Filters label node, so a click landing during the reload is not lost', async ({ page, request }) => { + // WebKit drops a click whose mousedown hit a text node that was replaced before mouseup; the reload reports the + // (unchanged) active count, so the label must keep its node. + await expect(await request.post('/__test__/reset', { data: { extra_images: 80 } })).toBeOK(); + await page.goto(workspace); + const button = filtersButton(page); + await button.click();await expect(refineDialog(page).getByLabel('Folder', { exact: true }).locator('option[value="bulk"]')).toHaveCount(1);await page.keyboard.press('Escape'); + await refine(page, { Source: '0', Folder: 'bulk' }); + await expect(button).toHaveText('Filters · 2'); + const label = await button.evaluateHandle(node => node.firstChild); + const reloaded = page.waitForResponse(response => { + const url = new URL(response.url());return imagePages(url) && url.searchParams.get('q') === 'png'; + }); + await search(page).fill('png'); + await reloaded; + await expect(button).toHaveText('Filters · 2'); + expect(await button.evaluate((node, text) => node.firstChild === text && text.isConnected, label)).toBe(true); +}); + test('Refine offers only the chosen source folders', async ({ page, request }) => { await expect(await request.post('/__test__/add-nested-image')).toBeOK(); await page.goto(workspace); diff --git a/tests/browser/workspace_concept_fixture.py b/tests/browser/workspace_concept_fixture.py index df9d5e3..ea48be8 100644 --- a/tests/browser/workspace_concept_fixture.py +++ b/tests/browser/workspace_concept_fixture.py @@ -170,6 +170,9 @@ _TABLES = ( "execution_context_texts", "saved_views", "image_favorites", + "media_image_tags", + "media_annotation_bytes", + "media_annotation_files", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_library_editor_fixture.py b/tests/browser/workspace_library_editor_fixture.py index 2ea38b1..7b852a1 100644 --- a/tests/browser/workspace_library_editor_fixture.py +++ b/tests/browser/workspace_library_editor_fixture.py @@ -28,6 +28,9 @@ EDITOR_METADATA = { _TABLES = ( "saved_views", "image_favorites", + "media_image_tags", + "media_annotation_bytes", + "media_annotation_files", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_visual_fixture.py b/tests/browser/workspace_visual_fixture.py index 9ba1d48..9a7cdad 100644 --- a/tests/browser/workspace_visual_fixture.py +++ b/tests/browser/workspace_visual_fixture.py @@ -114,6 +114,9 @@ def seed_workspace_visual(api: PromptManagerAPI, root: Path) -> None: for table in ( "saved_views", "image_favorites", + "media_image_tags", + "media_annotation_bytes", + "media_annotation_files", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_vocabulary_fixture.py b/tests/browser/workspace_vocabulary_fixture.py index f0e047d..69e0fb9 100644 --- a/tests/browser/workspace_vocabulary_fixture.py +++ b/tests/browser/workspace_vocabulary_fixture.py @@ -81,6 +81,9 @@ MAX_EXTRA_CATEGORIES = 200 _TABLES = ( "saved_views", "image_favorites", + "media_image_tags", + "media_annotation_bytes", + "media_annotation_files", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/js/workspace-toolbar.test.js b/tests/js/workspace-toolbar.test.js index 789569d..d825d6a 100644 --- a/tests/js/workspace-toolbar.test.js +++ b/tests/js/workspace-toolbar.test.js @@ -765,6 +765,20 @@ test('the Filters button shows the active refinement count', async () => { assert.equal(filtersButton.textContent, 'Filters'); for (const bad of [-1, 1.5, '1', null, Number.NaN]) assert.throws(() => toolbar.setActiveFilters(bad), TypeError); assert.equal(filtersButton.textContent, 'Filters'); + // An unchanged count leaves the label's text node alone: every listing reload reports the count, and WebKit + // drops a click whose mousedown hit a text node that was replaced before mouseup. + let writes = 0; + const proto = Object.getPrototypeOf(filtersButton); + const descriptor = Object.getOwnPropertyDescriptor(proto, 'textContent'); + Object.defineProperty(filtersButton, 'textContent', { + get() { return descriptor.get.call(this); }, + set(text) { writes += 1; descriptor.set.call(this, text); }, + configurable: true, + }); + toolbar.setActiveFilters(0); + assert.equal(writes, 0); + toolbar.setActiveFilters(3); + assert.deepEqual([writes, filtersButton.textContent], [1, 'Filters · 3']); assert.throws(() => mountToolbar(new dom.Element('div'), { scope, categories: [], sorts: [['newest', 'Recently added']], activeFilters: -1, onChange() {}, onFilters() {}, onSaveView() {} }), TypeError); } finally { diff --git a/tests/js/workspace-vocabulary-state.test.js b/tests/js/workspace-vocabulary-state.test.js index 51706fa..6be188c 100644 --- a/tests/js/workspace-vocabulary-state.test.js +++ b/tests/js/workspace-vocabulary-state.test.js @@ -507,10 +507,49 @@ test('category preview text states the real impact and whether it renames or con '“Landscapes” is used by 2 other saved prompts. Apply renames it to “Scenery”.'); }); +test('accepts image_count and names images in the delete text', async () => { + const {createVocabularyStream, createVocabularyCommand, readTagImpact, deleteTagText} = await helpers; + const rows = [{...tag(1, 'mist', 2), image_count: 3}, tag(2, 'dusk', 0)]; + const stream = createVocabularyStream({api: queued([tagPage(rows)]).api, kind: 'tags'}); + assert.equal(await stream.loadMore(), 2); + assert.deepEqual(stream.state.items.map(row => row.image_count), [3, undefined]); + for (const bad of [{...tag(1), image_count: -1}, {...tag(1), image_count: '3'}, {...tag(1), image_count: 1.5}]) { + const rejected = createVocabularyStream({api: queued([tagPage([bad])]).api, kind: 'tags'}); + await rejected.loadMore(); + assert.ok(rejected.state.error instanceof TypeError, JSON.stringify(bad)); + assert.deepEqual(rejected.state.items, []); + } + const withImages = {...impact, images: 2}; + const api = {get: async () => ({success: true, library_epoch: EPOCH, impact: withImages})}; + assert.deepEqual((await readTagImpact(api, {id: 4, name: 'mist'})).impact, withImages); + await assert.rejects(readTagImpact({get: async () => ({success: true, library_epoch: EPOCH, impact: {...impact, images: -2}})}, + {id: 4, name: 'mist'}), TypeError); + const command = createVocabularyCommand({api: { + patch: async () => renameReceipt({tag: {...tag(1, 'dusk', 2), image_count: 1}, images_changed: 1}), + post: async () => ({success: true, library_epoch: EPOCH, deleted: true, components_changed: 0, pairs_changed: 0, images_changed: 2}), + }}); + await command.renameTag({id: 1, name: 'dawn', prompt_count: 2}, 'dusk', EPOCH); + assert.equal(command.state.phase, 'done'); + assert.equal(command.state.receipt.images_changed, 1); + command.reset(); + await command.deleteTag({id: 1, name: 'dusk'}, EPOCH); + assert.equal(command.state.phase, 'done'); + const malformed = createVocabularyCommand({api: {post: async () => ({success: true, library_epoch: EPOCH, deleted: true, + components_changed: 0, pairs_changed: 0, images_changed: 'two'})}}); + await malformed.deleteTag({id: 1, name: 'dusk'}, EPOCH); + assert.equal(malformed.state.phase, 'unknown-outcome'); + assert.equal(deleteTagText('mist', withImages), + 'Deleting “mist” removes the name and its links from 3 positive prompts, 3 negative variants (1 archived), 1 other saved prompt and 2 images. The prompts, negative variants, images and history themselves are kept.'); + assert.equal(deleteTagText('mist', {...zero, images: 1}), + 'Deleting “mist” removes the name and its links from 1 image. The prompts, negative variants, images and history themselves are kept.'); + assert.equal(deleteTagText('unused', {...zero, images: 0}), + '“unused” is not linked to any saved prompt or image. Deleting it removes only the name.'); +}); + test('delete text names the tag, its real links and what stays unchanged', async () => { const {deleteTagText} = await helpers; assert.equal(deleteTagText('mist', impact), 'Deleting “mist” removes the name and its links from 3 positive prompts, 3 negative variants (1 archived) and 1 other saved prompt. The prompts, negative variants, images and history themselves are kept.'); assert.equal(deleteTagText('unused', zero), - '“unused” is not linked to any saved prompt. Deleting it removes only the name.'); + '“unused” is not linked to any saved prompt or image. Deleting it removes only the name.'); }); diff --git a/tests/test_annotation_schema.py b/tests/test_annotation_schema.py new file mode 100644 index 0000000..cc37b9a --- /dev/null +++ b/tests/test_annotation_schema.py @@ -0,0 +1,602 @@ +"""Image annotation schema group (#249): install, validation and restore safety.""" + +import hashlib +import re +import sqlite3 +import tempfile +import unittest +from contextlib import closing +from pathlib import Path +from unittest.mock import patch + +try: + from tests.open_handles import assert_closed +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + +from database.models import PromptModel + +try: # guarded so a missing module fails assertions, not test loading + from database.annotation_schema import ( + annotation_schema_error, + migrate_annotation_state, + ) +except ImportError: + annotation_schema_error = None + migrate_annotation_state = None + +METADATA = "media_annotation_metadata" +FILES = "media_annotation_files" +BYTES = "media_annotation_bytes" +IMAGE_TAGS = "media_image_tags" +TABLES = (METADATA, FILES, BYTES, IMAGE_TAGS) +# The exact DDL the plan pins; the module must produce precisely these objects. +DDL = { + METADATA: ( + f"CREATE TABLE {METADATA}(singleton_id INTEGER PRIMARY KEY " + "CHECK(singleton_id=1), schema_version INTEGER NOT NULL " + "CHECK(schema_version=1), image_index_revision INTEGER NOT NULL DEFAULT 0 " + "CHECK(image_index_revision>=0))" + ), + FILES: ( + f"CREATE TABLE {FILES}(id TEXT PRIMARY KEY CHECK(length(id)=36), " + "locator TEXT NOT NULL UNIQUE, created_at TEXT NOT NULL)" + ), + BYTES: ( + f"CREATE TABLE {BYTES}(id TEXT PRIMARY KEY CHECK(length(id)=36), " + f"file_id TEXT NOT NULL REFERENCES {FILES}(id) ON DELETE CASCADE, " + "sha256 TEXT NOT NULL CHECK(length(sha256)=64 AND sha256 NOT GLOB " + "'*[^0-9a-f]*'), size INTEGER NOT NULL CHECK(size>=0), " + "mtime_ns INTEGER NOT NULL, verified_at TEXT NOT NULL, " + "first_seen_at TEXT NOT NULL, tag_revision INTEGER NOT NULL DEFAULT 0 " + "CHECK(tag_revision>=0), UNIQUE(file_id, sha256))" + ), + IMAGE_TAGS: ( + f"CREATE TABLE {IMAGE_TAGS}(byte_id TEXT NOT NULL REFERENCES {BYTES}(id) " + "ON DELETE CASCADE, tag_id INTEGER NOT NULL REFERENCES tags(id) " + "ON DELETE CASCADE, position INTEGER NOT NULL CHECK(position>=0), " + "PRIMARY KEY(byte_id, tag_id))" + ), +} +INDEXES = { + "idx_media_annotation_bytes_fingerprint": ( + f"CREATE INDEX idx_media_annotation_bytes_fingerprint ON {BYTES}" + "(file_id, size, mtime_ns)" + ), + "idx_media_image_tags_tag": ( + f"CREATE INDEX idx_media_image_tags_tag ON {IMAGE_TAGS}(tag_id)" + ), + "idx_media_image_tags_position": ( + f"CREATE UNIQUE INDEX idx_media_image_tags_position ON {IMAGE_TAGS}" + "(byte_id, position)" + ), +} +AUTOINDEXES = {FILES: 2, BYTES: 2, IMAGE_TAGS: 1} +LEGACY_TABLES = ( + "prompts", + "generated_images", + "tags", + "prompt_tags", + "prompt_pair_tags", + "image_favorites", +) +FILE_ID = "0f8e4b2a-1c3d-4e5f-8a9b-0c1d2e3f4a5b" +BYTE_ID = "6a1b2c3d-4e5f-4a6b-8c9d-0e1f2a3b4c5d" +OTHER_ID = "7b2c3d4e-5f6a-4b7c-9d0e-1f2a3b4c5d6e" +DIGEST = hashlib.sha256(b"img").hexdigest() +NOW = "2026-10-22T00:00:00.000000+00:00" +MAX_LOCATOR = 4096 + + +def _compact(sql): + return re.sub(r"\s+", "", sql or "").lower() + + +def expected_catalog(): + objects = {("table", name, _compact(sql)) for name, sql in DDL.items()} + objects |= {("index", name, _compact(sql)) for name, sql in INDEXES.items()} + objects |= { + ("index", f"sqlite_autoindex_{table}_{n}", "") + for table, count in AUTOINDEXES.items() + for n in range(1, count + 1) + } + return objects + + +def catalog(conn): + names = tuple(TABLES) + tuple(INDEXES) + marks = ",".join("?" * len(names)) + table_marks = ",".join("?" * len(TABLES)) + return { + (kind, name, _compact(sql)) + for kind, name, sql in conn.execute( + "SELECT type,name,sql FROM sqlite_master " + f"WHERE name IN ({marks}) OR tbl_name IN ({table_marks})", # nosec B608 + names + tuple(TABLES), + ) + } + + +def insert_group_rows(conn, tag_name="img-only", locator="/out/img.png"): + """One file, one byte record and one image tag carrying *tag_name*.""" + tag_id = conn.execute("INSERT INTO tags(name) VALUES(?)", (tag_name,)).lastrowid + conn.execute(f"INSERT INTO {FILES} VALUES(?,?,?)", (FILE_ID, locator, NOW)) + conn.execute( + f"INSERT INTO {BYTES} VALUES(?,?,?,?,?,?,?,?)", + (BYTE_ID, FILE_ID, DIGEST, 3, 1_700_000_000_000_000_000, NOW, NOW, 1), + ) + conn.execute(f"INSERT INTO {IMAGE_TAGS} VALUES(?,?,0)", (BYTE_ID, tag_id)) + return tag_id + + +def rows(conn, sql, *params): + """Plain tuples whatever row factory *conn* uses.""" + return [tuple(row) for row in conn.execute(sql, params)] + + +def image_tags(conn): + return rows( + conn, + f"SELECT t.name, it.position FROM {IMAGE_TAGS} it " + "JOIN tags t ON t.id = it.tag_id ORDER BY it.position", + ) + + +def metadata(conn): + return rows(conn, f"SELECT schema_version,image_index_revision FROM {METADATA}") + + +def table_rows(conn, tables): + found = {} + for table in tables: + present = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (table,) + ).fetchone() + if present: + found[table] = rows(conn, f"SELECT * FROM {table}") # nosec B608 + return found + + +def populate_legacy(conn): + conn.execute( + "INSERT INTO prompts(id,text,hash,category,notes,run_count) VALUES" + "(7,'positive','p','keep','notes',12)" + ) + conn.execute("INSERT INTO prompts(id,text,hash) VALUES(9,'negative','n')") + conn.execute( + "INSERT INTO generated_images(id,prompt_id,image_path,filename,file_path) " + "VALUES(4,7,'/tmp/image.png','image.png','/tmp/image.png')" + ) + conn.execute("INSERT INTO tags(id,name) VALUES(3,'keep-tag')") + conn.execute("INSERT INTO prompt_tags VALUES(7,3)") + conn.execute( + "INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id) " + "VALUES(6,7,9)" + ) + conn.execute("INSERT INTO prompt_pair_tags VALUES(6,3)") + conn.execute("INSERT INTO image_favorites VALUES('root','a.png',?)", (NOW,)) + + +class AnnotationSchemaTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.addCleanup(assert_closed, self, self.temp.name) + self.addCleanup(PromptModel.reset_schema_cache) + self.models = [] + self.addCleanup(self.close_models) + self.path = str(Path(self.temp.name) / "live.db") + + def close_models(self): + for model in self.models: + model.close_all() + + def model(self, path=None): + model = PromptModel(path or self.path) + self.models.append(model) + return model + + def require_module(self): + self.assertIsNotNone( + annotation_schema_error, "database.annotation_schema is not implemented" + ) + + def drop_group(self, path): + """Remove the group from a closed database file (an old backup).""" + with closing(sqlite3.connect(path)) as raw, raw: + raw.execute("PRAGMA foreign_keys=OFF") + for table in reversed(TABLES): + raw.execute(f"DROP TABLE {table}") + + def test_fresh_install_and_reopen_is_idempotent(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + self.assertEqual(catalog(conn), expected_catalog()) + self.assertIsNone(annotation_schema_error(conn, allow_missing=False)) + self.assertEqual(metadata(conn), [(1, 0)]) + insert_group_rows(conn) + self.assertEqual( + PromptModel.verify_database_file(self.path, require_library_state=True), + (True, "ok"), + ) + model.close_all() + for _ in range(2): + PromptModel.reset_schema_cache() + model = self.model() + with model.get_connection() as conn: + self.assertEqual(catalog(conn), expected_catalog()) + self.assertIsNone(annotation_schema_error(conn, allow_missing=False)) + self.assertEqual(metadata(conn), [(1, 0)]) + self.assertEqual(image_tags(conn), [("img-only", 0)]) + # Running the migration again on a current database changes nothing. + migrate_annotation_state(conn) + self.assertEqual(metadata(conn), [(1, 0)]) + self.assertEqual(image_tags(conn), [("img-only", 0)]) + model.close_all() + # A tag delete and a file delete cascade into the junction and bytes. + PromptModel.reset_schema_cache() + model = self.model() + with model.get_connection() as conn: + conn.execute("DELETE FROM tags WHERE name='img-only'") + self.assertEqual(image_tags(conn), []) + conn.execute(f"DELETE FROM {FILES}") + self.assertEqual(rows(conn, f"SELECT * FROM {BYTES}"), []) + self.assertIsNone(annotation_schema_error(conn, allow_missing=False)) + + def test_old_database_without_group_migrates_without_touching_legacy_rows(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + populate_legacy(conn) + model.close_all() + self.drop_group(self.path) + with closing(sqlite3.connect(self.path)) as raw: + self.assertIsNone(annotation_schema_error(raw)) + self.assertEqual( + annotation_schema_error(raw, allow_missing=False), + "Missing image annotation schema", + ) + legacy_rows_before = table_rows(raw, LEGACY_TABLES) + self.assertTrue(PromptModel.verify_database_file(self.path)[0]) + self.assertFalse( + PromptModel.verify_database_file(self.path, require_library_state=True)[0] + ) + PromptModel.reset_schema_cache() + + model = self.model() + + with model.get_connection() as conn: + self.assertEqual(catalog(conn), expected_catalog()) + self.assertIsNone(annotation_schema_error(conn, allow_missing=False)) + self.assertEqual(metadata(conn), [(1, 0)]) + legacy_rows_after = table_rows(conn, LEGACY_TABLES) + # No backfill: existing prompt tags never become image tags. + counts = [ + conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] + for table in (FILES, BYTES, IMAGE_TAGS) + ] + self.assertEqual(counts, [0, 0, 0]) + self.assertEqual(legacy_rows_after, legacy_rows_before) + self.assertEqual( + PromptModel.verify_database_file(self.path, require_library_state=True), + (True, "ok"), + ) + + def tampered(self, name, statements, *, checks_off=False): + """A closed current database whose group was altered after the fact.""" + path = str(Path(self.temp.name) / f"{name.replace(' ', '-')}.db") + model = self.model(path) + with model.get_connection() as conn: + populate_legacy(conn) + insert_group_rows(conn) + model.close_all() + with closing(sqlite3.connect(path)) as raw, raw: + raw.execute("PRAGMA foreign_keys=OFF") + if checks_off: + raw.execute("PRAGMA ignore_check_constraints=ON") + for statement in statements: + if isinstance(statement, tuple): + raw.execute(*statement) + else: + raw.execute(statement) + return path + + def assert_rejected_unrepaired(self, path): + with closing(sqlite3.connect(path)) as raw: + schema_before = raw.execute( + "SELECT type,name,tbl_name,sql FROM sqlite_master ORDER BY name" + ).fetchall() + rows_before = table_rows(raw, TABLES) + for allow_missing in (False, True): + error = annotation_schema_error(raw, allow_missing=allow_missing) + self.assertIsNotNone(error) + self.assertIn("image annotation", error.lower()) + self.assertFalse(PromptModel.verify_database_file(path)[0]) + self.assertFalse( + PromptModel.verify_database_file(path, require_library_state=True)[0] + ) + PromptModel.reset_schema_cache() + with self.assertRaises(sqlite3.DatabaseError): + self.model(path) + PromptModel.reset_schema_cache() + with closing(sqlite3.connect(path)) as raw: + self.assertEqual( + raw.execute( + "SELECT type,name,tbl_name,sql FROM sqlite_master ORDER BY name" + ).fetchall(), + schema_before, + ) + self.assertEqual(table_rows(raw, TABLES), rows_before) + + def test_partial_or_malformed_group_is_rejected_without_repair(self): + self.require_module() + cases = { + "missing metadata": (f"DROP TABLE {METADATA}",), + "missing files": (f"DROP TABLE {FILES}",), + "missing bytes": (f"DROP TABLE {BYTES}",), + "missing image tags": (f"DROP TABLE {IMAGE_TAGS}",), + "reserved view": ( + f"DROP TABLE {FILES}", + f"CREATE VIEW {FILES} AS SELECT 'x' AS id", + ), + "extra trigger": ( + f"CREATE TRIGGER media_annotation_trigger AFTER INSERT ON {FILES} " + "BEGIN SELECT 1; END", + ), + "extra index": ( + f"CREATE INDEX idx_media_annotation_extra ON {FILES}(created_at)", + ), + "extra column": ( + f"DROP TABLE {IMAGE_TAGS}", + DDL[IMAGE_TAGS].replace( + "position INTEGER", "note TEXT, position INTEGER" + ), + ), + "wrong primary key": ( + f"DROP TABLE {FILES}", + DDL[FILES].replace("id TEXT PRIMARY KEY", "id TEXT"), + ), + "wrong foreign key": ( + f"DROP TABLE {BYTES}", + DDL[BYTES].replace("ON DELETE CASCADE", "ON DELETE RESTRICT"), + INDEXES["idx_media_annotation_bytes_fingerprint"], + ), + "wrong check": ( + f"DROP TABLE {METADATA}", + DDL[METADATA].replace("schema_version=1", "schema_version=2"), + f"INSERT INTO {METADATA} VALUES(1,2,0)", + ), + "missing unique": ( + f"DROP TABLE {FILES}", + DDL[FILES].replace("NOT NULL UNIQUE", "NOT NULL"), + ), + "missing index": ("DROP INDEX idx_media_image_tags_tag",), + "non-unique position index": ( + "DROP INDEX idx_media_image_tags_position", + INDEXES["idx_media_image_tags_position"].replace("UNIQUE ", ""), + ), + "wrong index columns": ( + "DROP INDEX idx_media_annotation_bytes_fingerprint", + INDEXES["idx_media_annotation_bytes_fingerprint"].replace( + ", mtime_ns", "" + ), + ), + "orphan tag": (f"INSERT INTO {IMAGE_TAGS} VALUES('{BYTE_ID}',404,1)",), + "orphan byte": ( + ( + f"INSERT INTO {BYTES} VALUES(?,?,?,1,1,?,?,0)", + ( + OTHER_ID, + "9" * 8 + "-0000-4000-8000-" + "0" * 12, + DIGEST, + NOW, + NOW, + ), + ), + ), + "bad file uuid": ( + (f"UPDATE {FILES} SET id=?", ("x" * 36,)), + (f"UPDATE {BYTES} SET file_id=?", ("x" * 36,)), + ), + "uppercase byte uuid": ( + f"UPDATE {BYTES} SET id=upper(id)", + f"UPDATE {IMAGE_TAGS} SET byte_id=upper(byte_id)", + ), + "null file id": ((f"INSERT INTO {FILES} VALUES(NULL,?,?)", ("/n", NOW)),), + "blob ids": ( + f"UPDATE {FILES} SET id=CAST(id AS BLOB)", + f"UPDATE {BYTES} SET file_id=CAST(file_id AS BLOB)", + ), + "no metadata row": (f"DELETE FROM {METADATA}",), + "long created_at": ((f"UPDATE {FILES} SET created_at=?", ("x" * 65,)),), + "blob verified_at": (f"UPDATE {BYTES} SET verified_at=X'00'",), + "text mtime": (f"UPDATE {BYTES} SET mtime_ns='soon'",), + "duplicate position": ( + "DROP INDEX idx_media_image_tags_position", + "INSERT INTO tags(id,name) VALUES(77,'dup')", + f"INSERT INTO {IMAGE_TAGS} VALUES('{BYTE_ID}',77,0)", + ), + "long locator": ( + (f"UPDATE {FILES} SET locator=?", ("x" * (MAX_LOCATOR + 1),)), + ), + "multibyte long locator": ( + (f"UPDATE {FILES} SET locator=?", ("é" * (MAX_LOCATOR + 1),)), + ), + "nul locator bypass": ( + (f"UPDATE {FILES} SET locator=?", ("x\0" + "y" * MAX_LOCATOR,)), + ), + } + for label, statements in cases.items(): + with self.subTest(label): + self.assert_rejected_unrepaired(self.tampered(label, statements)) + unchecked = { + "two metadata rows": (f"INSERT INTO {METADATA} VALUES(2,1,0)",), + "wrong schema version": (f"UPDATE {METADATA} SET schema_version=2",), + "negative index revision": ( + f"UPDATE {METADATA} SET image_index_revision=-1", + ), + "bad digest": (f"UPDATE {BYTES} SET sha256=upper(sha256)",), + "short digest": (f"UPDATE {BYTES} SET sha256=substr(sha256,1,63)",), + "negative size": (f"UPDATE {BYTES} SET size=-1",), + "negative tag revision": (f"UPDATE {BYTES} SET tag_revision=-1",), + "negative position": (f"UPDATE {IMAGE_TAGS} SET position=-1",), + } + for label, statements in unchecked.items(): + with self.subTest(label): + path = self.tampered(label, statements, checks_off=True) + self.assert_rejected_unrepaired(path) + accepted = { + "locator at bound": ("x" * MAX_LOCATOR,), + "multibyte locator at bound": ("é" * MAX_LOCATOR,), + } + for label, (locator,) in accepted.items(): + with self.subTest(label): + path = self.tampered( + label, ((f"UPDATE {FILES} SET locator=?", (locator,)),) + ) + with closing(sqlite3.connect(path)) as raw: + self.assertIsNone(annotation_schema_error(raw, allow_missing=False)) + self.assertEqual( + PromptModel.verify_database_file(path, require_library_state=True), + (True, "ok"), + ) + + def test_install_failure_rolls_back_and_retries(self): + self.require_module() + with patch( + "database.annotation_schema.annotation_schema_error", + side_effect=[None, "injected final validation failure"], + ): + with self.assertRaisesRegex(sqlite3.DatabaseError, "injected"): + self.model() + with closing(sqlite3.connect(self.path)) as raw: + self.assertEqual(catalog(raw), set()) + self.assertIsNotNone( + raw.execute( + "SELECT 1 FROM sqlite_master WHERE name='saved_views'" + ).fetchone() + ) + PromptModel.reset_schema_cache() + model = self.model() + self.assertEqual( + PromptModel.verify_database_file(model.db_path, require_library_state=True), + (True, "ok"), + ) + with model.get_connection() as conn: + self.assertEqual(metadata(conn), [(1, 0)]) + + def test_wal_backup_restore_round_trips_group(self): + self.require_module() + source_path = str(Path(self.temp.name) / "source.db") + source = self.model(source_path) + with source.get_connection() as conn: + conn.execute("PRAGMA wal_autocheckpoint=0") + conn.execute("INSERT INTO prompts(id,text,hash) VALUES(7,'p','p')") + insert_group_rows(conn) + conn.commit() + source_files = rows(conn, f"SELECT * FROM {FILES}") + source_bytes = rows(conn, f"SELECT * FROM {BYTES}") + self.assertGreater(Path(source_path + "-wal").stat().st_size, 0) + live = self.model() + with live.get_connection() as conn: + conn.execute("INSERT INTO prompts(id,text,hash) VALUES(1,'live','live')") + insert_group_rows(conn, tag_name="live-only", locator="/out/live.png") + source_before = Path(source_path).read_bytes() + self.assertEqual( + PromptModel.verify_database_file(source_path, require_library_state=True), + (True, "ok"), + ) + # The rows still sit in the WAL; a backup file must carry them too. + backup_file = str(Path(self.temp.name) / "backup.db") + self.assertTrue(source.backup_database(backup_file)) + with closing(sqlite3.connect(backup_file)) as raw: + self.assertEqual(image_tags(raw), [("img-only", 0)]) + self.assertIsNone(annotation_schema_error(raw, allow_missing=False)) + + safety = live.restore_from_file(source_path) + + self.assertTrue(Path(safety).is_file()) + self.assertEqual(Path(source_path).read_bytes(), source_before) + with live.get_connection() as conn: + restored_image_tags = image_tags(conn) + self.assertEqual(rows(conn, f"SELECT * FROM {FILES}"), source_files) + self.assertEqual(rows(conn, f"SELECT * FROM {BYTES}"), source_bytes) + self.assertEqual(metadata(conn), [(1, 0)]) + self.assertIsNone(annotation_schema_error(conn, allow_missing=False)) + self.assertEqual(rows(conn, "SELECT id FROM prompts"), [(7,)]) + self.assertEqual(restored_image_tags, [("img-only", 0)]) + with closing(sqlite3.connect(safety)) as raw: + self.assertEqual(image_tags(raw), [("live-only", 0)]) + self.assertEqual(raw.execute("SELECT id FROM prompts").fetchall(), [(1,)]) + # The same rows restore from the backup file as well. + PromptModel.reset_schema_cache() + live.restore_from_file(backup_file) + with live.get_connection() as conn: + self.assertEqual(image_tags(conn), [("img-only", 0)]) + PromptModel.close_all_instances() + assert_closed(self, self.temp.name) + + def test_staged_failure_preserves_source_live_safety_and_handles(self): + self.require_module() + live = self.model() + with live.get_connection() as conn: + conn.execute("INSERT INTO prompts(id,text,hash) VALUES(1,'live','live')") + insert_group_rows(conn, tag_name="live-only") + # A source with a malformed group is refused before any change. + bad = self.tampered( + "restore bad", + (f"UPDATE {BYTES} SET sha256=upper(sha256)",), + checks_off=True, + ) + # A staging failure while installing the group must be refused too. + old = str(Path(self.temp.name) / "old.db") + self.model(old).close_all() + self.drop_group(old) + for source, patched in ((bad, False), (old, True)): + with self.subTest(source=Path(source).name): + before = Path(source).read_bytes() + safety = sorted(Path(self.temp.name).glob("live.db.backup_*")) + if patched: + context = patch( + "database.models.migrate_annotation_state", + side_effect=sqlite3.OperationalError("staging failed"), + ) + expected = sqlite3.OperationalError + else: + context = patch.object(PromptModel, "backup_database") + expected = ValueError + with context as hook, self.assertRaises(expected): + live.restore_from_file(source) + if not patched: + hook.assert_not_called() + self.assertEqual(Path(source).read_bytes(), before) + self.assertEqual( + sorted(Path(self.temp.name).glob("live.db.backup_*")), safety + ) + assert_closed(self, source) + with live.get_connection() as conn: + self.assertEqual(image_tags(conn), [("live-only", 0)]) + self.assertEqual(metadata(conn), [(1, 0)]) + self.assertFalse( + PromptModel.verify_database_file(old, require_library_state=True)[0] + ) + self.assertTrue(PromptModel.verify_database_file(old)[0]) + + def test_verify_requires_group_only_with_require_library_state(self): + self.require_module() + self.model().close_all() + self.assertEqual( + PromptModel.verify_database_file(self.path, require_library_state=True), + (True, "ok"), + ) + self.drop_group(self.path) + self.assertEqual(PromptModel.verify_database_file(self.path), (True, "ok")) + self.assertEqual( + PromptModel.verify_database_file(self.path, require_library_state=True), + (False, "Missing image annotation schema"), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_annotations.py b/tests/test_annotations.py new file mode 100644 index 0000000..a5aa600 --- /dev/null +++ b/tests/test_annotations.py @@ -0,0 +1,598 @@ +"""Per-image tag sets (#249 A2): byte-exact storage in the shared vocabulary.""" + +import hashlib +import os +import re +import sqlite3 +import tempfile +import unittest +from unittest.mock import patch + +try: + from tests.open_handles import assert_closed +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + +from database.annotation_schema import annotation_schema_error +from database.library_commands import LibraryCommandError +from database.operations import PromptDatabase +from utils.hashing import generate_prompt_hash + +try: # guarded so a missing module fails assertions, not test loading + from database import annotations as annotations_module + from database.annotations import ( + MAX_IMAGE_TAGS, + AnnotationCommandError, + normalize_image_tags, + ) +except ImportError: + annotations_module = None + MAX_IMAGE_TAGS = None + AnnotationCommandError = None + normalize_image_tags = None + +UUID = re.compile(r"[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}") +DIGEST_A = hashlib.sha256(b"bytes a").hexdigest() +DIGEST_B = hashlib.sha256(b"bytes b").hexdigest() +SIZE = 7 +MTIME = 1_700_000_000_123_456_789 +ANNOTATION_TABLES = ( + "media_annotation_files", + "media_annotation_bytes", + "media_image_tags", + "media_annotation_metadata", +) +LEGACY_TABLES = ( + "tags", + "prompt_tags", + "prompt_pair_tags", + "generated_images", + "image_prompt_pairs", +) + + +class AnnotationTestCase(unittest.TestCase): + def setUp(self): + self.assertIsNotNone( + annotations_module, "database.annotations is not implemented" + ) + directory = tempfile.TemporaryDirectory() + self.addCleanup(directory.cleanup) + self.addCleanup(assert_closed, self, directory.name) + self.root = directory.name + self.db = PromptDatabase(os.path.join(self.root, "annotations.db")) + self.addCleanup(self.db.close_all) + self.epoch = self.db.get_execution_epoch() + + def locator(self, name: str) -> str: + return os.path.normcase(os.path.join(os.path.realpath(self.root), name)) + + def replace(self, loc, digest=DIGEST_A, tags=(), *, revision=0, **over): + options = { + "locator": loc, + "sha256": digest, + "size": SIZE, + "mtime_ns": MTIME, + "expected_epoch": self.epoch, + "expected_tag_revision": revision, + "tags": list(tags) if isinstance(tags, tuple) else tags, + } + options.update(over) + return self.db.replace_image_tags(**options) + + def add(self, loc, digest=DIGEST_A, tags=(), **over): + options = { + "locator": loc, + "sha256": digest, + "size": SIZE, + "mtime_ns": MTIME, + "expected_epoch": self.epoch, + "tags": list(tags) if isinstance(tags, tuple) else tags, + } + options.update(over) + return self.db.add_image_tags(**options) + + def read(self, loc, digest=DIGEST_A): + return self.db.read_image_tags(loc, digest) + + def rows(self, *tables): + """Every row of each table; ``verified_at`` is a fingerprint, not content.""" + conn = self.db.model.get_connection() + result = [] + for table in tables: + columns = [ + row[1] + for row in conn.execute(f"PRAGMA table_info({table})") + if row[1] != "verified_at" + ] + result.append( + [ + tuple(row) + for row in conn.execute( + f"SELECT {','.join(columns)} FROM {table} ORDER BY 1,2" + ) + ] + ) + return result + + def counts(self): + conn = self.db.model.get_connection() + return { + table: conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] + for table in ANNOTATION_TABLES + } + + def index_revision(self): + return ( + self.db.model.get_connection() + .execute("SELECT image_index_revision FROM media_annotation_metadata") + .fetchone()[0] + ) + + def schema_clean(self): + self.assertIsNone( + annotation_schema_error(self.db.model.get_connection(), allow_missing=False) + ) + + +class NormalizeImageTagsTests(unittest.TestCase): + def setUp(self): + self.assertIsNotNone( + normalize_image_tags, "database.annotations is not implemented" + ) + + def test_case_insensitive_dedupe_keeps_first_spelling_and_position(self): + self.assertEqual( + normalize_image_tags([" Forest ", "mist", "forest", "MIST", "Rain"]), + ["Forest", "mist", "Rain"], + ) + self.assertEqual(normalize_image_tags([]), []) + self.assertEqual(MAX_IMAGE_TAGS, 100) + + def test_limit_and_invalid_names_rejected(self): + for bad in ( + None, + "forest", + {"a": 1}, + [1], + [None], + ["a,b"], + [""], + [" "], + ["x" * 101], + [f"t{i}" for i in range(MAX_IMAGE_TAGS + 1)], + ): + with self.subTest(bad=bad), self.assertRaises(ValueError): + normalize_image_tags(bad) + # Duplicates fold before the cap is applied. + many = [f"t{i}" for i in range(MAX_IMAGE_TAGS)] + self.assertEqual(normalize_image_tags(many + ["T0", "T99"]), many) + + +class ReplaceImageTagsTests(AnnotationTestCase): + def test_first_save_creates_owner_and_byte_atomically_including_empty_set(self): + loc = self.locator("a.png") + receipt = self.replace(loc, tags=["Forest", "mist", "forest"]) + self.assertEqual(receipt["tags"], ["Forest", "mist"]) + self.assertEqual(receipt["tag_revision"], 1) + self.assertEqual( + receipt, + { + "library_epoch": self.epoch, + "sha256": DIGEST_A, + "tags": ["Forest", "mist"], + "tag_revision": 1, + "replayed": False, + }, + ) + conn = self.db.model.get_connection() + files = conn.execute("SELECT * FROM media_annotation_files").fetchall() + self.assertEqual(len(files), 1) + self.assertRegex(files[0]["id"], UUID) + self.assertEqual(files[0]["locator"], loc) + byte = conn.execute("SELECT * FROM media_annotation_bytes").fetchone() + self.assertRegex(byte["id"], UUID) + self.assertEqual( + (byte["file_id"], byte["sha256"], byte["size"], byte["mtime_ns"]), + (files[0]["id"], DIGEST_A, SIZE, MTIME), + ) + self.assertEqual(byte["tag_revision"], 1) + self.assertEqual(byte["verified_at"], byte["first_seen_at"]) + self.assertEqual( + [ + tuple(row) + for row in conn.execute( + "SELECT t.name, mit.position FROM media_image_tags mit " + "JOIN tags t ON t.id = mit.tag_id ORDER BY mit.position" + ) + ], + [("Forest", 0), ("mist", 1)], + ) + self.assertEqual(self.index_revision(), 1) + # An explicit empty set is a real save: owner and byte record exist. + empty = self.replace(self.locator("b.png"), DIGEST_B, []) + self.assertEqual((empty["tags"], empty["tag_revision"]), ([], 1)) + self.assertEqual(self.counts()["media_annotation_files"], 2) + self.assertEqual(self.counts()["media_annotation_bytes"], 2) + self.assertEqual( + self.db.read_image_tags(self.locator("b.png"), DIGEST_B), + {"tags": [], "tag_revision": 1, "prior_count": 0}, + ) + self.assertEqual(self.index_revision(), 2) + self.schema_clean() + + def test_save_changes_only_target_byte_set(self): + a, b, c = self.locator("a.png"), self.locator("b.png"), self.locator("c.png") + prompt = self.db.save_prompt( + "Shared forest", + tags=["shared"], + prompt_hash=generate_prompt_hash("Shared forest"), + ) + pair = self.db.save_prompt_pair("Shared forest", "blurry", tags=["pairtag"]) + image_a = self.db.link_image_to_prompt(prompt, a) + image_b = self.db.link_image_to_prompt(prompt, b) + self.db.link_image_to_pair(pair, image_a, {"seed": 1}) + self.db.link_image_to_pair(pair, image_b, {"seed": 2}) + self.replace(b, DIGEST_A, ["sibling"]) + self.replace(c, DIGEST_A, ["copy"]) # equal bytes at another path + before = self.rows(*LEGACY_TABLES) + sibling_before = self.db.read_image_tags(b, DIGEST_A) + copy_before = self.db.read_image_tags(c, DIGEST_A) + + receipt = self.replace(a, DIGEST_A, ["only a"]) + + self.assertEqual((receipt["tags"], receipt["tag_revision"]), (["only a"], 1)) + after = self.rows(*LEGACY_TABLES) + # The registry gained the new name; nothing else moved. + self.assertEqual(after[1:], before[1:]) + self.assertEqual([row for row in after[0] if row[1] != "only a"], before[0]) + self.assertEqual(self.db.read_image_tags(b, DIGEST_A), sibling_before) + self.assertEqual(self.db.read_image_tags(c, DIGEST_A), copy_before) + self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["shared"]) + self.assertEqual(self.db.get_prompt_pair(pair)["tags"], ["pairtag"]) + + def test_changed_bytes_at_same_locator_start_empty_and_keep_prior(self): + loc = self.locator("a.png") + self.replace(loc, DIGEST_A, ["old"]) + self.assertEqual( + self.db.read_image_tags(loc, DIGEST_B), + {"tags": [], "tag_revision": 0, "prior_count": 1}, + ) + receipt = self.replace(loc, DIGEST_B, ["new"], mtime_ns=MTIME + 5) + self.assertEqual((receipt["tags"], receipt["tag_revision"]), (["new"], 1)) + self.assertEqual( + self.db.read_image_tags(loc, DIGEST_A), + {"tags": ["old"], "tag_revision": 1, "prior_count": 1}, + ) + self.assertEqual( + self.db.read_image_tags(loc, DIGEST_B), + {"tags": ["new"], "tag_revision": 1, "prior_count": 1}, + ) + self.assertEqual(self.counts()["media_annotation_files"], 1) + self.assertEqual(self.counts()["media_annotation_bytes"], 2) + self.assertEqual( + self.db.read_image_tags(self.locator("unseen.png"), DIGEST_A), + {"tags": [], "tag_revision": 0, "prior_count": 0}, + ) + + def test_revision_conflict_returns_current_set(self): + loc = self.locator("a.png") + self.replace(loc, tags=["a"]) + before = self.rows(*ANNOTATION_TABLES, "tags") + with self.assertRaises(AnnotationCommandError) as ctx: + self.replace(loc, tags=["b"], revision=0) + self.assertEqual( + (ctx.exception.code, ctx.exception.current["tag_revision"]), + ("tag-conflict", 1), + ) + self.assertIsInstance(ctx.exception, LibraryCommandError) + self.assertEqual(ctx.exception.epoch, self.epoch) + self.assertEqual(ctx.exception.current["tags"], ["a"]) + self.assertEqual(ctx.exception.current["sha256"], DIGEST_A) + with self.assertRaises(AnnotationCommandError) as ahead: + self.replace(loc, tags=["a"], revision=5) + self.assertEqual(ahead.exception.code, "tag-conflict") + self.assertEqual(self.rows(*ANNOTATION_TABLES, "tags"), before) + self.assertFalse(self.db.model.get_connection().in_transaction) + + def test_identical_replay_after_commit_is_success_replayed(self): + loc = self.locator("a.png") + first = self.replace(loc, tags=["a", "b"]) + self.assertEqual((first["tag_revision"], first["replayed"]), (1, False)) + before = self.rows(*ANNOTATION_TABLES) + # The lost-response retry: same set, the revision it was sent against. + retry = self.replace(loc, tags=["a", "b"], revision=0) + self.assertEqual(retry, {**first, "replayed": True}) + # Saving the current set at the current revision writes nothing either. + again = self.replace(loc, tags=["a", "b"], revision=1) + self.assertEqual(again, {**first, "replayed": True}) + self.assertEqual(self.rows(*ANNOTATION_TABLES), before) + # Order is part of the set. + reordered = self.replace(loc, tags=["b", "a"], revision=1) + self.assertEqual( + (reordered["tags"], reordered["tag_revision"], reordered["replayed"]), + (["b", "a"], 2, False), + ) + self.assertEqual(self.index_revision(), 2) + + def test_replay_with_new_fingerprint_moves_only_the_index_revision(self): + loc = self.locator("a.png") + first = self.replace(loc, tags=["a"]) + index = self.index_revision() + same = self.replace(loc, tags=["a"], revision=1) + self.assertEqual(same, {**first, "replayed": True}) + self.assertEqual(self.index_revision(), index) # identical: nothing moved + touched = self.replace(loc, tags=["a"], revision=1, mtime_ns=MTIME + 5) + self.assertEqual(touched, {**first, "replayed": True}) + self.assertEqual(self.index_revision(), index + 1) # currency changed + self.assertEqual( + self.db.current_image_tags({loc: (SIZE, MTIME + 5)}), {loc: ["a"]} + ) + self.assertEqual(self.db.current_image_tags({loc: (SIZE, MTIME)}), {}) + self.assertEqual(self.read(loc)["tag_revision"], 1) + + def test_epoch_checked_inside_transaction(self): + loc = self.locator("a.png") + seen = [] + original = PromptDatabase._library_epoch + + def spy(db, conn, expected_epoch): + seen.append(conn.in_transaction) + return original(db, conn, expected_epoch) + + with patch.object(PromptDatabase, "_library_epoch", spy): + self.replace(loc, tags=["a"]) + self.add(loc, tags=["b"]) + self.assertEqual(seen, [True, True]) + stale = "11111111-2222-4333-8444-555555555555" + for call in ( + lambda: self.replace( + self.locator("x.png"), tags=["a"], expected_epoch=stale + ), + lambda: self.add(self.locator("x.png"), tags=["a"], expected_epoch=stale), + ): + with self.assertRaises(LibraryCommandError) as ctx: + call() + self.assertEqual(ctx.exception.code, "storage-changed") + self.assertEqual(ctx.exception.epoch, self.epoch) + self.assertEqual(self.counts()["media_annotation_files"], 1) + for bad_epoch in (None, 7, "not-a-uuid"): + with self.assertRaises(ValueError): + self.replace(loc, tags=["a"], expected_epoch=bad_epoch) + + def test_limit_and_invalid_names_rejected_without_partial_write(self): + loc = self.locator("a.png") + before = self.rows(*ANNOTATION_TABLES, "tags") + for tags in ( + ["ok", "a,b"], + ["ok", ""], + ["ok", 3], + "ok", + [f"t{i}" for i in range(101)], + ): + with self.subTest(tags=tags), self.assertRaises(ValueError): + self.replace(loc, tags=tags) + with self.subTest(tags=tags, command="add"), self.assertRaises(ValueError): + self.add(loc, tags=tags) + for field, value in ( + ("locator", ""), + ("locator", None), + ("locator", "x" * 4097), + ("sha256", DIGEST_A.upper()), + ("sha256", DIGEST_A[:-1]), + ("sha256", None), + ("size", -1), + ("size", True), + ("size", "7"), + ("mtime_ns", 1.5), + ("mtime_ns", None), + ("expected_tag_revision", -1), + ("expected_tag_revision", False), + ): + with self.subTest(field=field, value=value), self.assertRaises(ValueError): + self.replace(loc, tags=["ok"], **{field: value}) + self.assertEqual(self.rows(*ANNOTATION_TABLES, "tags"), before) + self.assertEqual(self.index_revision(), 0) + + def test_writers_supply_uuid_ids_and_refuse_to_write_without_one(self): + loc = self.locator("a.png") + before = self.rows(*ANNOTATION_TABLES) + with patch.object(annotations_module, "uuid4", return_value=""): + with self.assertRaises(ValueError): + self.replace(loc, tags=["a"]) + self.assertEqual(self.rows(*ANNOTATION_TABLES), before) + self.replace(loc, tags=["a"]) + self.add(self.locator("b.png"), DIGEST_B, ["b"]) + conn = self.db.model.get_connection() + for table in ("media_annotation_files", "media_annotation_bytes"): + for (value,) in conn.execute(f"SELECT id FROM {table}"): + self.assertRegex(value, UUID) + self.schema_clean() + + def test_never_resolves_owner_through_prompt_lookup(self): + loc = self.locator("a.png") + prompt = self.db.save_prompt("Linked", tags=["prompt-tag"]) + self.db.link_image_to_prompt(prompt, loc) + with patch.object( + PromptDatabase, + "get_prompt_id_for_image", + side_effect=AssertionError("prompt lookup is not an owner"), + ): + self.replace(loc, tags=["img"]) + self.add(loc, tags=["more"]) + self.assertEqual( + self.db.read_image_tags(loc, DIGEST_A)["tags"], ["img", "more"] + ) + self.assertEqual( + self.db.current_image_tags({loc: (SIZE, MTIME)}), + {loc: ["img", "more"]}, + ) + self.assertTrue( + self.db.refresh_image_fingerprint(loc, DIGEST_A, SIZE, MTIME + 1) + ) + list(self.db.iter_image_tag_records()) + self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["prompt-tag"]) + + +class AddImageTagsTests(AnnotationTestCase): + def test_add_merges_without_removing_and_caps_at_100(self): + loc = self.locator("a.png") + self.replace(loc, tags=["Forest", "mist"]) + receipt = self.add(loc, tags=["forest", "Rain", "rain"]) + self.assertEqual( + receipt, + { + "library_epoch": self.epoch, + "sha256": DIGEST_A, + "tags": ["Forest", "mist", "Rain"], + "tag_revision": 2, + "replayed": False, + "added": 1, + }, + ) + nothing = self.add(loc, tags=["MIST"]) + self.assertEqual( + (nothing["tags"], nothing["tag_revision"], nothing["replayed"]), + (["Forest", "mist", "Rain"], 2, True), + ) + self.assertEqual(nothing["added"], 0) + before = self.rows(*ANNOTATION_TABLES, "tags") + with self.assertRaises(AnnotationCommandError) as ctx: + self.add(loc, tags=[f"t{i}" for i in range(98)]) + self.assertEqual(ctx.exception.code, "tag-limit") + self.assertEqual(ctx.exception.current["tag_revision"], 2) + self.assertEqual(self.rows(*ANNOTATION_TABLES, "tags"), before) + exact = self.add(loc, tags=[f"t{i}" for i in range(97)]) + self.assertEqual((len(exact["tags"]), exact["added"]), (100, 97)) + # An unseen file: add creates the owner and byte record like a save. + fresh = self.add(self.locator("b.png"), DIGEST_B, ["one"]) + self.assertEqual( + (fresh["tags"], fresh["tag_revision"], fresh["added"]), (["one"], 1, 1) + ) + self.schema_clean() + + +class CurrentImageTagsTests(AnnotationTestCase): + def test_current_tags_require_matching_fingerprint(self): + loc = self.locator("a.png") + other = self.locator("b.png") + self.replace(loc, tags=["a"]) + self.replace(other, DIGEST_B, [], mtime_ns=MTIME + 9) + self.assertEqual( + self.db.current_image_tags({loc: (SIZE, MTIME), other: (SIZE, MTIME + 9)}), + {loc: ["a"], other: []}, + ) + self.assertEqual(self.db.current_image_tags({loc: (SIZE, MTIME + 1)}), {}) + self.assertEqual(self.db.current_image_tags({loc: (SIZE + 1, MTIME)}), {}) + self.assertEqual(self.db.current_image_tags({}), {}) + self.assertEqual( + self.db.current_image_tags({self.locator("none.png"): (SIZE, MTIME)}), {} + ) + # Refreshing the verified fingerprint makes the record current again. + self.assertTrue( + self.db.refresh_image_fingerprint(loc, DIGEST_A, SIZE, MTIME + 1) + ) + self.assertFalse(self.db.refresh_image_fingerprint(loc, DIGEST_B, SIZE, MTIME)) + self.assertFalse( + self.db.refresh_image_fingerprint( + self.locator("x.png"), DIGEST_A, SIZE, MTIME + ) + ) + self.assertEqual( + self.db.current_image_tags({loc: (SIZE, MTIME + 1)}), {loc: ["a"]} + ) + self.assertEqual(self.db.current_image_tags({loc: (SIZE, MTIME)}), {}) + self.assertEqual( + self.db.read_image_tags(loc, DIGEST_A), + {"tags": ["a"], "tag_revision": 1, "prior_count": 0}, + ) + for bad in (None, [], {loc: (SIZE,)}, {loc: (SIZE, 1.5)}, {3: (SIZE, MTIME)}): + with self.subTest(bad=bad), self.assertRaises(ValueError): + self.db.current_image_tags(bad) + + def test_newest_verification_wins_on_equal_fingerprints(self): + loc = self.locator("a.png") + stamps = iter(f"2026-10-22T00:00:0{i}.000000+00:00" for i in range(10)) + with patch.object(annotations_module, "utc_now_iso", side_effect=stamps): + self.replace(loc, DIGEST_A, ["old bytes"]) + self.replace(loc, DIGEST_B, ["new bytes"]) # same size and mtime + self.assertEqual( + self.db.current_image_tags({loc: (SIZE, MTIME)}), {loc: ["new bytes"]} + ) + with patch.object( + annotations_module, + "utc_now_iso", + return_value="2026-10-22T00:01:00.000000+00:00", + ): + self.assertTrue( + self.db.refresh_image_fingerprint(loc, DIGEST_A, SIZE, MTIME) + ) + self.assertEqual( + self.db.current_image_tags({loc: (SIZE, MTIME)}), {loc: ["old bytes"]} + ) + + def test_lookup_is_chunked_and_records_stream_in_position_order(self): + locators = [self.locator(f"{i:04}.png") for i in range(503)] + self.replace(locators[0], tags=["z", "y"]) + self.replace(locators[502], DIGEST_B, ["x"], mtime_ns=MTIME + 2) + self.replace(locators[250], DIGEST_B, [], mtime_ns=MTIME + 1) + statements = [] + conn = self.db.model.get_connection() + conn.set_trace_callback(statements.append) + try: + current = self.db.current_image_tags( + { + loc: (SIZE, MTIME + (2 if i == 502 else 0)) + for i, loc in enumerate(locators) + } + ) + finally: + conn.set_trace_callback(None) + self.assertEqual(current, {locators[0]: ["z", "y"], locators[502]: ["x"]}) + lookups = [s for s in statements if "media_annotation_files" in s] + self.assertEqual(len(lookups), 2) + self.assertEqual( + list(self.db.iter_image_tag_records()), + [ + (locators[0], SIZE, MTIME, ["z", "y"]), + (locators[250], SIZE, MTIME + 1, []), + (locators[502], SIZE, MTIME + 2, ["x"]), + ], + ) + + +class AnnotationIsolationTests(AnnotationTestCase): + def test_cascade_from_vocabulary_delete_keeps_byte_records(self): + loc = self.locator("a.png") + self.replace(loc, tags=["gone", "kept"]) + self.db.delete_tag_all_prompts("gone") + self.assertEqual( + self.db.read_image_tags(loc, DIGEST_A), + { + "tags": ["kept"], + "tag_revision": 2, + "prior_count": 0, + }, # moved by the delete + ) + self.schema_clean() + + def test_junction_failure_leaves_nothing_behind(self): + loc = self.locator("a.png") + self.replace(loc, tags=["a"]) + conn = self.db.model.get_connection() + conn.execute( + "CREATE TRIGGER reject_image_tag BEFORE INSERT ON media_image_tags " + "BEGIN SELECT RAISE(ABORT,'image tag failure'); END" + ) + conn.commit() + before = self.rows(*ANNOTATION_TABLES, "tags") + with self.assertRaises(sqlite3.IntegrityError): + self.replace(loc, tags=["b"], revision=1) + with self.assertRaises(sqlite3.IntegrityError): + self.add(self.locator("new.png"), DIGEST_B, ["c"]) + self.assertFalse(conn.in_transaction) + self.assertEqual(self.rows(*ANNOTATION_TABLES, "tags"), before) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_api_vocabulary.py b/tests/test_api_vocabulary.py index a1d4a5a..062c77d 100644 --- a/tests/test_api_vocabulary.py +++ b/tests/test_api_vocabulary.py @@ -1,6 +1,8 @@ """Vocabulary read routes page with signed, epoch-scoped cursors.""" +import hashlib import json +import os import unittest from unittest.mock import patch @@ -184,6 +186,7 @@ class TestVocabularyAPI(ImageAPITestCase): "id": created["tag"]["id"], "name": "sunset", "prompt_count": 0, + "image_count": 0, }, }, ) @@ -202,6 +205,7 @@ class TestVocabularyAPI(ImageAPITestCase): "other_prompts": 0, "negative_variants": 1, "archived_negative_variants": 0, + "images": 0, }, }, ) @@ -221,9 +225,11 @@ class TestVocabularyAPI(ImageAPITestCase): "id": created["tag"]["id"], "name": "sunset", "prompt_count": 2, + "image_count": 0, }, "components_changed": 1, "pairs_changed": 1, + "images_changed": 0, }, ) self.assertEqual(self.api.db.get_prompt_by_id(parent)["tags"], ["sunset"]) @@ -242,6 +248,7 @@ class TestVocabularyAPI(ImageAPITestCase): "deleted": True, "components_changed": 1, "pairs_changed": 1, + "images_changed": 0, }, ) self.assertEqual(self.api.db.get_all_tags(), []) @@ -253,6 +260,95 @@ class TestVocabularyAPI(ImageAPITestCase): self.assertEqual(response.status, 400, bad) self.assertNotIn("code", result) + def image(self, name, tags, digest="a", mtime_ns=1): + """Tag the bytes ``digest`` of the file ``name``; return its locator.""" + locator = os.path.normcase(str(self.output_dir / name)) + self.api.db.replace_image_tags( + locator=locator, + sha256=hashlib.sha256(digest.encode()).hexdigest(), + size=1, + mtime_ns=mtime_ns, + expected_epoch=self.api.db.get_execution_epoch(), + expected_tag_revision=0, + tags=tags, + ) + return locator + + async def test_tag_rows_report_image_count_by_locator(self): + self._save_prompt("Prompt only", tags=["prompt-only", "shared"]) + self.image("one.png", ["twice", "shared"], digest="old") + self.image("one.png", ["twice"], digest="new", mtime_ns=2) + self.image("two.png", ["shared"], digest="new") + self.image("three.png", ["prior-only"], digest="first") + self.image("three.png", [], digest="second", mtime_ns=2) + response, page = await self.get("tags") + self.assertEqual(response.status, 200, page) + rows = {row["name"]: row for row in page["tags"]} + self.assertEqual( + { + name: (row["prompt_count"], row["image_count"]) + for name, row in rows.items() + }, + { + "prompt-only": (1, 0), + "shared": (1, 2), + "twice": (0, 1), + "prior-only": (0, 1), + }, + ) + row = rows["shared"] + self.assertEqual(row["image_count"], 2) + self.assertEqual(set(row), {"id", "name", "prompt_count", "image_count"}) + + async def test_impact_and_delete_counts_include_images(self): + epoch = self.api.db.get_execution_epoch() + self._save_prompt("Parent", tags=["moody"]) + self.image("a.png", ["moody", "keep"]) + self.image("b.png", ["moody"]) + _, page = await self.get("tags") + moody = next(row for row in page["tags"] if row["name"] == "moody") + self.assertEqual(moody["image_count"], 2) + response, impact = await self.get(f"tags/{moody['id']}/impact") + self.assertEqual(response.status, 200, impact) + self.assertEqual( + impact["impact"], + { + "positive_prompts": 1, + "other_prompts": 0, + "negative_variants": 0, + "archived_negative_variants": 0, + "images": 2, + }, + ) + response, deleted = await self.send( + "POST", + f"tags/{moody['id']}/delete", + {"expected_epoch": epoch, "name": "moody"}, + ) + self.assertEqual(response.status, 200, deleted) + self.assertEqual( + deleted, + { + "success": True, + "library_epoch": epoch, + "deleted": True, + "components_changed": 1, + "pairs_changed": 0, + "images_changed": 2, + }, + ) + keep = next( + row for row in (await self.get("tags"))[1]["tags"] if row["name"] == "keep" + ) + response, renamed = await self.send( + "PATCH", + f"tags/{keep['id']}", + {"expected_epoch": epoch, "name": "keep", "new_name": "kept"}, + ) + self.assertEqual(response.status, 200, renamed) + self.assertEqual(renamed["images_changed"], 1) + self.assertEqual(renamed["tag"]["image_count"], 1) + async def test_create_route_validates_names(self): epoch = self.api.db.get_execution_epoch() before = self.tag_rows() diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index dc7436a..deb1346 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -452,6 +452,44 @@ class TestRestoreFromFile(BackupTestCase): ).fetchone() ) + def test_restore_old_backup_without_annotations(self): + tables = ( + "media_image_tags", + "media_annotation_bytes", + "media_annotation_files", + "media_annotation_metadata", + ) + src = self._make_source_db(["predates image annotations"]) + with closing(sqlite3.connect(src)) as conn: + conn.execute("PRAGMA foreign_keys=OFF") + for table in tables: + conn.execute(f"DROP TABLE {table}") + conn.commit() + self.assertTrue(PromptModel.verify_database_file(src)[0]) + ok, reason = PromptModel.verify_database_file(src, require_library_state=True) + self.assertFalse(ok) + self.assertEqual(reason, "Missing image annotation schema") + + self.db.model.restore_from_file(src) + + self.assertEqual(_texts(self.path), ["predates image annotations"]) + with closing(sqlite3.connect(self.path)) as conn: + self.assertEqual( + conn.execute("SELECT * FROM media_annotation_metadata").fetchall(), + [(1, 1, 0)], + ) + ok, reason = PromptModel.verify_database_file( + self.path, require_library_state=True + ) + self.assertTrue(ok, reason) + with closing(sqlite3.connect(src)) as conn: # the source is untouched + self.assertEqual( + conn.execute( + "SELECT COUNT(*) FROM sqlite_master WHERE name LIKE 'media_%'" + ).fetchone()[0], + 0, + ) + def test_restore_preserves_favorites_and_closes_handles(self): self.db.set_image_favorite("output", "live.png") with tempfile.TemporaryDirectory() as directory: diff --git a/tests/test_database_migrations.py b/tests/test_database_migrations.py index 461e50b..544d0fb 100644 --- a/tests/test_database_migrations.py +++ b/tests/test_database_migrations.py @@ -200,6 +200,104 @@ class TestIdempotentMigrations(MigrationTestCase): self.assertEqual(self._columns("prompt_library_state"), ["prompt_id", "title"]) + def test_legacy_database_gains_annotation_group(self): + self._legacy_db([]) + + self._open() + + for table in ( + "media_annotation_metadata", + "media_annotation_files", + "media_annotation_bytes", + "media_image_tags", + ): + self.assertIn(table, self._table_names()) + conn = sqlite3.connect(self.path) + try: + self.assertEqual( + conn.execute("SELECT * FROM media_annotation_metadata").fetchall(), + [(1, 1, 0)], + ) + # No backfill: a legacy database starts with no image tags. + self.assertEqual( + conn.execute("SELECT COUNT(*) FROM media_image_tags").fetchone()[0], 0 + ) + finally: + conn.close() + + def test_rebuild_keeps_annotation_rows(self): + # The prompts rebuild runs with foreign keys off; image tags that + # reference the tags table must survive it untouched. + file_id = "0f8e4b2a-1c3d-4e5f-8a9b-0c1d2e3f4a5b" + byte_id = "6a1b2c3d-4e5f-4a6b-8c9d-0e1f2a3b4c5d" + self._legacy_db( + [ + "INSERT INTO prompts (text, hash, workflow_name)" + " VALUES ('a', 'h1', 'wf')", + "CREATE TABLE tags (id INTEGER PRIMARY KEY AUTOINCREMENT," + " name TEXT NOT NULL UNIQUE)", + "INSERT INTO tags (id, name) VALUES (3, 'img-only')", + "CREATE TABLE media_annotation_metadata (singleton_id INTEGER" + " PRIMARY KEY CHECK(singleton_id=1), schema_version INTEGER NOT" + " NULL CHECK(schema_version=1), image_index_revision INTEGER NOT" + " NULL DEFAULT 0 CHECK(image_index_revision>=0))", + "INSERT INTO media_annotation_metadata VALUES (1, 1, 4)", + "CREATE TABLE media_annotation_files (id TEXT PRIMARY KEY" + " CHECK(length(id)=36), locator TEXT NOT NULL UNIQUE," + " created_at TEXT NOT NULL)", + f"INSERT INTO media_annotation_files VALUES ('{file_id}'," + " '/out/a.png', '2026-10-22T00:00:00.000000+00:00')", + "CREATE TABLE media_annotation_bytes (id TEXT PRIMARY KEY" + " CHECK(length(id)=36), file_id TEXT NOT NULL REFERENCES" + " media_annotation_files(id) ON DELETE CASCADE, sha256 TEXT NOT" + " NULL CHECK(length(sha256)=64 AND sha256 NOT GLOB" + " '*[^0-9a-f]*'), size INTEGER NOT NULL CHECK(size>=0), mtime_ns" + " INTEGER NOT NULL, verified_at TEXT NOT NULL, first_seen_at TEXT" + " NOT NULL, tag_revision INTEGER NOT NULL DEFAULT 0" + " CHECK(tag_revision>=0), UNIQUE(file_id, sha256))", + "CREATE INDEX idx_media_annotation_bytes_fingerprint ON" + " media_annotation_bytes(file_id, size, mtime_ns)", + f"INSERT INTO media_annotation_bytes VALUES ('{byte_id}'," + f" '{file_id}', '{'a' * 64}', 3, 1700000000000000000," + " '2026-10-22T00:00:00.000000+00:00'," + " '2026-10-22T00:00:00.000000+00:00', 2)", + "CREATE TABLE media_image_tags (byte_id TEXT NOT NULL REFERENCES" + " media_annotation_bytes(id) ON DELETE CASCADE, tag_id INTEGER" + " NOT NULL REFERENCES tags(id) ON DELETE CASCADE, position" + " INTEGER NOT NULL CHECK(position>=0), PRIMARY KEY(byte_id," + " tag_id))", + "CREATE INDEX idx_media_image_tags_tag ON media_image_tags(tag_id)", + "CREATE UNIQUE INDEX idx_media_image_tags_position ON" + " media_image_tags(byte_id, position)", + f"INSERT INTO media_image_tags VALUES ('{byte_id}', 3, 0)", + ] + ) + + self._open() + + self.assertNotIn("workflow_name", self._columns("prompts")) + conn = sqlite3.connect(self.path) + try: + self.assertEqual( + conn.execute("SELECT * FROM media_annotation_metadata").fetchall(), + [(1, 1, 4)], + ) + self.assertEqual( + conn.execute( + "SELECT t.name, it.position FROM media_image_tags it" + " JOIN tags t ON t.id = it.tag_id" + ).fetchall(), + [("img-only", 0)], + ) + self.assertEqual( + conn.execute( + "SELECT tag_revision FROM media_annotation_bytes" + ).fetchall(), + [(2,)], + ) + finally: + conn.close() + def test_migration_error_is_logged_and_raised(self): # The legacy table has no CHECK on rating; the rebuilt one does, so the # copy into prompts_new fails inside the migration. diff --git a/tests/test_media_capture.py b/tests/test_media_capture.py new file mode 100644 index 0000000..e10a647 --- /dev/null +++ b/tests/test_media_capture.py @@ -0,0 +1,886 @@ +"""Bounded private capture store (#249 A3): exact bytes, hard limits, cleanup.""" + +import hashlib +import os +import stat +import struct +import tempfile +import threading +import time +import unittest +import zlib +from pathlib import Path +from unittest.mock import patch + +try: + from tests.open_handles import assert_closed, open_handles_under +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed, open_handles_under + +try: # guarded so a missing module fails assertions, not test loading + from utils import media_capture as capture_module + from utils.media_capture import ( + CAPTURE_EXTENSIONS, + CaptureCapacity, + CaptureChanged, + CaptureDenied, + CaptureStore, + CaptureTooLarge, + ) +except ImportError: + capture_module = None + CAPTURE_EXTENSIONS = None + CaptureCapacity = CaptureChanged = CaptureDenied = CaptureTooLarge = None + CaptureStore = None + +EPOCH = "epoch-1" +SCOPE = "a" * 64 +WAIT = 30.0 # generous ceiling for cross-thread handoffs, never a short clock + + +def _chunk(kind: bytes, data: bytes) -> bytes: + return ( + struct.pack(">I", len(data)) + + kind + + data + + struct.pack(">I", zlib.crc32(kind + data)) + ) + + +def png_bytes(width: int, height: int, pixels: bytes = b"\x00" * 16) -> bytes: + """A structurally valid PNG whose header declares *width* x *height*.""" + ihdr = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0) + return ( + b"\x89PNG\r\n\x1a\n" + + _chunk(b"IHDR", ihdr) + + _chunk(b"IDAT", pixels) + + _chunk(b"IEND", b"") + ) + + +class FakeClock: + def __init__(self, now: float = 1_000.0): + self.now = now + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class CaptureStoreTestCase(unittest.TestCase): + def setUp(self): + self.assertIsNotNone(capture_module, "utils.media_capture is not implemented") + parent = tempfile.TemporaryDirectory() + self.addCleanup(parent.cleanup) + self.addCleanup(assert_closed, self, parent.name) + # The resolved path: macOS /var -> /private/var, Windows 8.3 names. + self.parent = Path(os.path.realpath(parent.name)) + self.root = self.parent / "output" + self.root.mkdir() + self.clock = FakeClock() + + def make_store(self, **options) -> "CaptureStore": + options.setdefault("clock", self.clock) + store = CaptureStore(self.parent / "captures", **options) + self.addCleanup(store.close) + return store + + def write(self, relative: str, data: bytes) -> Path: + path = self.root / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return path + + def symlink(self, target, link): + try: + os.symlink(str(target), str(link)) + except (OSError, NotImplementedError) as error: + self.skipTest(f"symlinks unavailable here: {error}") + + def wait_for_change_time_after(self, ns: int) -> None: + """Block until a new file's ``st_ctime_ns`` passes *ns*. + + File timestamps tick coarsely (Linux jiffies, Windows ~1-16 ms). Each + probe is a fresh name: Windows reports the creation time (or, on + Python 3.12+, the change time) and tunnels the creation time of a + reused name for 15 s. + """ + deadline = time.monotonic() + WAIT + for index in range(1_000_000): + probe = self.parent / f"tick-{index}" + probe.write_bytes(b"x") + advanced = os.stat(probe).st_ctime_ns > ns + probe.unlink() + if advanced: + return + if time.monotonic() >= deadline: + break + time.sleep(0.001) + self.fail("the filesystem clock did not advance") # pragma: no cover + + def capture(self, store, relative="a.png", **options): + options.setdefault("epoch", EPOCH) + options.setdefault("scope", SCOPE) + return store.capture(self.root, relative, **options) + + def assert_empty(self, store): + self.assertFalse(any(Path(store.directory).iterdir())) + self.assertEqual( + store.stats(), {"captures": 0, "reserved_bytes": 0, "workers": 0} + ) + + # -- exact bytes and identity ------------------------------------------- + + def test_capture_hashes_and_copies_exact_bytes_owner_only(self): + source_bytes = png_bytes(3, 2, os.urandom(40)) + source = self.write("a.png", source_bytes) + store = self.make_store() + + cap = self.capture(store, root_index=1) + + self.assertEqual(Path(cap.path).read_bytes(), source_bytes) + self.assertEqual(cap.sha256, hashlib.sha256(source_bytes).hexdigest()) + self.assertEqual(cap.mtime_ns, os.stat(source).st_mtime_ns) + self.assertEqual(cap.size, len(source_bytes)) + self.assertEqual( + cap.locator, + os.path.normcase(os.path.join(os.path.realpath(self.root), "a.png")), + ) + self.assertEqual((cap.root_index, cap.scope, cap.epoch), (1, SCOPE, EPOCH)) + self.assertEqual((cap.media_type, cap.width, cap.height), ("image", 3, 2)) + self.assertTrue(cap.inferable) + self.assertIsNone(cap.inferable_reason) + self.assertIsInstance(cap.token, str) + self.assertGreaterEqual(len(cap.token), 32) + self.assertNotIn(cap.token, str(cap.path)) + self.assertNotIn(str(self.root), str(cap.path)) + capture_path = Path(cap.path) + self.assertEqual(capture_path.parent, Path(store.directory)) + self.assertEqual(capture_path.suffix, ".png") + self.assertTrue( + Path(store.directory).name.startswith("prompt-manager-captures-") + ) + if os.name == "posix": + self.assertEqual(stat.S_IMODE(os.stat(store.directory).st_mode) & 0o077, 0) + self.assertEqual(stat.S_IMODE(os.stat(cap.path).st_mode) & 0o077, 0) + self.assertIs(store.get(cap.token), cap) + self.assertEqual( + store.stats(), + {"captures": 1, "reserved_bytes": len(source_bytes), "workers": 0}, + ) + self.assertTrue(store.release(cap.token)) + self.assertFalse(store.release(cap.token)) + self.assertIsNone(store.get(cap.token)) + self.assert_empty(store) + assert_closed(self, self.parent) + + def test_streams_in_bounded_chunks_never_the_whole_file(self): + data = png_bytes(1, 1, os.urandom(5 * capture_module.CHUNK_SIZE + 7)) + self.write("big.png", data) + store = self.make_store(max_file=len(data)) + sizes = [] + original = CaptureStore._read + + def hook(store_, stream, size): + sizes.append(size) + return original(store_, stream, size) + + with patch.object(CaptureStore, "_read", hook): + cap = self.capture(store, "big.png") + self.assertTrue(sizes) + self.assertLessEqual(max(sizes), capture_module.CHUNK_SIZE) + self.assertEqual(cap.sha256, hashlib.sha256(data).hexdigest()) + + def test_symlinked_file_or_root_escape_is_denied(self): + outside = self.parent / "outside" + outside.mkdir() + (outside / "real.png").write_bytes(png_bytes(1, 1)) + self.write("a.png", png_bytes(1, 1)) + self.write("notes.txt", b"text") + (self.root / "folder").mkdir() + store = self.make_store() + + denied = [ + "../outside/real.png", + "a/../a.png", + "./a.png", + "a.png/", + "", + "/a.png", + "C:/a.png", + "a\\b.png", + "a\x00.png", + "notes.txt", + "folder", + "x" * 1021 + ".png", + ] + for relative in denied: + with self.subTest(relative=relative): + with self.assertRaises(CaptureDenied): + self.capture(store, relative) + with self.assertRaises(CaptureDenied): + self.capture(store, "a.png", root_index=True) + with self.assertRaises(CaptureChanged): + self.capture(store, "missing.png") + self.assert_empty(store) + + self.symlink(outside / "real.png", self.root / "link.png") + self.symlink(outside, self.root / "linked-dir") + with self.assertRaises(CaptureDenied): + self.capture(store, "link.png") + with self.assertRaises(CaptureDenied): + self.capture(store, "linked-dir/real.png") + self.assert_empty(store) + # A symlink to a file inside the root is still an alias, never a locator. + self.symlink(self.root / "a.png", self.root / "alias.png") + with self.assertRaises(CaptureDenied): + self.capture(store, "alias.png") + self.assert_empty(store) + + def test_locator_matches_listing_identity_for_symlinked_root(self): + real_root = self.parent / "real-output" + real_root.mkdir() + (real_root / "sub").mkdir() + source_bytes = png_bytes(2, 2) + (real_root / "sub" / "a.png").write_bytes(source_bytes) + link_root = self.parent / "link-output" + self.symlink(real_root, link_root) + store = self.make_store() + + cap = store.capture(link_root, "sub/a.png", epoch=EPOCH, scope=SCOPE) + + expected = os.path.normcase( + os.path.join(os.path.realpath(link_root), "sub", "a.png") + ) + self.assertEqual(cap.locator, expected) + self.assertEqual( + cap.locator, os.path.normcase(os.path.realpath(link_root / "sub/a.png")) + ) + self.assertTrue( + cap.locator.startswith(os.path.normcase(os.path.realpath(real_root))) + ) + self.assertEqual(cap.sha256, hashlib.sha256(source_bytes).hexdigest()) + + # -- change detection ---------------------------------------------------- + + def test_change_during_copy_is_detected_and_temp_removed(self): + source_bytes = png_bytes(1, 1, b"\x11" * 40) + source = self.write("a.png", source_bytes) + store = self.make_store(max_file=4096) + original = CaptureStore._read + + def appending(store_, stream, size): + data = original(store_, stream, 16) + if data and not getattr(appending, "done", False): + appending.done = True + with open(source, "ab") as handle: + handle.write(b"\x22" * 8) + return data + + with patch.object(CaptureStore, "_read", appending): + with self.assertRaises(CaptureChanged): + store.capture(self.root, "a.png", epoch=EPOCH, scope=SCOPE) + self.assert_empty(store) + + # Same size, different bytes, mtime restored: only the change time + # (POSIX ctime) can tell. Windows exposes no change time through stat, + # so there the torn copy completes and its digest describes the copy. + source.write_bytes(source_bytes) + before = os.stat(source) + self.wait_for_change_time_after(before.st_ctime_ns) + rewritten = bytes(b ^ 0xFF for b in source_bytes) + + def rewriting(store_, stream, size): + data = original(store_, stream, 16) + if data and not getattr(rewriting, "done", False): + rewriting.done = True + with open(source, "r+b") as handle: + handle.write(rewritten) + os.utime(source, ns=(before.st_atime_ns, before.st_mtime_ns)) + return data + + # Only a change time can tell: POSIX always has one; Windows exposes + # it from stat on Python 3.12+ (where the OS supports the fast path) + # and never from fstat, so detection there is best effort. The test + # asserts the outcome, not the platform: a detected change leaves no + # temp file; an undetected one (Windows only) completes with a digest + # that describes the copied bytes. + with patch.object(CaptureStore, "_read", rewriting): + try: + cap = store.capture(self.root, "a.png", epoch=EPOCH, scope=SCOPE) + except CaptureChanged: + detected = True + else: + detected = False + self.assertEqual(os.name, "nt", "only Windows may miss the rewrite") + copied = Path(cap.path).read_bytes() + self.assertEqual(cap.sha256, hashlib.sha256(copied).hexdigest()) + store.release(cap.token) + if os.name == "posix": + self.assertTrue(detected) + self.assert_empty(store) + self.assertEqual(os.stat(source).st_mtime_ns, before.st_mtime_ns) + + if os.name == "posix": + # The path now names another file (new inode) once the copy ends. + source.write_bytes(source_bytes) + + def replacing(store_, stream, size): + data = original(store_, stream, size) + if not data and not getattr(replacing, "done", False): + replacing.done = True + other = self.root / "other.png" + other.write_bytes(source_bytes) + os.replace(other, source) + return data + + with patch.object(CaptureStore, "_read", replacing): + with self.assertRaises(CaptureChanged): + store.capture(self.root, "a.png", epoch=EPOCH, scope=SCOPE) + self.assert_empty(store) + assert_closed(self, self.parent) + + def test_size_cap_on_stat_and_stream(self): + store = self.make_store(max_file=64) + self.write("large.png", bytes(65)) + self.write("fits.png", bytes(64)) + + def never(store_, stream, size): # pragma: no cover - guard only + raise AssertionError("an over-cap file must not be read") + + with patch.object(CaptureStore, "_read", never): + with self.assertRaises(CaptureTooLarge): + self.capture(store, "large.png") + self.assert_empty(store) + + cap = self.capture(store, "fits.png") + self.assertEqual(cap.size, 64) + store.release(cap.token) + + growing = self.write("growing.png", bytes(40)) + original = CaptureStore._read + + def grow(store_, stream, size): + data = original(store_, stream, 16) + if data and not getattr(grow, "done", False): + grow.done = True + with open(growing, "ab") as handle: + handle.write(bytes(40)) + return data + + with patch.object(CaptureStore, "_read", grow): + with self.assertRaises(CaptureTooLarge): + self.capture(store, "growing.png") + self.assert_empty(store) + + # -- admission ----------------------------------------------------------- + + def test_count_byte_and_worker_admission_is_nonblocking(self): + for name in ("a", "b", "c"): + self.write(f"{name}.png", bytes(60)) + store = self.make_store(max_captures=2, max_bytes=1000) + first = self.capture(store, "a.png") + self.capture(store, "b.png") + with self.assertRaises(CaptureCapacity): + self.capture(store, "c.png") + self.assertTrue(store.release(first.token)) + third = self.capture(store, "c.png") + self.assertEqual(store.stats()["captures"], 2) + store.invalidate_all() + self.assertIsNone(store.get(third.token)) + self.assert_empty(store) + + bytes_store = self.make_store(max_bytes=100) + kept = self.capture(bytes_store, "a.png") + with self.assertRaises(CaptureCapacity): + self.capture(bytes_store, "b.png") + self.assertEqual(bytes_store.stats()["reserved_bytes"], 60) + bytes_store.release(kept.token) + self.capture(bytes_store, "b.png") + + workers = self.make_store(max_workers=1) + entered, proceed = threading.Event(), threading.Event() + original = CaptureStore._read + results = {} + + def blocking(store_, stream, size): + if not entered.is_set(): + entered.set() + self.assertTrue(proceed.wait(WAIT)) + return original(store_, stream, size) + + def worker(): + try: + results["capture"] = self.capture(workers, "a.png") + except Exception as error: # pragma: no cover - reported below + results["error"] = error + + with patch.object(CaptureStore, "_read", blocking): + thread = threading.Thread(target=worker) + thread.start() + try: + self.assertTrue(entered.wait(WAIT)) + self.assertEqual(workers.stats()["workers"], 1) + with self.assertRaises(CaptureCapacity): + self.capture(workers, "b.png") + finally: + proceed.set() + thread.join(WAIT) + self.assertFalse(thread.is_alive()) + self.assertNotIn("error", results) + self.assertEqual(workers.stats()["workers"], 0) + self.assertEqual(results["capture"].size, 60) + self.capture(workers, "b.png") + + # -- lifetime ------------------------------------------------------------ + + def test_idle_and_hard_expiry_and_in_use_retention(self): + self.write("a.png", png_bytes(1, 1)) + self.write("b.png", png_bytes(1, 1)) + store = self.make_store(idle=900, hard=1800) + cap = self.capture(store, "a.png") + + self.clock.advance(899) + self.assertIs(store.get(cap.token), cap) # touches idle + self.assertAlmostEqual(store.expires_in(cap.token), 900) + self.clock.advance(899) + self.assertIs(store.get(cap.token), cap) + self.assertAlmostEqual(store.expires_in(cap.token), 2) + self.clock.advance(2) + self.assertIsNone(store.get(cap.token)) # hard limit despite touches + self.assertIsNone(store.expires_in(cap.token)) + self.assertFalse(Path(cap.path).exists()) + self.assert_empty(store) + + idle = self.capture(store, "a.png") + self.clock.advance(900) + store.sweep() + self.assertIsNone(store.get(idle.token)) + self.assert_empty(store) + + used = self.capture(store, "b.png") + self.assertIs(store.acquire(used.token), used) + self.assertIs(store.acquire(used.token), used) + self.clock.advance(901) + store.sweep() + self.assertIs(store.get(used.token), used) # in use is not idle + self.clock.advance(900) + store.sweep() + self.assertIsNone(store.get(used.token)) # expired for new requests + self.assertIsNone(store.acquire(used.token)) + self.assertTrue(Path(used.path).exists()) # retained while in use + self.assertEqual(store.stats()["captures"], 1) + store.release_use(used.token) + self.assertTrue(Path(used.path).exists()) + store.release_use(used.token) + self.assertFalse(Path(used.path).exists()) + self.assert_empty(store) + store.release_use(used.token) # unknown token: no error, nothing freed + self.assertIsNone(store.acquire("unknown")) + + def test_release_while_in_use_keeps_file_until_last_use_ends(self): + self.write("a.png", png_bytes(1, 1)) + store = self.make_store() + cap = self.capture(store, "a.png") + self.assertIs(store.acquire(cap.token), cap) + self.assertTrue(store.release(cap.token)) + self.assertIsNone(store.get(cap.token)) + self.assertTrue(Path(cap.path).exists()) + store.invalidate_all() + self.assertTrue(Path(cap.path).exists()) + store.release_use(cap.token) + self.assertFalse(Path(cap.path).exists()) + self.assert_empty(store) + + def test_invalidate_all_and_close_remove_owned_directory_only(self): + self.write("a.png", png_bytes(1, 1)) + self.write("b.png", png_bytes(1, 1)) + sibling = self.parent / "captures" / "sibling.txt" + sibling.parent.mkdir() + sibling.write_bytes(b"keep") + store = CaptureStore(self.parent / "captures", clock=self.clock) + try: + a = self.capture(store, "a.png") + b = self.capture(store, "b.png") + self.assertEqual(Path(store.directory).parent, sibling.parent) + store.invalidate_all() + self.assertIsNone(store.get(a.token)) + self.assertIsNone(store.get(b.token)) + self.assert_empty(store) + self.assertTrue(Path(store.directory).is_dir()) + c = self.capture(store, "a.png") + self.assertTrue(Path(c.path).exists()) + finally: + store.close() + self.assertFalse(Path(store.directory).exists()) + self.assertEqual(sibling.read_bytes(), b"keep") + self.assertTrue((self.root / "a.png").exists()) + store.close() # idempotent + with self.assertRaises(RuntimeError): + self.capture(store, "a.png") + self.assertIsNone(store.get(c.token)) + assert_closed(self, self.parent) + + def test_windows_handles_closed_before_unlink(self): + source_bytes = png_bytes(1, 1, b"\x33" * 40) + source = self.write("a.png", source_bytes) + store = self.make_store() + original_unlink = os.unlink + unlinked = [] + + def checked_unlink(path, *args, **kwargs): + self.assertEqual(open_handles_under(path), []) + self.assertTrue(store._lock.acquire(blocking=False)) # no I/O under lock + store._lock.release() + unlinked.append(os.fspath(path)) + return original_unlink(path, *args, **kwargs) + + with patch("os.unlink", checked_unlink): + cap = self.capture(store, "a.png") + self.assertEqual(open_handles_under(cap.path), []) + self.assertEqual(open_handles_under(source), []) + store.release(cap.token) + original = CaptureStore._read + + def appending(store_, stream, size): + data = original(store_, stream, 16) + if data and not getattr(appending, "done", False): + appending.done = True + with open(source, "ab") as handle: + handle.write(b"\x44") + return data + + with patch.object(CaptureStore, "_read", appending): + with self.assertRaises(CaptureChanged): + self.capture(store, "a.png") + self.assertEqual(open_handles_under(source), []) + self.assertEqual(len(unlinked), 2) + self.assertTrue(all(p.startswith(str(store.directory)) for p in unlinked)) + self.assert_empty(store) + assert_closed(self, self.parent) + + # -- inference readiness ------------------------------------------------- + + def test_dimension_limits_mark_generate_unsupported(self): + cases = { + "edge.png": (png_bytes(4096, 4096), True), + "wide.png": (png_bytes(4097, 1), False), + "tall.png": (png_bytes(1, 4097), False), + } + for name, (data, _) in cases.items(): + self.write(name, data) + self.write("garbage.jpg", b"\xff\xd8not really a jpeg") + store = self.make_store() + + for name, (data, supported) in cases.items(): + with self.subTest(name=name): + cap = self.capture(store, name) + self.assertEqual(cap.inferable, supported) + self.assertEqual((cap.width, cap.height) != (None, None), True) + if supported: + self.assertIsNone(cap.inferable_reason) + else: + self.assertIn("4096", cap.inferable_reason) + store.release(cap.token) + self.assertEqual( + (capture_module.MAX_INFERENCE_SIDE, capture_module.MAX_INFERENCE_PIXELS), + (4096, 16_777_216), + ) + cap = self.capture(store, "garbage.jpg") + self.assertEqual((cap.media_type, cap.width, cap.height), ("image", None, None)) + self.assertFalse(cap.inferable) + self.assertIn("header", cap.inferable_reason) + self.assertEqual( + cap.sha256, hashlib.sha256(b"\xff\xd8not really a jpeg").hexdigest() + ) + + def test_video_capture_is_not_inferable(self): + data = b"\x00\x00\x00\x18ftypmp42" + os.urandom(32) + self.write("clip.mp4", data) + store = self.make_store() + + def never(path): # pragma: no cover - guard only + raise AssertionError("video captures are never probed") + + with patch.object(capture_module, "read_inspection_file", never): + cap = self.capture(store, "clip.mp4") + self.assertEqual((cap.media_type, cap.width, cap.height), ("video", None, None)) + self.assertFalse(cap.inferable) + self.assertEqual( + cap.inferable_reason, + "Auto-tagging is available for supported still images only.", + ) + self.assertEqual(Path(cap.path).suffix, ".mp4") + self.assertEqual(cap.sha256, hashlib.sha256(data).hexdigest()) + + # -- verification without a copy ------------------------------------- + + def test_verify_hashes_current_file_without_retaining_a_copy(self): + source_bytes = png_bytes(1, 1, b"\x55" * 30) + source = self.write("a.png", source_bytes) + store = self.make_store(max_file=128, max_workers=1) + + facts = store.verify(self.root, "a.png") + + self.assertEqual(facts.sha256, hashlib.sha256(source_bytes).hexdigest()) + self.assertEqual(facts.size, len(source_bytes)) + self.assertEqual(facts.mtime_ns, os.stat(source).st_mtime_ns) + self.assertEqual( + facts.locator, + os.path.normcase(os.path.join(os.path.realpath(self.root), "a.png")), + ) + self.assert_empty(store) + self.write("large.png", bytes(129)) + with self.assertRaises(CaptureTooLarge): + store.verify(self.root, "large.png") + with self.assertRaises(CaptureDenied): + store.verify(self.root, "../a.png") + self.assert_empty(store) + + # -- arguments and failure paths --------------------------------------- + + def test_constructor_and_arguments_are_validated(self): + for options in ( + {"max_file": 0}, + {"max_captures": True}, + {"max_bytes": -1}, + {"max_workers": "2"}, + {"idle": 0}, + {"hard": -5}, + ): + with self.subTest(options=options): + with self.assertRaises(ValueError): + CaptureStore(self.parent / "captures", **options) + self.write("a.png", png_bytes(1, 1)) + store = self.make_store() + with self.assertRaises(ValueError): + store.capture(self.root, "a.png", epoch="", scope=SCOPE) + with self.assertRaises(ValueError): + store.capture(self.root, "a.png", epoch=EPOCH, scope=None) + with self.assertRaises(CaptureDenied): + self.capture(store, root_index=-1) + with self.assertRaises(CaptureDenied): + store.capture( + self.parent / "missing-root", "a.png", epoch=EPOCH, scope=SCOPE + ) + self.assert_empty(store) + + default = CaptureStore(clock=self.clock) + try: + self.assertEqual( + Path(default.directory).parent, Path(tempfile.gettempdir()) + ) + self.assertTrue(Path(default.directory).is_dir()) + finally: + default.close() + self.assertFalse(Path(default.directory).exists()) + + def test_directory_target_and_unreadable_file(self): + (self.root / "dir.png").mkdir() + store = self.make_store() + with self.assertRaises(CaptureDenied): + self.capture(store, "dir.png") + self.assert_empty(store) + if os.name != "posix" or os.geteuid() == 0: + return # read permission cannot be withdrawn here + locked = self.write("locked.png", png_bytes(1, 1)) + os.chmod(locked, 0) + try: + with self.assertRaises(CaptureChanged): + self.capture(store, "locked.png") + finally: + os.chmod(locked, 0o600) + self.assert_empty(store) + + def test_removal_failures_are_logged_and_never_touch_other_files(self): + self.write("a.png", png_bytes(1, 1)) + store = CaptureStore(self.parent / "captures", clock=self.clock) + stray = Path(store.directory) / "stray.bin" + try: + cap = self.capture(store, "a.png") + self.assertIs(store.acquire(cap.token), cap) + self.assertAlmostEqual(store.expires_in(cap.token), 1800) # in use + store.release_use(cap.token) + original_unlink = os.unlink + + def failing(path, *args, **kwargs): + if os.fspath(path) == cap.path: + raise PermissionError("busy") + return original_unlink(path, *args, **kwargs) + + with patch("os.unlink", failing): + with self.assertLogs("utils.media_capture", "WARNING") as logs: + self.assertTrue(store.release(cap.token)) + self.assertIn("could not be removed", logs.output[0]) + self.assertIsNone(store.get(cap.token)) + self.assertTrue(Path(cap.path).exists()) + os.unlink(cap.path) + stray.write_bytes(b"not ours") + finally: + with self.assertLogs("utils.media_capture", "WARNING") as logs: + store.close() + self.assertIn("could not be removed", logs.output[0]) + self.assertEqual(stray.read_bytes(), b"not ours") + stray.unlink() + os.rmdir(store.directory) + assert_closed(self, self.parent) + + # -- platform identity checks ------------------------------------------- + + def test_path_and_handle_stats_are_never_compared_across_apis(self): + """Windows (Python 3.12+) reports the change time from stat but the + creation time from fstat, and may encode dev/ino differently; only + size and mtime_ns may be compared across the two APIs.""" + source_bytes = png_bytes(2, 2, b"\x66" * 40) + self.write("a.png", source_bytes) + store = self.make_store() + original_fstat = os.fstat + + def shifted(fd): + st = original_fstat(fd) + return os.stat_result( + ( + st.st_mode, + st.st_ino + 7, + st.st_dev + 7, + st.st_nlink, + st.st_uid, + st.st_gid, + st.st_size, + st.st_atime, + st.st_mtime, + st.st_ctime + 1, + ), + { + "st_atime_ns": st.st_atime_ns, + "st_mtime_ns": st.st_mtime_ns, + "st_ctime_ns": st.st_ctime_ns + 1_000_000_000, + }, + ) + + with patch.object(capture_module.os, "fstat", shifted): + cap = self.capture(store, "a.png") + facts = store.verify(self.root, "a.png") + self.assertEqual(cap.sha256, hashlib.sha256(source_bytes).hexdigest()) + self.assertEqual(facts.sha256, cap.sha256) + store.release(cap.token) + self.assert_empty(store) + + def test_change_time_moving_between_fstat_reads_is_detected(self): + """The Windows 3.13 / POSIX signal, forced on every platform: the + descriptor's change time differs after the copy, size and mtime equal.""" + self.write("a.png", png_bytes(1, 1, b"\x77" * 40)) + store = self.make_store() + original_fstat = os.fstat + calls = [] + + def drifting(fd): + st = original_fstat(fd) + calls.append(fd) + if len(calls) < 2: + return st + return os.stat_result( + tuple(st)[:10], + { + "st_atime_ns": st.st_atime_ns, + "st_mtime_ns": st.st_mtime_ns, + "st_ctime_ns": st.st_ctime_ns + 1, + }, + ) + + with patch.object(capture_module.os, "fstat", drifting): + with self.assertRaises(CaptureChanged): + self.capture(store, "a.png") + self.assertEqual(len(calls), 2) + self.assert_empty(store) + + def test_fresh_file_captures_twice_on_every_platform(self): + data = png_bytes(1, 1, os.urandom(3 * capture_module.CHUNK_SIZE)) + self.write("fresh.png", data) + store = self.make_store() + first = self.capture(store, "fresh.png") + second = self.capture(store, "fresh.png") + self.assertEqual(first.sha256, hashlib.sha256(data).hexdigest()) + self.assertEqual(second.sha256, first.sha256) + self.assertNotEqual(first.token, second.token) + self.assertEqual(store.stats()["captures"], 2) + + def test_error_messages_never_name_server_paths(self): + self.write("a.png", png_bytes(1, 1)) + self.write("dir.png/x", b"") # dir.png is a directory + store = self.make_store(max_file=64) + attempts = [ + lambda: self.capture(store, "missing.png"), + lambda: self.capture(store, "a.png"), # over the 64-byte cap + lambda: self.capture(store, "dir.png"), + lambda: self.capture(store, "../a.png"), + lambda: self.capture(store, "notes.txt"), + lambda: store.verify(self.root, "missing.png"), + ] + if os.name == "posix" and os.geteuid() != 0: + locked = self.write("locked.png", bytes(10)) + os.chmod(locked, 0) + self.addCleanup(os.chmod, locked, 0o600) + attempts.append(lambda: self.capture(store, "locked.png")) + for attempt in attempts: + with self.assertRaises(capture_module.CaptureError) as caught: + attempt() + message = str(caught.exception) + self.assertTrue(message) + self.assertNotIn(str(self.root), message) + self.assertNotIn(str(self.parent), message) + self.assertNotIn(os.sep, message) + self.assert_empty(store) + + def test_close_during_a_capture_removes_the_directory_once_it_finishes(self): + self.write("a.png", png_bytes(1, 1, bytes(64))) + store = CaptureStore(self.parent / "captures", clock=self.clock) + original = CaptureStore._read + + def closing(store_, stream, size): + data = original(store_, stream, 16) + if data and not getattr(closing, "done", False): + closing.done = True + store.close() + return data + + with patch.object(CaptureStore, "_read", closing): + with self.assertRaises(RuntimeError): + self.capture(store, "a.png") + self.assertFalse(Path(store.directory).exists()) + self.assertEqual(store.stats()["reserved_bytes"], 0) + store.close() + assert_closed(self, self.parent) + + # -- constants --------------------------------------------------------- + + def test_defaults_match_the_plan(self): + store = self.make_store() + self.assertEqual( + ( + store.max_file, + store.max_captures, + store.max_bytes, + store.max_workers, + store.idle, + store.hard, + ), + (256 * 2**20, 8, 512 * 2**20, 2, 900, 1800), + ) + from py.api.inspection_service import SERVABLE_EXTENSIONS + + self.assertEqual(CAPTURE_EXTENSIONS, SERVABLE_EXTENSIONS) + + def test_exceptions_carry_codes(self): + self.assertEqual(CaptureCapacity.code, "review-capacity") + self.assertEqual(CaptureChanged.code, "file-changed") + self.assertEqual(CaptureTooLarge.code, "unsupported-media") + self.assertEqual(CaptureDenied.code, "source-changed") + for cls in (CaptureCapacity, CaptureChanged, CaptureTooLarge, CaptureDenied): + self.assertTrue(issubclass(cls, capture_module.CaptureError)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pair_schema_restore.py b/tests/test_pair_schema_restore.py index 4fb22f3..dff5420 100644 --- a/tests/test_pair_schema_restore.py +++ b/tests/test_pair_schema_restore.py @@ -251,7 +251,19 @@ class PairSchemaTests(unittest.TestCase): [r[1] for r in conn.execute("PRAGMA table_info(prompt_library_state)")], ["prompt_id", "title"], ) + self.assertEqual( + [ + tuple(r) + for r in conn.execute("SELECT * FROM media_annotation_metadata") + ], + [(1, 1, 0)], + ) self.assertTrue(PromptModel.verify_database_file(live.db_path)[0]) + self.assertTrue( + PromptModel.verify_database_file(live.db_path, require_library_state=True)[ + 0 + ] + ) def test_restore_staged_verification_failure_does_not_replace_live_pages(self): live = self.model() diff --git a/tests/test_pair_tag_maintenance.py b/tests/test_pair_tag_maintenance.py index 4a65691..6402013 100644 --- a/tests/test_pair_tag_maintenance.py +++ b/tests/test_pair_tag_maintenance.py @@ -1,5 +1,6 @@ """Global tag maintenance preserves pair and standalone tag relationships.""" +import hashlib import os import sqlite3 import tempfile @@ -10,6 +11,7 @@ try: except ImportError: # discovered with tests/ as the top-level directory from open_handles import assert_closed +from database.library_commands import LibraryCommandError from database.operations import PromptDatabase @@ -18,6 +20,7 @@ class TestPairTagMaintenance(unittest.TestCase): directory = tempfile.TemporaryDirectory() self.addCleanup(directory.cleanup) self.addCleanup(assert_closed, self, directory.name) + self.root = directory.name self.db = PromptDatabase(os.path.join(directory.name, "pair-tags.db")) self.addCleanup(self.db.close_all) @@ -27,6 +30,36 @@ class TestPairTagMaintenance(unittest.TestCase): def tags(self, pair_id: int) -> list: return self.db.get_prompt_pair(pair_id)["tags"] + def image(self, name: str, tags: list, digest: str = "a", revision: int = 0): + """Tag the bytes ``digest`` of the file ``name``; return its locator.""" + locator = os.path.normcase(os.path.join(os.path.realpath(self.root), name)) + replace = getattr(self.db, "replace_image_tags", None) + self.assertTrue(callable(replace), "replace_image_tags is missing") + replace( + locator=locator, + sha256=hashlib.sha256(digest.encode()).hexdigest(), + size=1, + mtime_ns=1, + expected_epoch=self.db.get_execution_epoch(), + expected_tag_revision=revision, + tags=tags, + ) + return locator + + def image_tags(self, locator: str, digest: str = "a") -> list: + return self.db.read_image_tags( + locator, hashlib.sha256(digest.encode()).hexdigest() + )["tags"] + + def image_rows(self) -> list: + conn = self.db.model.get_connection() + return [ + tuple(row) + for row in conn.execute( + "SELECT byte_id, tag_id, position FROM media_image_tags ORDER BY 1,3" + ) + ] + def test_rename_to_existing_tag_preserves_pair_links_and_deduplicates(self): first = self.pair("portrait", ["old", "new", "keep"]) second = self.pair("landscape", ["old"]) @@ -136,6 +169,207 @@ class TestPairTagMaintenance(unittest.TestCase): self.assertEqual(self.tags(other), ["source"]) self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["source"]) + def test_rename_and_merge_move_image_links_with_dedup(self): + pair = self.pair("portrait", ["old", "new"]) + both = self.image("both.png", ["new", "old", "keep"]) + only_old = self.image("old.png", ["old"]) + result = self.db.rename_tag_all_prompts("old", "new") + self.assertEqual(result["affected_count"], 0) # prompts only, as before + self.assertEqual(self.tags(pair), ["new"]) + self.assertEqual(self.image_tags(both), ["new", "keep"]) + self.assertEqual(self.image_tags(only_old), ["new"]) + self.assertNotIn("old", self.db.get_all_tags()) + self.db.rename_tag_all_prompts("keep", "kept") + self.assertEqual(self.image_tags(both), ["new", "kept"]) + merged = self.db.merge_tags(["kept", "new"], "final") + self.assertEqual(merged["tags_merged"], 2) + self.assertEqual(merged["affected_count"], 0) + self.assertEqual(self.image_tags(both), ["final"]) + self.assertEqual(self.image_tags(only_old), ["final"]) + self.assertEqual(self.tags(pair), ["final"]) + epoch = self.db.get_execution_epoch() + conn = self.db.model.get_connection() + final = conn.execute("SELECT id FROM tags WHERE name='final'").fetchone()[0] + native = self.db.rename_vocabulary_tag( + final, expected_epoch=epoch, name="final", new_name="native" + ) + self.assertEqual(native["images_changed"], 2) + self.assertEqual(self.image_tags(both), ["native"]) + self.assertEqual(self.image_tags(only_old), ["native"]) + self.image("again.png", ["native", "merge-me"]) + source = conn.execute("SELECT id FROM tags WHERE name='merge-me'").fetchone() + receipt = self.db.rename_vocabulary_tag( + source[0], expected_epoch=epoch, name="merge-me", new_name="native" + ) + self.assertTrue(receipt["merged"]) + self.assertEqual(receipt["images_changed"], 1) + self.assertEqual(receipt["tag"]["image_count"], 3) + self.assertEqual(len(self.image_rows()), 3) + + def test_delete_cascades_image_links_and_keeps_legacy_counts(self): + pair = self.pair("portrait", ["remove", "keep"]) + prompt = self.db.save_prompt("standalone", tags=["remove", "keep"]) + locator = self.image("a.png", ["keep", "remove"]) + result = self.db.delete_tag_all_prompts("remove") + self.assertEqual(result["affected_count"], 1) + self.assertEqual(self.tags(pair), ["keep"]) + self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["keep"]) + self.assertEqual(self.image_tags(locator), ["keep"]) + self.assertEqual(len(self.image_rows()), 1) + epoch = self.db.get_execution_epoch() + conn = self.db.model.get_connection() + keep = conn.execute("SELECT id FROM tags WHERE name='keep'").fetchone()[0] + impact = self.db.vocabulary_tag_impact(keep)["impact"] + self.assertEqual(impact["images"], 1) + deleted = self.db.delete_vocabulary_tag(keep, expected_epoch=epoch, name="keep") + self.assertEqual( + ( + deleted["components_changed"], + deleted["pairs_changed"], + deleted["images_changed"], + ), + (1, 1, 1), + ) + self.assertEqual(self.image_tags(locator), []) + self.assertEqual(self.image_rows(), []) + self.assertEqual( + conn.execute("SELECT COUNT(*) FROM media_annotation_bytes").fetchone()[0], 1 + ) + + def index_revision(self) -> int: + return ( + self.db.model.get_connection() + .execute("SELECT image_index_revision FROM media_annotation_metadata") + .fetchone()[0] + ) + + def revision(self, locator: str, digest: str = "a") -> int: + return self.db.read_image_tags( + locator, hashlib.sha256(digest.encode()).hexdigest() + )["tag_revision"] + + def tag_id(self, name: str) -> int: + conn = self.db.model.get_connection() + return conn.execute("SELECT id FROM tags WHERE name=?", (name,)).fetchone()[0] + + def test_vocabulary_maintenance_bumps_image_revisions_so_stale_saves_conflict( + self, + ): + forest = self.image("a.png", ["forest"]) + misty = self.image("b.png", ["mist"]) + self.pair("portrait", ["unlinked"]) + index = self.index_revision() + epoch = self.db.get_execution_epoch() + + # A rename of a tag no image carries moves nothing image-side. + self.db.rename_vocabulary_tag( + self.tag_id("unlinked"), expected_epoch=epoch, name="unlinked", new_name="u" + ) + self.assertEqual(self.index_revision(), index) + + self.db.rename_vocabulary_tag( + self.tag_id("forest"), expected_epoch=epoch, name="forest", new_name="woods" + ) + self.assertEqual(self.image_tags(forest), ["woods"]) + self.assertEqual((self.revision(forest), self.revision(misty)), (2, 1)) + self.assertEqual(self.index_revision(), index + 1) + # The Tag dialog still open at revision 1 cannot recreate "forest". + with self.assertRaises(LibraryCommandError) as caught: + self.db.replace_image_tags( + locator=forest, + sha256=hashlib.sha256(b"a").hexdigest(), + size=1, + mtime_ns=1, + expected_epoch=epoch, + expected_tag_revision=1, + tags=["forest", "mist"], + ) + self.assertEqual(caught.exception.code, "tag-conflict") + self.assertEqual(caught.exception.current["tags"], ["woods"]) + self.assertEqual(caught.exception.current["tag_revision"], 2) + self.assertNotIn("forest", self.db.get_all_tags()) + + self.db.rename_tag_all_prompts("woods", "trees") # legacy rename + self.assertEqual((self.revision(forest), self.revision(misty)), (3, 1)) + self.db.merge_tags(["trees"], "mist") # legacy merge: only the source side + self.assertEqual(self.image_tags(forest), ["mist"]) + self.assertEqual((self.revision(forest), self.revision(misty)), (4, 1)) + self.db.rename_tag_all_prompts("mist", "fog") # both carry it now + self.assertEqual((self.revision(forest), self.revision(misty)), (5, 2)) + self.db.delete_vocabulary_tag( + self.tag_id("fog"), expected_epoch=epoch, name="fog" + ) + self.assertEqual((self.revision(forest), self.revision(misty)), (6, 3)) + self.assertEqual((self.image_tags(forest), self.image_tags(misty)), ([], [])) + self.image("c.png", ["gone"]) + self.db.delete_tag_all_prompts("gone") # legacy delete + self.assertEqual(self.revision(forest), 6) + self.assertEqual( + self.index_revision(), index + 7 + ) # 6 maintenance ops + the c.png save + + def test_image_only_and_prior_only_names_participate(self): + only = self.image("only.png", ["image-only"]) + prior = self.image("prior.png", ["prior-only"], digest="old") + self.image("prior.png", [], digest="new") # the file's bytes changed + self.db.save_prompt("standalone", tags=["target"]) + merged = self.db.merge_tags(["image-only"], "target") + self.assertEqual((merged["tags_merged"], merged["affected_count"]), (1, 0)) + self.assertEqual(self.image_tags(only), ["target"]) + self.assertEqual( + self.db.rename_tag_all_prompts("prior-only", "renamed")["affected_count"], 0 + ) + self.assertEqual(self.image_tags(prior, "old"), ["renamed"]) + self.assertEqual(self.image_tags(prior, "new"), []) + self.assertEqual(set(self.db.get_all_tags()), {"target", "renamed"}) + epoch = self.db.get_execution_epoch() + conn = self.db.model.get_connection() + renamed = conn.execute("SELECT id FROM tags WHERE name='renamed'").fetchone()[0] + receipt = self.db.rename_vocabulary_tag( + renamed, expected_epoch=epoch, name="renamed", new_name="target" + ) + self.assertTrue(receipt["merged"]) + self.assertEqual(receipt["images_changed"], 1) + self.assertEqual(self.image_tags(prior, "old"), ["target"]) + self.assertEqual(receipt["tag"]["image_count"], 2) + + def test_image_junction_failure_rolls_back_prompt_pair_and_name_changes(self): + pair = self.pair("portrait", ["source", "target"]) + prompt = self.db.save_prompt("standalone", tags=["source"]) + locator = self.image("a.png", ["source"]) + conn = self.db.model.get_connection() + tables = ("tags", "prompt_tags", "prompt_pair_tags", "media_image_tags") + before = [ + conn.execute(f"SELECT * FROM {table} ORDER BY 1,2").fetchall() + for table in tables + ] + source = conn.execute("SELECT id FROM tags WHERE name='source'").fetchone()[0] + conn.execute( + "CREATE TRIGGER reject_image_tag_change BEFORE UPDATE " + "ON media_image_tags BEGIN SELECT RAISE(ABORT,'image tag failure'); END" + ) + conn.commit() + with self.assertRaises(sqlite3.IntegrityError): + self.db.rename_vocabulary_tag( + source, + expected_epoch=self.db.get_execution_epoch(), + name="source", + new_name="target", + ) + with self.assertRaises(sqlite3.IntegrityError): + self.db.merge_tags(["source"], "target") + with self.assertRaises(sqlite3.IntegrityError): + self.db.rename_tag_all_prompts("source", "target") + self.assertFalse(conn.in_transaction) + after = [ + conn.execute(f"SELECT * FROM {table} ORDER BY 1,2").fetchall() + for table in tables + ] + self.assertEqual(after, before) + self.assertEqual(self.tags(pair), ["source", "target"]) + self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["source"]) + self.assertEqual(self.image_tags(locator), ["source"]) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_vocabulary.py b/tests/test_vocabulary.py index 8ab623c..d8252eb 100644 --- a/tests/test_vocabulary.py +++ b/tests/test_vocabulary.py @@ -1,5 +1,7 @@ """Bounded vocabulary reads count exactly what Library Browse shows.""" +import hashlib +import os import re import sqlite3 import tempfile @@ -61,7 +63,48 @@ class TestVocabularyReads(unittest.TestCase): by_name = {row["name"]: row for row in page["tags"]} self.assertEqual(by_name["z"]["prompt_count"], 0) self.assertEqual(by_name["a"]["prompt_count"], 2) - self.assertEqual(set(by_name["c"]), {"id", "name", "prompt_count"}) + self.assertEqual( + set(by_name["c"]), {"id", "name", "prompt_count", "image_count"} + ) + self.assertEqual({row["image_count"] for row in page["tags"]}, {0}) + + def image(self, name, tags, digest="a", mtime_ns=1): + """Tag the bytes ``digest`` of the file ``name``; return its locator.""" + locator = os.path.normcase(os.path.join(os.path.realpath(self.temp.name), name)) + self.read( + "replace_image_tags", + locator=locator, + sha256=hashlib.sha256(digest.encode()).hexdigest(), + size=1, + mtime_ns=mtime_ns, + expected_epoch=self.db.get_execution_epoch(), + expected_tag_revision=0, + tags=tags, + ) + return locator + + def test_tag_rows_report_image_count_by_locator(self): + self.save("Prompt only", tags=["prompt-only"]) + self.image("one.png", ["twice", "shared"], digest="old") + self.image( + "one.png", ["twice"], digest="new", mtime_ns=2 + ) # two bytes, one file + self.image("two.png", ["shared"], digest="new") # equal digest, other locator + self.image("three.png", ["prior-only"], digest="first") + self.image("three.png", [], digest="second", mtime_ns=2) + rows = {row["name"]: row for row in self.all_tags()} + self.assertEqual(rows["twice"]["image_count"], 1) + row = rows["shared"] + self.assertEqual(row["image_count"], 2) + self.assertEqual(rows["prior-only"]["image_count"], 1) + self.assertEqual(rows["prompt-only"]["image_count"], 0) + self.assertEqual( + {name: r["prompt_count"] for name, r in rows.items()}, + {"prompt-only": 1, "twice": 0, "shared": 0, "prior-only": 0}, + ) + self.assertEqual( + self.db.list_library(tags=["shared"])["total"], row["prompt_count"] + ) def test_prompt_count_equals_library_browse_total_for_every_tag(self): self.save("Parent only", tags=["parent-only"]) @@ -302,7 +345,12 @@ class TestVocabularyTagCommands(unittest.TestCase): self.assertTrue(created["created"]) self.assertEqual( created["tag"], - {"id": self.tag_id("sunset"), "name": "sunset", "prompt_count": 0}, + { + "id": self.tag_id("sunset"), + "name": "sunset", + "prompt_count": 0, + "image_count": 0, + }, ) after = self.junction_rows() self.assertEqual(after[1:], before[1:]) @@ -345,9 +393,15 @@ class TestVocabularyTagCommands(unittest.TestCase): { "library_epoch": self.epoch, "merged": False, - "tag": {"id": tag_id, "name": "dramatic", "prompt_count": 2}, + "tag": { + "id": tag_id, + "name": "dramatic", + "prompt_count": 2, + "image_count": 0, + }, "components_changed": 2, "pairs_changed": 2, + "images_changed": 0, }, ) self.assertIsNone(self.tag_id("moody")) @@ -397,7 +451,8 @@ class TestVocabularyTagCommands(unittest.TestCase): ) self.assertFalse(result["merged"]) self.assertEqual( - result["tag"], {"id": tag_id, "name": "cat", "prompt_count": 1} + result["tag"], + {"id": tag_id, "name": "cat", "prompt_count": 1, "image_count": 0}, ) self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["cat"]) @@ -467,6 +522,7 @@ class TestVocabularyTagCommands(unittest.TestCase): "other_prompts": 1, "negative_variants": 1, "archived_negative_variants": 1, + "images": 0, }, }, ) @@ -482,6 +538,7 @@ class TestVocabularyTagCommands(unittest.TestCase): "deleted": True, "components_changed": 2, "pairs_changed": 2, + "images_changed": 0, }, ) self.assertIsNone(self.tag_id("moody")) @@ -494,6 +551,69 @@ class TestVocabularyTagCommands(unittest.TestCase): ) self.assertEqual(caught.exception.code, "vocabulary-conflict") + def image(self, name, tags, digest="a", mtime_ns=1): + locator = os.path.normcase(os.path.join(os.path.realpath(self.temp.name), name)) + self.command( + "replace_image_tags", + locator=locator, + sha256=hashlib.sha256(digest.encode()).hexdigest(), + size=1, + mtime_ns=mtime_ns, + expected_epoch=self.epoch, + expected_tag_revision=0, + tags=tags, + ) + return locator + + def test_impact_and_delete_counts_include_images(self): + self.owners() + self.image("a.png", ["moody", "keep"]) + self.image("a.png", ["moody"], digest="prior", mtime_ns=2) + self.image("b.png", ["moody"]) + tag_id = self.tag_id("moody") + impact = self.command("vocabulary_tag_impact", tag_id)["impact"] + self.assertEqual(impact["images"], 2) + self.assertEqual( + self.command("vocabulary_tag_impact", self.tag_id("keep"))["impact"], + { + "positive_prompts": 0, + "other_prompts": 0, + "negative_variants": 0, + "archived_negative_variants": 0, + "images": 1, + }, + ) + result = self.command( + "delete_vocabulary_tag", tag_id, expected_epoch=self.epoch, name="moody" + ) + self.assertEqual( + ( + result["components_changed"], + result["pairs_changed"], + result["images_changed"], + ), + (2, 2, 2), + ) + with self.db.model.get_connection() as conn: + self.assertEqual( + conn.execute("SELECT COUNT(*) FROM media_image_tags").fetchone()[0], 1 + ) + self.assertEqual( + conn.execute("SELECT COUNT(*) FROM media_annotation_bytes").fetchone()[ + 0 + ], + 3, + ) + renamed = self.command( + "rename_vocabulary_tag", + self.tag_id("keep"), + expected_epoch=self.epoch, + name="keep", + new_name="kept", + ) + self.assertEqual(renamed["images_changed"], 1) + self.assertEqual(renamed["tag"]["image_count"], 1) + _OLD_STAMP = "2000-01-01T00:00:00+00:00" _DML = re.compile(r"\s*(INSERT|UPDATE|DELETE|REPLACE)\b", re.IGNORECASE) diff --git a/utils/media_capture.py b/utils/media_capture.py new file mode 100644 index 0000000..86988db --- /dev/null +++ b/utils/media_capture.py @@ -0,0 +1,632 @@ +"""Bounded private capture store for image annotations (#249 A3). + +One explicit Tag or review action captures exactly one selected media file: +the rooted file is streamed to a private, owner-only temporary copy while its +whole-file SHA-256 is computed, and the identity of the opened descriptor is +checked against the path before and after the copy. Preview, inference and the +metadata read use only that copy, so what the user reviewed is what the model +saw and what Save re-checks. Everything is bounded: bytes per file, completed +captures, reserved bytes, concurrent capture/verify workers, idle and hard +lifetimes. + +The contract is "these bytes at this location, last checked at this time". +Stat equality is never content proof: a same-size rewrite that restores +``mtime_ns`` is visible only through the change time between the two +``fstat`` reads, which Windows before Python 3.13 reports as the creation time +(no signal). +""" + +from __future__ import annotations + +import hashlib +import logging +import os +import secrets +import stat as stat_module +import tempfile +import threading +import time +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Callable + +from PIL import Image + +try: + from .prompt_inspector import MetadataReadError, read_inspection_file + from .video_metadata import VIDEO_EXTENSIONS, is_video_path +except ImportError: # imported with the repo root on sys.path + from utils.prompt_inspector import MetadataReadError, read_inspection_file + from utils.video_metadata import VIDEO_EXTENSIONS, is_video_path + +logger = logging.getLogger(__name__) + +CHUNK_SIZE = 1024 * 1024 +DEFAULT_MAX_FILE = 256 * 2**20 +DEFAULT_MAX_CAPTURES = 8 +DEFAULT_MAX_BYTES = 512 * 2**20 +DEFAULT_MAX_WORKERS = 2 +DEFAULT_IDLE_SECONDS = 900 +DEFAULT_HARD_SECONDS = 1800 +MAX_INFERENCE_SIDE = 4096 +MAX_INFERENCE_PIXELS = 16_777_216 +MAX_RELATIVE_LENGTH = 1024 + +VIDEO_ONLY_REASON = "Auto-tagging is available for supported still images only." +HEADER_REASON = "The image header could not be read." +DIMENSION_REASON = ( + "Images larger than 4096 px per side or 16,777,216 px are not supported." +) + +# The listing's media allowlist (``py.api.inspection_service.SERVABLE_EXTENSIONS``), +# restated here so utilities never import the aiohttp API package; a test pins +# the two sets equal. +_STILL_EXTENSIONS = frozenset( + {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"} +) +CAPTURE_EXTENSIONS = _STILL_EXTENSIONS | frozenset(VIDEO_EXTENSIONS) + +_DIRECTORY_PREFIX = "prompt-manager-captures-" +_OPEN_FLAGS = ( + os.O_RDONLY + | getattr(os, "O_BINARY", 0) + | getattr(os, "O_NOFOLLOW", 0) + | getattr(os, "O_CLOEXEC", 0) +) + + +class CaptureError(Exception): + """Base of the capture outcomes; ``code`` is the API error code.""" + + code = "capture-error" + + +class CaptureCapacity(CaptureError): + """Count, byte or worker admission failed (nonblocking).""" + + code = "review-capacity" + + +class CaptureTooLarge(CaptureError): + """The file exceeds the per-file byte cap on stat or on the stream.""" + + code = "unsupported-media" + + +class CaptureChanged(CaptureError): + """The file changed, vanished or was replaced across the copy (retryable).""" + + code = "file-changed" + + +class CaptureDenied(CaptureError): + """The target is not a listed rooted media file (alias, escape, non-media).""" + + code = "source-changed" + + +@dataclass(frozen=True) +class FileFacts: + """Whole-file facts of the current rooted file, without a retained copy.""" + + locator: str + sha256: str + size: int + mtime_ns: int + + +@dataclass(frozen=True) +class Capture: + """One completed private copy; ``path`` is the owner-only temporary file.""" + + token: str + locator: str + root_index: int + scope: str + epoch: str + sha256: str + size: int + mtime_ns: int + path: str + media_type: str + width: int | None + height: int | None + inferable: bool + inferable_reason: str | None + + +class _Entry: + __slots__ = ("capture", "created_at", "last_used_at", "uses", "doomed") + + def __init__(self, capture: Capture, now: float): + self.capture = capture + self.created_at = now + self.last_used_at = now + self.uses = 0 + self.doomed = False + + +def _same_facts(before: os.stat_result, after: os.stat_result) -> bool: + """The facts every stat API reports alike: size and modification time. + + This is the only comparison made across APIs (a path ``stat`` against the + opened descriptor's ``fstat``). Windows reports the creation time from + ``fstat`` but, on Python 3.12+, the change time from ``stat``, and may + encode device and inode differently, so those fields are compared only + between two reads of the same kind (``_same_identity``). + """ + return before.st_size == after.st_size and before.st_mtime_ns == after.st_mtime_ns + + +def _same_identity(before: os.stat_result, after: os.stat_result) -> bool: + """Full identity between two reads of the same kind (both ``fstat`` or both + path stats): size, mtime, change time, and device/inode where meaningful.""" + if not _same_facts(before, after) or before.st_ctime_ns != after.st_ctime_ns: + return False + if before.st_ino and after.st_ino and before.st_ino != after.st_ino: + return False + if before.st_dev and after.st_dev and before.st_dev != after.st_dev: + return False + return True + + +def _reason(error: OSError) -> str: + """The OS error text without the path, so messages never name server paths.""" + return error.strerror or type(error).__name__ + + +def _unlink_quietly(path: str) -> None: + try: + os.unlink(path) + except FileNotFoundError: + pass + except OSError as error: + logger.warning("Capture file %s could not be removed: %s", path, error) + + +class CaptureStore: + """Bounded, lock-protected registry of private media captures.""" + + def __init__( + self, + temp_parent: os.PathLike | str | None = None, + *, + max_file: int = DEFAULT_MAX_FILE, + max_captures: int = DEFAULT_MAX_CAPTURES, + max_bytes: int = DEFAULT_MAX_BYTES, + max_workers: int = DEFAULT_MAX_WORKERS, + idle: float = DEFAULT_IDLE_SECONDS, + hard: float = DEFAULT_HARD_SECONDS, + clock: Callable[[], float] = time.monotonic, + ): + for name, value in ( + ("max_file", max_file), + ("max_captures", max_captures), + ("max_bytes", max_bytes), + ("max_workers", max_workers), + ): + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{name} must be a positive integer") + if not (idle > 0 and hard > 0): + raise ValueError("idle and hard must be positive") + self.max_file = max_file + self.max_captures = max_captures + self.max_bytes = max_bytes + self.max_workers = max_workers + self.idle = idle + self.hard = hard + self._clock = clock + parent = None if temp_parent is None else os.fspath(temp_parent) + if parent is not None: + os.makedirs(parent, exist_ok=True) + # mkdtemp creates the directory with owner-only permissions. + self.directory = tempfile.mkdtemp(prefix=_DIRECTORY_PREFIX, dir=parent) + self._lock = threading.Lock() + self._entries: dict[str, _Entry] = {} + self._completed_bytes = 0 + self._pending = 0 + self._pending_bytes = 0 + self._workers = 0 + self._closed = False + self._stale: list[str] = [] + + # -- public API ------------------------------------------------------------ + + def capture( + self, + root: os.PathLike | str, + relative: str, + *, + epoch: str, + scope: str, + root_index: int = 0, + ) -> Capture: + """Copy and hash one rooted media file into a private capture.""" + if isinstance(root_index, bool) or not isinstance(root_index, int): + raise CaptureDenied("Invalid gallery root") + if root_index < 0: + raise CaptureDenied("Invalid gallery root") + for name, value in (("epoch", epoch), ("scope", scope)): + if not isinstance(value, str) or not value: + raise ValueError(f"{name} must be a non-empty string") + parts, extension = _relative_parts(relative) + candidate, locator = _resolve(root, parts) + before = self._pre_stat(candidate) + self._admit(before.st_size, reserve=True) + try: + capture = self._copy( + candidate, locator, extension, before, root_index, epoch, scope + ) + finally: + self._release_admission(before.st_size, reserve=True) + return capture + + def verify(self, root: os.PathLike | str, relative: str) -> FileFacts: + """Hash the current rooted file under the same bounds, keeping no copy.""" + parts, _extension = _relative_parts(relative) + candidate, locator = _resolve(root, parts) + before = self._pre_stat(candidate) + self._admit(0, reserve=False) + try: + sha256 = self._stream(candidate, locator, before, None) + finally: + self._release_admission(0, reserve=False) + return FileFacts(locator, sha256, before.st_size, before.st_mtime_ns) + + def get(self, token: str) -> Capture | None: + """The live capture for *token*, touching its idle clock; else None. + + ``get`` does not pin the file: a reader that streams ``capture.path`` + (the preview route) must bracket the read with ``acquire`` and + ``release_use``, or a concurrent ``sweep``/``release`` may unlink the + file under it (and, on Windows, leak it until ``close``). + """ + try: + with self._lock: + entry = self._live_locked(token) + if entry is None: + return None + entry.last_used_at = self._clock() + return entry.capture + finally: + self._flush() + + def acquire(self, token: str) -> Capture | None: + """Mark the capture in use (inference); it outlives expiry until released.""" + try: + with self._lock: + entry = self._live_locked(token) + if entry is None: + return None + entry.uses += 1 + entry.last_used_at = self._clock() + return entry.capture + finally: + self._flush() + + def release_use(self, token: str) -> None: + """End one ``acquire``; a doomed capture is removed with its last use.""" + with self._lock: + entry = self._entries.get(token) + if entry is not None: + entry.uses = max(0, entry.uses - 1) + entry.last_used_at = self._clock() + if entry.doomed and entry.uses == 0: + self._remove_locked(entry) + self._flush() + + def release(self, token: str) -> bool: + """Explicit release; True when the token was live.""" + with self._lock: + entry = self._entries.get(token) + live = entry is not None and not entry.doomed + if live: + self._doom_locked(entry) + self._flush() + return live + + def invalidate_all(self) -> None: + """Expire every capture now (storage epoch changed, restore, shutdown).""" + with self._lock: + for entry in list(self._entries.values()): + self._doom_locked(entry) + self._flush() + + def sweep(self) -> None: + """Remove idle and hard-expired captures that are not in use.""" + with self._lock: + self._sweep_locked(self._clock()) + self._flush() + + def expires_in(self, token: str) -> float | None: + """Seconds until *token* expires without further use; None when gone.""" + try: + with self._lock: + entry = self._live_locked(token) + if entry is None: + return None + now = self._clock() + remaining = self.hard - (now - entry.created_at) + if entry.uses == 0: + remaining = min(remaining, self.idle - (now - entry.last_used_at)) + return max(0.0, remaining) + finally: + self._flush() + + def stats(self) -> dict: + with self._lock: + return { + "captures": len(self._entries), + "reserved_bytes": self._completed_bytes + self._pending_bytes, + "workers": self._workers, + } + + def close(self) -> None: + """Remove every owned capture file and the owned directory. Idempotent. + + A capture still in flight finishes first: its worker removes its own + temp file and then the directory (``_release_admission``). + """ + with self._lock: + self._closed = True + for entry in list(self._entries.values()): + self._remove_locked(entry) + busy = self._pending > 0 + self._flush() + if not busy: + self._remove_directory() + + def _remove_directory(self) -> None: + try: + os.rmdir(self.directory) + except FileNotFoundError: + pass + except OSError as error: + logger.warning( + "Capture directory %s could not be removed: %s", self.directory, error + ) + + # -- admission ------------------------------------------------------------- + + def _admit(self, size: int, *, reserve: bool) -> None: + try: + self._admit_locked(size, reserve=reserve) + finally: + self._flush() + + def _admit_locked(self, size: int, *, reserve: bool) -> None: + with self._lock: + if self._closed: + raise RuntimeError("The capture store is closed") + self._sweep_locked(self._clock()) + if self._workers >= self.max_workers: + raise CaptureCapacity("Too many captures are in progress") + if reserve: + if len(self._entries) + self._pending >= self.max_captures: + raise CaptureCapacity("Too many captures are held") + reserved = self._completed_bytes + self._pending_bytes + if reserved + size > self.max_bytes: + raise CaptureCapacity("Capture byte budget exhausted") + self._pending += 1 + self._pending_bytes += size + self._workers += 1 + + def _release_admission(self, size: int, *, reserve: bool) -> None: + with self._lock: + self._workers -= 1 + if reserve: + self._pending -= 1 + self._pending_bytes -= size + last_out = self._closed and self._pending == 0 and self._workers == 0 + if last_out: + self._remove_directory() + + def _register(self, capture: Capture) -> None: + with self._lock: + if self._closed: + raise RuntimeError("The capture store is closed") + self._completed_bytes += capture.size + self._entries[capture.token] = _Entry(capture, self._clock()) + + # -- lifetime (all under the lock) ----------------------------------------- + + def _expired(self, entry: _Entry, now: float) -> bool: + if now - entry.created_at >= self.hard: + return True + return entry.uses == 0 and now - entry.last_used_at >= self.idle + + def _live_locked(self, token: str) -> _Entry | None: + entry = self._entries.get(token) + if entry is None or entry.doomed: + return None + if self._expired(entry, self._clock()): + self._doom_locked(entry) + return None + return entry + + def _sweep_locked(self, now: float) -> None: + for entry in list(self._entries.values()): + if entry.doomed or self._expired(entry, now): + self._doom_locked(entry) + + def _doom_locked(self, entry: _Entry) -> None: + entry.doomed = True + if entry.uses == 0: + self._remove_locked(entry) + + def _remove_locked(self, entry: _Entry) -> None: + """Forget the entry; its file is unlinked by ``_flush`` after the lock.""" + if self._entries.pop(entry.capture.token, None) is not None: + self._completed_bytes -= entry.capture.size + self._stale.append(entry.capture.path) + + def _flush(self) -> None: + """Unlink files released under the lock, without holding it (slow I/O).""" + with self._lock: + stale, self._stale = self._stale, [] + for path in stale: + _unlink_quietly(path) + + # -- copy and hash ----------------------------------------------------------- + + def _pre_stat(self, candidate: str) -> os.stat_result: + try: + facts = os.lstat(candidate) + except FileNotFoundError: + raise CaptureChanged("The file is no longer there") from None + except OSError as error: + logger.debug("Capture pre-stat failed for %s: %s", candidate, error) + raise CaptureChanged( + f"The file could not be read ({_reason(error)})" + ) from None + if stat_module.S_ISLNK(facts.st_mode): + raise CaptureDenied("Symbolic links are not captured") + if not stat_module.S_ISREG(facts.st_mode): + raise CaptureDenied("Only regular files are captured") + if facts.st_size > self.max_file: + raise CaptureTooLarge("The file is larger than the capture limit") + return facts + + def _copy( + self, + candidate: str, + locator: str, + extension: str, + before: os.stat_result, + root_index: int, + epoch: str, + scope: str, + ) -> Capture: + fd, temp_path = tempfile.mkstemp( + prefix="capture-", suffix=extension, dir=self.directory + ) + sink = os.fdopen(fd, "wb") + completed = False + try: + sha256 = self._stream(candidate, locator, before, sink) + sink.close() + media_type = "video" if is_video_path(candidate) else "image" + width, height, inferable, reason = self._probe(temp_path, media_type) + capture = Capture( + token=secrets.token_urlsafe(32), + locator=locator, + root_index=root_index, + scope=scope, + epoch=epoch, + sha256=sha256, + size=before.st_size, + mtime_ns=before.st_mtime_ns, + path=temp_path, + media_type=media_type, + width=width, + height=height, + inferable=inferable, + inferable_reason=reason, + ) + self._register(capture) + completed = True + return capture + finally: + if not sink.closed: + sink.close() + if not completed: + _unlink_quietly(temp_path) + + def _read(self, stream, size: int) -> bytes: + """One bounded read of the source; a seam for change-detection tests.""" + return stream.read(size) + + def _stream(self, candidate: str, locator: str, before, sink) -> str: + """Hash (and copy into *sink*) the file; every handle closes here.""" + digest = hashlib.sha256() + total = 0 + try: + fd = os.open(candidate, _OPEN_FLAGS) + except OSError as error: + logger.debug("Capture open failed for %s: %s", candidate, error) + raise CaptureChanged( + f"The file could not be opened ({_reason(error)})" + ) from None + with os.fdopen(fd, "rb") as stream: + opened = os.fstat(stream.fileno()) + if not _same_facts(before, opened): + raise CaptureChanged("The file changed before it was read") + while True: + chunk = self._read(stream, CHUNK_SIZE) + if not chunk: + break + total += len(chunk) + if total > self.max_file: + raise CaptureTooLarge("The file grew past the capture limit") + digest.update(chunk) + if sink is not None: + sink.write(chunk) + after = os.fstat(stream.fileno()) + if total != before.st_size or not _same_identity(opened, after): + raise CaptureChanged("The file changed while it was read") + try: + current = os.lstat(candidate) + except OSError: + raise CaptureChanged("The file is no longer there") from None + if not _same_identity(before, current): + raise CaptureChanged("The path names another file now") + if os.path.normcase(os.path.realpath(candidate)) != locator: + raise CaptureChanged("The path resolves elsewhere now") + return digest.hexdigest() + + def _probe(self, path: str, media_type: str): + """Header-only dimensions of the capture; never decodes pixels. + + Uses the inspection reader, so a PNG whose prompt metadata is over the + inspection budgets is reported as unreadable (not inferable) although + its pixels decode; B2 may switch to a dimensions-only header read. + """ + if media_type == "video": + return None, None, False, VIDEO_ONLY_REASON + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", Image.DecompressionBombWarning) + read = read_inspection_file(Path(path)) + except (MetadataReadError, Image.DecompressionBombError): + return None, None, False, HEADER_REASON + width, height = read.width, read.height + if width is None or height is None: + return None, None, False, HEADER_REASON + if ( + width > MAX_INFERENCE_SIDE + or height > MAX_INFERENCE_SIDE + or width * height > MAX_INFERENCE_PIXELS + ): + return width, height, False, DIMENSION_REASON + return width, height, True, None + + +def _relative_parts(relative) -> tuple[list[str], str]: + """Split a safe relative POSIX media path; anything else is denied.""" + if not isinstance(relative, str) or not relative: + raise CaptureDenied("Invalid file path") + if len(relative) > MAX_RELATIVE_LENGTH or "\x00" in relative: + raise CaptureDenied("Invalid file path") + if "\\" in relative or relative.startswith("/"): + raise CaptureDenied("Invalid file path") + parts = relative.split("/") + for part in parts: + if part in ("", ".", "..") or ":" in part: + raise CaptureDenied("Invalid file path") + extension = os.path.splitext(parts[-1])[1].lower() + if extension not in CAPTURE_EXTENSIONS: + raise CaptureDenied("Unsupported media type") + return parts, extension + + +def _resolve(root, parts: list[str]) -> tuple[str, str]: + """The candidate path and its locator; aliases and escapes are denied.""" + real_root = os.path.realpath(os.fspath(root)) + if not os.path.isdir(real_root): + raise CaptureDenied("Unknown gallery root") + candidate = os.path.join(real_root, *parts) + locator = os.path.normcase(candidate) + if os.path.normcase(os.path.realpath(candidate)) != locator: + raise CaptureDenied("The file is not a listed rooted file") + return candidate, locator diff --git a/web/js/workspace/toolbar.mjs b/web/js/workspace/toolbar.mjs index fc84104..49cebf5 100644 --- a/web/js/workspace/toolbar.mjs +++ b/web/js/workspace/toolbar.mjs @@ -107,7 +107,10 @@ export function mountToolbar(container, { scope, categories = [], sorts, value = category.value = [...category.options].some(node => node.value === selected) ? selected : ''; } function setActiveFilters(count) { - filters.textContent = filtersLabel(count); + const next = filtersLabel(count); + // Every listing reload reports the count; only a changed label rewrites the text node, because WebKit drops a + // click whose mousedown hit a text node that was replaced before mouseup. + if (filters.textContent !== next) filters.textContent = next; } setValue(value); diff --git a/web/js/workspace/vocabulary-state.mjs b/web/js/workspace/vocabulary-state.mjs index 6f35cc5..18447f4 100644 --- a/web/js/workspace/vocabulary-state.mjs +++ b/web/js/workspace/vocabulary-state.mjs @@ -16,7 +16,8 @@ import { countOf } from './plural.mjs'; const LIMIT = 100; const BASE = '/prompt_manager/vocabulary'; const KINDS = { - tags: {key: row => row.id, valid: row => id(row.id) && typeof row.name === 'string' && row.name.length > 0}, + // image_count (#249 image tags) is optional until every server reports it. + tags: {key: row => row.id, valid: row => id(row.id) && typeof row.name === 'string' && row.name.length > 0 && optionalCount(row.image_count)}, categories: {key: row => row.value, valid: row => typeof row.value === 'string' && row.value.length > 0}, }; const RECOVERY = {'storage-changed': 'storage-changed', 'vocabulary-conflict': 'conflict'}; @@ -29,6 +30,7 @@ const IMPACT_FIELDS = ['positive_prompts', 'other_prompts', 'negative_variants', const object = value => value !== null && typeof value === 'object' && !Array.isArray(value); const id = value => Number.isSafeInteger(value) && value > 0; const count = value => Number.isSafeInteger(value) && value >= 0; +const optionalCount = value => value === undefined || count(value); const epochOf = value => typeof value === 'string' && /^[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}$/.test(value); const hexToken = value => typeof value === 'string' && /^[0-9a-f]{64}$/.test(value); @@ -108,7 +110,7 @@ export function createVocabularyStream({api, kind, onChange = () => {}}) { } function validateTag(value) {return object(value) && KINDS.tags.valid(value) && count(value.prompt_count);} -function validateImpact(value) {return object(value) && IMPACT_FIELDS.every(name => count(value[name]));} +function validateImpact(value) {return object(value) && IMPACT_FIELDS.every(name => count(value[name])) && optionalCount(value.images);} function requireRow(row) {requireValue(object(row) && id(row.id) && typeof row.name === 'string', 'Invalid tag row');} function requireEpoch(epoch) {requireValue(epochOf(epoch), 'Invalid library epoch');} function requireText(value, name) {requireValue(typeof value === 'string', `${name} must be a string`);} @@ -116,8 +118,9 @@ function requireText(value, name) {requireValue(typeof value === 'string', `${na const RECEIPTS = { 'create-tag': value => typeof value.created === 'boolean' && validateTag(value.tag), 'rename-tag': value => typeof value.merged === 'boolean' && validateTag(value.tag) && - count(value.components_changed) && count(value.pairs_changed), - 'delete-tag': value => value.deleted === true && count(value.components_changed) && count(value.pairs_changed), + count(value.components_changed) && count(value.pairs_changed) && optionalCount(value.images_changed), + 'delete-tag': value => value.deleted === true && count(value.components_changed) && count(value.pairs_changed) && + optionalCount(value.images_changed), 'apply-category': (value, body) => value.source === body.source && value.target === body.target && typeof value.consolidated === 'boolean' && validateImpact(value.impact), }; @@ -270,6 +273,7 @@ function impactList(impact) { parts.push(`${countOf(negatives, 'negative variant')}${archived}`); } if (impact.other_prompts) parts.push(countOf(impact.other_prompts, 'other saved prompt')); + if (impact.images) parts.push(countOf(impact.images, 'image')); return parts.length > 1 ? `${parts.slice(0, -1).join(', ')} and ${parts.at(-1)}` : parts.join(''); } @@ -282,6 +286,6 @@ export function categoryPreviewText(shown) { /** The Delete tag confirmation sentence (decision D14). */ export function deleteTagText(name, impact) { const list = impactList(impact); - if (!list) return `“${name}” is not linked to any saved prompt. Deleting it removes only the name.`; + if (!list) return `“${name}” is not linked to any saved prompt or image. Deleting it removes only the name.`; return `Deleting “${name}” removes the name and its links from ${list}. The prompts, negative variants, images and history themselves are kept.`; }