diff --git a/database/image_categories.py b/database/image_categories.py new file mode 100644 index 00000000..b8657f1d --- /dev/null +++ b/database/image_categories.py @@ -0,0 +1,242 @@ +"""Image categories (#332): an image file's own category. + +An image shows its own category when it has one, otherwise the category its +link supplies (its recipe's, else its positive prompt's); this table holds +only the own values. A row is keyed by the file key: the deepest configured +output source that contains the file (its real, case-normalized path) and the +path below it, case-folded where the OS folds case +(``py/api/images.py`` ``image_category_key``). Without nested sources the key +equals a favorite's key. Nothing follows a moved file: its row stays and is +hidden until a file appears at the same path again. + +Writes validate strictly (the route strips the value and applies +``validate_category``). Startup checks the structure only, so a value saved +under older rules never stops the library opening; restore also checks the +values. Optional in old backups, required once a database is current. +""" + +import re +import sqlite3 +from typing import Iterable + +TABLE = "image_categories" +INDEX = "idx_image_categories_category" +MAX_KEY_LENGTH = 4096 +MAX_CATEGORY_LENGTH = 100 +_DDL = ( + f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT NOT NULL, " + "category TEXT NOT NULL, updated_at TEXT NOT NULL, " + "PRIMARY KEY (root_key, relative_path))" +) +_INDEX_DDL = f"CREATE INDEX {INDEX} ON {TABLE}(category)" +_AUTOINDEX = f"sqlite_autoindex_{TABLE}_1" +_COLUMNS = ( + ("root_key", "TEXT", 1, 1), + ("relative_path", "TEXT", 1, 2), + ("category", "TEXT", 1, 0), + ("updated_at", "TEXT", 1, 0), +) +# A category may not hold a character below U+0020 (validate_category). +_CONTROL_GLOB = "*[" + "".join(chr(code) for code in range(1, 32)) + "]*" + +CATEGORY_OF_SQL = ( + f"SELECT category FROM {TABLE} WHERE root_key = ? AND relative_path = ?" +) +CATEGORIES_OF_ROOT_SQL = ( + f"SELECT relative_path, category FROM {TABLE} WHERE root_key = ?" +) +UPSERT_SQL = ( + f"INSERT INTO {TABLE} (root_key, relative_path, category, updated_at) " + "VALUES (?, ?, ?, ?) ON CONFLICT(root_key, relative_path) " + "DO UPDATE SET category = excluded.category, updated_at = excluded.updated_at" +) +DELETE_SQL = f"DELETE FROM {TABLE} WHERE root_key = ? AND relative_path = ?" +# The duplicate cleanup's copy (#332 MC11): the removed copy's row is copied +# to the kept copy only when the kept copy has none; it is never deleted. +COPY_SQL = ( + f"INSERT OR IGNORE INTO {TABLE} (root_key, relative_path, category, updated_at) " + f"SELECT ?, ?, category, updated_at FROM {TABLE} " + "WHERE root_key = ? AND relative_path = ?" +) +STATEMENTS = (CATEGORY_OF_SQL, CATEGORIES_OF_ROOT_SQL, UPSERT_SQL, DELETE_SQL, COPY_SQL) + + +def _compact(sql: str) -> str: + return re.sub(r"\s+", "", sql or "").lower() + + +def _structure_error(conn: sqlite3.Connection, allow_missing: bool) -> str | None: + objects = conn.execute( + "SELECT type,name,sql FROM sqlite_master WHERE name=? " + "OR (type IN ('index','trigger') AND tbl_name=?)", + (TABLE, TABLE), + ).fetchall() + if not objects: + return None if allow_missing else "Missing image categories schema" + tables = [sql for kind, _name, sql in objects if kind == "table"] + others = {(kind, name) for kind, name, _sql in objects if kind != "table"} + if len(tables) != 1 or others != {("index", _AUTOINDEX), ("index", INDEX)}: + return "Image categories name is reserved by an unsupported schema object" + if _compact(tables[0]) != _compact(_DDL): + return "Image categories has an invalid table definition" + index_sql = next(sql for kind, name, sql in objects if name == INDEX) + if _compact(index_sql) != _compact(_INDEX_DDL): + return "Image categories has an invalid index" + autoindex_sql = next(sql for kind, name, sql in objects if name == _AUTOINDEX) + if autoindex_sql is not None: + return "Image categories has an invalid index" + columns = tuple( + (row[1], row[2].upper(), row[3], row[5]) + for row in conn.execute(f"PRAGMA table_info({TABLE})") + ) + if columns != _COLUMNS: + return "Image categories has invalid columns" + return None + + +def _values_error(conn: sqlite3.Connection) -> str | None: + # Bytes bound characters from above (UTF-8 is at most 4 bytes each). + limit = MAX_KEY_LENGTH * 4 + if conn.execute( + f"SELECT 1 FROM {TABLE} WHERE typeof(root_key) != 'text' " # nosec B608 + "OR typeof(relative_path) != 'text' OR typeof(category) != 'text' " + "OR typeof(updated_at) != 'text' " + f"OR length(CAST(root_key AS BLOB)) > {limit} " + f"OR length(CAST(relative_path AS BLOB)) > {limit} " + "OR length(CAST(updated_at AS BLOB)) > 64 " + "OR instr(CAST(category AS BLOB), x'00') > 0 " + f"OR length(category) < 1 OR length(category) > {MAX_CATEGORY_LENGTH} " + "OR category != TRIM(category) OR category GLOB ? LIMIT 1", + (_CONTROL_GLOB,), + ).fetchone(): + return "Image categories contains invalid values" + for root, path in conn.execute( + f"SELECT CAST(root_key AS BLOB), CAST(relative_path AS BLOB) " # nosec B608 + f"FROM {TABLE} WHERE length(CAST(root_key AS BLOB)) > {MAX_KEY_LENGTH} " + f"OR length(CAST(relative_path AS BLOB)) > {MAX_KEY_LENGTH}" + ): + for raw in (root, path): + if len(bytes(raw).decode("utf-8", "replace")) > MAX_KEY_LENGTH: + return "Image categories contains invalid values" + return None + + +def image_categories_schema_error( + conn: sqlite3.Connection, + *, + allow_missing: bool = True, + check_values: bool = True, +) -> str | None: + """Reject an unsupported image categories schema (and, optionally, values). + + Reads only. A missing table is acceptable for backups that predate it + unless ``allow_missing`` is false. ``check_values`` also refuses a stored + value the write path could never have written. + """ + error = _structure_error(conn, allow_missing) + if error is not None or not check_values: + return error + exists = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (TABLE,) + ).fetchone() + return _values_error(conn) if exists else None + + +def migrate_image_categories(conn: sqlite3.Connection) -> None: + """Install the table and its index atomically; structure is checked + before and after. Reads only ``sqlite_master`` and this table's PRAGMAs + and writes only its own two objects (no backfill).""" + conn.execute("SAVEPOINT image_categories_schema") + try: + error = image_categories_schema_error(conn, check_values=False) + if error: + raise sqlite3.DatabaseError(error) + conn.execute(_DDL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1)) + conn.execute( + _INDEX_DDL.replace("CREATE INDEX", "CREATE INDEX IF NOT EXISTS", 1) + ) + error = image_categories_schema_error( + conn, allow_missing=False, check_values=False + ) + if error: + raise sqlite3.DatabaseError(error) + conn.execute("RELEASE image_categories_schema") + except BaseException: + conn.execute("ROLLBACK TO image_categories_schema") + conn.execute("RELEASE image_categories_schema") + raise + + +def copy_image_category(conn: sqlite3.Connection, removed_key, kept_key) -> int: + """Give the kept copy the removed copy's category when it has none. + + Runs inside the caller's transaction; never deletes the removed copy's + row, so a give-back or a failed delete loses nothing. Returns 1 when a + row was inserted, else 0. + """ + if tuple(removed_key) == tuple(kept_key): + return 0 + return max(conn.execute(COPY_SQL, (*kept_key, *removed_key)).rowcount, 0) + + +class ImageCategoryOperations: + """PromptDatabase mixin: set, clear and read images' own categories. + + Never reads or writes ``prompts``, ``prompt_pairs`` or + ``generated_images``. + """ + + def image_category(self, root_key: str, relative_key: str) -> str | None: + """One file's own category, or None.""" + row = ( + self.model.get_connection() + .execute(CATEGORY_OF_SQL, (root_key, relative_key)) + .fetchone() + ) + return row[0] if row else None + + def own_categories_of(self, root_keys: Iterable[str]) -> dict: + """``{(root key, relative key): category}`` under the given sources. + + One primary-key range read per source. + """ + found = {} + conn = self.model.get_connection() + for root_key in dict.fromkeys(root_keys): + for relative, category in conn.execute(CATEGORIES_OF_ROOT_SQL, (root_key,)): + found[(root_key, relative)] = category + return found + + def set_image_categories(self, file_keys, category: str) -> dict: + """Set *category* on every file key, in one transaction. + + ``file_keys``: ``(root key, relative key)`` pairs from any number of + sources. Returns ``{updated, unchanged}``; a file that already had + exactly this value is ``unchanged`` and keeps its ``updated_at``. + """ + from .models import utc_now_iso # models imports this module's schema + + keys = list(dict.fromkeys(tuple(key) for key in file_keys)) + counts = {"updated": 0, "unchanged": 0} + now = utc_now_iso() + with self._pair_transaction() as conn: + for key in keys: + row = conn.execute(CATEGORY_OF_SQL, key).fetchone() + if row is not None and row[0] == category: + counts["unchanged"] += 1 + continue + conn.execute(UPSERT_SQL, (*key, category, now)) + counts["updated"] += 1 + return counts + + def clear_image_categories(self, file_keys) -> int: + """Remove the own category of every file key (one transaction). + + Returns the number of rows removed. + """ + keys = list(dict.fromkeys(tuple(key) for key in file_keys)) + removed = 0 + with self._pair_transaction() as conn: + for key in keys: + removed += conn.execute(DELETE_SQL, key).rowcount + return removed diff --git a/database/models.py b/database/models.py index dad9cf07..25e98e44 100644 --- a/database/models.py +++ b/database/models.py @@ -34,6 +34,7 @@ from .review_schema import ( review_without_annotations_error, ) from .saved_views import migrate_saved_views, saved_views_schema_error +from .image_categories import image_categories_schema_error, migrate_image_categories from .workflow_documents import ( DOCUMENT_ID_COLUMNS, migrate_workflow_documents, @@ -350,6 +351,8 @@ class PromptModel: migrate_annotation_state(conn) migrate_review_state(conn) migrate_autotag_state(conn) + # Image categories (#332): independent of every table above. + migrate_image_categories(conn) # Last: the older generated_images rebuilds run only before file_path. migrate_workflow_documents(conn) @@ -1155,8 +1158,9 @@ class PromptModel: Args: path: File to inspect require_library_state: Also require the tables every current - database has (Library state, image favorites, saved views, - image annotations and the review queue), as a migrated + database has (Library state, image favorites, image + categories, saved views, image annotations and the review + queue), as a migrated restore stage must; old backups may lack them. documents: Read every workflow document and refuse a damaged one. A restore checks its upload without it and repairs @@ -1220,6 +1224,11 @@ class PromptModel: ) if favorites_error: return False, favorites_error + categories_error = image_categories_schema_error( + conn, allow_missing=not require_library_state + ) + if categories_error: + return False, categories_error views_error = saved_views_schema_error( conn, allow_missing=not require_library_state ) diff --git a/database/operations.py b/database/operations.py index 0f45097f..4460fcf3 100644 --- a/database/operations.py +++ b/database/operations.py @@ -27,6 +27,7 @@ from .health import HealthOperations, iter_consistency_issues from .duplicate_copies import DuplicateCopyOperations from .workflow_storage import WorkflowStorageOperations from .category_import import CategoryImportOperations, EmptyReplacement # noqa: F401 +from .image_categories import ImageCategoryOperations from .workflow_documents import ( DocumentReader, documents_of_rows, @@ -368,6 +369,7 @@ class PromptDatabase( DuplicateCopyOperations, WorkflowStorageOperations, CategoryImportOperations, + ImageCategoryOperations, ): """Database operations class for managing prompts.""" diff --git a/py/api/images.py b/py/api/images.py index 439a2800..dc28caec 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -800,6 +800,49 @@ def _favorite_root_key(root) -> str: return _favorite_relative_key(os.path.realpath(str(root))) +def _source_depth(source: str) -> int: + return len([part for part in re.split(r"[\\/]+", source) if part]) + + +def deepest_source(sources, absolute_path): + """``(source, POSIX path below it)`` of the deepest source holding the path. + + *sources* are compared by their real path (they are stored resolved); the + path is compared lexically after ``normpath`` and never resolved, so a + symlinked subfolder keeps its listed spelling. Containment folds case + where the OS does. None when no source contains the path. + """ + path = os.path.normpath(str(absolute_path)) + folded = os.path.normcase(path) + best = None + for source in sources: + real = os.path.normpath(os.path.realpath(str(source))) + prefix = os.path.join(os.path.normcase(real), "") + if not folded.startswith(prefix): + continue + if best is None or _source_depth(real) > _source_depth(best): + best = real + if best is None: + return None + relative = os.path.relpath(path, best) + return best, relative.replace(os.sep, "/") + + +def image_category_key(sources, absolute_path): + """The file key of an image's own category (#332), or None. + + ``(root key, relative key)``: the deepest output source that contains the + file and the path below it, keyed with the favorites helpers. Without + nested sources it equals the file's favorite key. A path no source + contains has no key (and so no own category). + """ + found = deepest_source(sources, absolute_path) + if found is None: + return None + source, relative = found + return _favorite_root_key(source), _favorite_relative_key(relative) + + def _name_suffix(name): """``Path(name).suffix`` as Python 3.10-3.13 compute it. diff --git a/tests/browser/server.py b/tests/browser/server.py index 71ea5669..160f4b0a 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -268,6 +268,7 @@ def create_app(root: Path) -> web.Application: "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_concept_fixture.py b/tests/browser/workspace_concept_fixture.py index 1c86dbbd..0c54e831 100644 --- a/tests/browser/workspace_concept_fixture.py +++ b/tests/browser/workspace_concept_fixture.py @@ -181,6 +181,7 @@ _TABLES = ( "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "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 8c0aed9e..ffc0534a 100644 --- a/tests/browser/workspace_library_editor_fixture.py +++ b/tests/browser/workspace_library_editor_fixture.py @@ -33,6 +33,7 @@ _TABLES = ( "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_review_fixture.py b/tests/browser/workspace_review_fixture.py index 45365dd7..16344eb2 100644 --- a/tests/browser/workspace_review_fixture.py +++ b/tests/browser/workspace_review_fixture.py @@ -85,6 +85,7 @@ _TABLES = ( "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/browser/workspace_tagging_fixture.py b/tests/browser/workspace_tagging_fixture.py index 01667ee8..1a6801c7 100644 --- a/tests/browser/workspace_tagging_fixture.py +++ b/tests/browser/workspace_tagging_fixture.py @@ -51,6 +51,7 @@ _TABLES = ( "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "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 f10cf457..6fbe1fcc 100644 --- a/tests/browser/workspace_visual_fixture.py +++ b/tests/browser/workspace_visual_fixture.py @@ -119,6 +119,7 @@ def seed_workspace_visual(api: PromptManagerAPI, root: Path) -> None: "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "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 17901872..95568762 100644 --- a/tests/browser/workspace_vocabulary_fixture.py +++ b/tests/browser/workspace_vocabulary_fixture.py @@ -86,6 +86,7 @@ _TABLES = ( "media_image_tags", "media_annotation_bytes", "media_annotation_files", + "image_categories", "image_prompt_pairs", "prompt_pair_tags", "prompt_pairs", diff --git a/tests/test_image_categories.py b/tests/test_image_categories.py new file mode 100644 index 00000000..e6560d56 --- /dev/null +++ b/tests/test_image_categories.py @@ -0,0 +1,513 @@ +"""Image categories (#332): schema, migration, restore check, store and file key.""" + +import hashlib +import ntpath +import os +import sqlite3 +import tempfile +import unittest +from contextlib import closing +from pathlib import Path +from unittest.mock import patch + +try: + from open_handles import assert_closed +except ImportError: + from tests.open_handles import assert_closed + +from database.image_categories import ( + COPY_SQL, + INDEX, + STATEMENTS, + TABLE, + copy_image_category, + image_categories_schema_error, + migrate_image_categories, +) +from database.models import PromptModel +from database.operations import PromptDatabase +from py.api.images import image_category_key + +_DDL = ( + f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT NOT NULL, " + "category TEXT NOT NULL, updated_at TEXT NOT NULL, " + "PRIMARY KEY (root_key, relative_path))" +) +_INDEX_DDL = f"CREATE INDEX {INDEX} ON {TABLE}(category)" + + +def category_checksum(conn) -> tuple: + """sha256 of prompts and prompt_pairs ``(id, category, updated_at)``.""" + sums = [] + for table in ("prompts", "prompt_pairs"): + digest = hashlib.sha256() + for row in conn.execute( # nosec B608 - fixed table names + f"SELECT id, category, updated_at FROM {table} ORDER BY id" + ): + digest.update(repr(tuple(row)).encode()) + sums.append(digest.hexdigest()) + return tuple(sums) + + +def table_checksums(conn) -> dict: + """``{name: (schema sql, sha256 of rows)}`` of every table but ours.""" + found = {} + for name, sql in conn.execute( + "SELECT name, sql FROM sqlite_master WHERE type='table' AND name != ?", + (TABLE,), + ).fetchall(): + digest = hashlib.sha256() + rows = conn.execute(f'SELECT * FROM "{name}"') # nosec B608 + for text in sorted(repr(tuple(row)) for row in rows): + digest.update(text.encode()) + found[name] = (sql, digest.hexdigest()) + return found + + +class _TempDir(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + self.path = os.path.join(self.tmpdir, "categories.db") + + def tearDown(self): + PromptModel.close_all_instances() + assert_closed(self, self.tmpdir) + for name in os.listdir(self.tmpdir): + os.unlink(os.path.join(self.tmpdir, name)) + os.rmdir(self.tmpdir) + + +class TestImageCategoriesSchema(_TempDir): + def test_migration_creates_table_and_index(self): + db = PromptDatabase(self.path) + db.close_all() + PromptModel.close_all_instances() + with closing(sqlite3.connect(self.path)) as conn: + names = { + (kind, name) + for kind, name in conn.execute( + "SELECT type, name FROM sqlite_master WHERE tbl_name=?", (TABLE,) + ) + } + self.assertEqual( + names, + { + ("table", TABLE), + ("index", INDEX), + ("index", f"sqlite_autoindex_{TABLE}_1"), + }, + ) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {TABLE}").fetchone()[0], 0 + ) + self.assertIsNone(image_categories_schema_error(conn, allow_missing=False)) + + def test_migration_leaves_prompt_categories_unchanged(self): + """23,000 links; the migration reads and writes only its own objects.""" + db = PromptDatabase(self.path) + db.close_all() + PromptModel.close_all_instances() + with closing(sqlite3.connect(self.path)) as conn: + conn.execute(f"DROP TABLE {TABLE}") # a library from before #332 + categories = ["Landscapes", "Portraits", None, " ", "Scans"] + conn.executemany( + "INSERT INTO prompts (text, category, hash, updated_at) " + "VALUES (?, ?, ?, ?)", + [ + (f"p{n}", categories[n % 5], f"h{n}", f"2026-01-01 00:00:{n % 60}") + for n in range(200) + ], + ) + conn.executemany( + "INSERT INTO prompt_pairs (positive_prompt_id, negative_prompt_id, " + "category) VALUES (?, ?, ?)", + [(n, n + 100, categories[(n + 1) % 5]) for n in range(1, 50)], + ) + conn.executemany( + "INSERT INTO generated_images (prompt_id, image_path, filename, " + "file_path) VALUES (?, ?, ?, ?)", + [ + (n % 200 + 1, f"/out/{n}.png", f"{n}.png", f"/out/{n}.png") + for n in range(23000) + ], + ) + conn.commit() + before = table_checksums(conn) + prompts_before = category_checksum(conn) + seen = [] + + def authorize(action, arg1, arg2, _db, _trigger): + seen.append((action, arg1, arg2)) + return sqlite3.SQLITE_OK + + conn.set_authorizer(authorize) + migrate_image_categories(conn) + conn.commit() + conn.set_authorizer(None) + own = {TABLE, "sqlite_master", "sqlite_schema", INDEX} + touched = set() + for action, arg1, arg2 in seen: + if action in (sqlite3.SQLITE_READ, sqlite3.SQLITE_INSERT): + touched.add(arg1) + if action == sqlite3.SQLITE_READ: + self.assertIn(arg1, own, (action, arg1, arg2)) + elif action in ( + sqlite3.SQLITE_INSERT, + sqlite3.SQLITE_UPDATE, + sqlite3.SQLITE_DELETE, + ): + self.assertIn(arg1, own, (action, arg1, arg2)) + elif action == sqlite3.SQLITE_CREATE_TABLE: + self.assertEqual(arg1, TABLE) + elif action == sqlite3.SQLITE_CREATE_INDEX: + self.assertIn(arg1, (INDEX, f"sqlite_autoindex_{TABLE}_1")) + self.assertEqual(arg2, TABLE) + elif action == sqlite3.SQLITE_PRAGMA: + self.assertEqual((arg1, arg2), ("table_info", TABLE)) + self.assertIn(TABLE, touched) # the authorizer saw the work + self.assertEqual(table_checksums(conn), before) + self.assertEqual(category_checksum(conn), prompts_before) + schema = conn.execute( + "SELECT sql FROM sqlite_master ORDER BY name" + ).fetchall() + migrate_image_categories(conn) # a second run changes nothing + conn.commit() + self.assertEqual( + conn.execute("SELECT sql FROM sqlite_master ORDER BY name").fetchall(), + schema, + ) + # The whole startup sequence keeps them too. + PromptDatabase(self.path).close_all() + PromptModel.close_all_instances() + with closing(sqlite3.connect(self.path)) as conn: + self.assertEqual(category_checksum(conn), prompts_before) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {TABLE}").fetchone()[0], 0 + ) + + def test_structure_accepted_and_refused(self): + variants = { + "missing": None, + "extra column": ( + f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT " + "NOT NULL, category TEXT NOT NULL, updated_at TEXT NOT NULL, x TEXT, " + "PRIMARY KEY (root_key, relative_path))", + _INDEX_DDL, + ), + "wrong index": (_DDL, f"CREATE INDEX {INDEX} ON {TABLE}(updated_at)"), + "no index": (_DDL,), + "extra index": ( + _DDL, + _INDEX_DDL, + f"CREATE INDEX extra ON {TABLE}(root_key)", + ), + "trigger": ( + _DDL, + _INDEX_DDL, + f"CREATE TRIGGER t AFTER INSERT ON {TABLE} BEGIN SELECT 1; END", + ), + "wrong ddl": ( + f"CREATE TABLE {TABLE} (root_key TEXT, relative_path TEXT, " + "category TEXT, updated_at TEXT, " + "PRIMARY KEY (root_key, relative_path))", + _INDEX_DDL, + ), + "view": (f"CREATE VIEW {TABLE} AS SELECT 1",), + } + for label, statements in variants.items(): + with self.subTest(label), closing(sqlite3.connect(":memory:")) as conn: + for sql in statements or (): + conn.execute(sql) + error = image_categories_schema_error(conn) + if statements is None: + self.assertIsNone(error) + self.assertIsNotNone( + image_categories_schema_error(conn, allow_missing=False) + ) + continue + self.assertIsNotNone(error) + self.assertIsNotNone( + image_categories_schema_error(conn, check_values=False) + ) + with self.assertRaises(sqlite3.DatabaseError): + migrate_image_categories(conn) + with closing(sqlite3.connect(":memory:")) as conn: + migrate_image_categories(conn) + migrate_image_categories(conn) # idempotent + self.assertIsNone(image_categories_schema_error(conn, allow_missing=False)) + + def test_values_refused_but_startup_checks_structure_only(self): + bad = { + "blank": "", + "101 characters": "x" * 101, + "control character": "a\tb", + "nul": "a\x00b", + "leading space": " Landscapes", + "trailing space": "Landscapes ", + } + for label, value in bad.items(): + with self.subTest(label), closing(sqlite3.connect(":memory:")) as conn: + migrate_image_categories(conn) + conn.execute( + f"INSERT INTO {TABLE} VALUES ('r', 'a.png', ?, 'now')", (value,) + ) + self.assertIsNotNone(image_categories_schema_error(conn)) + self.assertIsNone( + image_categories_schema_error(conn, check_values=False) + ) + migrate_image_categories(conn) # startup still opens it + with closing(sqlite3.connect(":memory:")) as conn: + migrate_image_categories(conn) + conn.execute( + f"INSERT INTO {TABLE} VALUES ('r', ?, 'Scans', 'now')", ("x" * 5000,) + ) + self.assertIsNotNone(image_categories_schema_error(conn)) + conn.execute(f"DELETE FROM {TABLE}") + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', x'00ff', 'Scans', 'now')") + self.assertIsNotNone(image_categories_schema_error(conn)) + conn.execute(f"DELETE FROM {TABLE}") + for value in ("x" * 100, "Scans", "Café · 風景", "a b"): + conn.execute( + f"INSERT OR REPLACE INTO {TABLE} VALUES ('r', 'a.png', ?, 'now')", + (value,), + ) + self.assertIsNone(image_categories_schema_error(conn), value) + + def test_restore_check(self): + PromptDatabase(self.path).close_all() + PromptModel.close_all_instances() + ok, reason = PromptModel.verify_database_file( + self.path, require_library_state=True + ) + self.assertTrue(ok, reason) + with closing(sqlite3.connect(self.path)) as conn: + conn.execute(f"DROP TABLE {TABLE}") + conn.commit() + # A backup from before #332 is accepted ... + ok, reason = PromptModel.verify_database_file(self.path) + self.assertTrue(ok, reason) + # ... but a current database must have the table. + ok, reason = PromptModel.verify_database_file( + self.path, require_library_state=True + ) + self.assertFalse(ok) + self.assertIn("image categories", reason.lower()) + with closing(sqlite3.connect(self.path)) as conn: + conn.execute(_DDL) + conn.execute(_INDEX_DDL) + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', 'a.png', ' x', 'now')") + conn.commit() + ok, reason = PromptModel.verify_database_file(self.path) + self.assertFalse(ok) + self.assertIn("image categories", reason.lower()) + + +class TestImageCategoryStore(_TempDir): + def setUp(self): + super().setUp() + self.db = PromptDatabase(self.path) + + def tearDown(self): + self.db.close_all() + super().tearDown() + + def rows(self): + return ( + self.db.model.get_connection() + .execute( + f"SELECT root_key, relative_path, category, updated_at FROM {TABLE} " + "ORDER BY root_key, relative_path" + ) + .fetchall() + ) + + def test_set_counts_updated_and_unchanged(self): + keys = [("/out", "a.png"), ("/other", "a.png"), ("/out", "a.png")] + self.assertEqual( + self.db.set_image_categories(keys, "Portraits"), + {"updated": 2, "unchanged": 0}, + ) + stamps = {(r[0], r[1]): r[3] for r in self.rows()} + self.assertEqual( + self.db.set_image_categories( + [("/out", "a.png"), ("/out", "b.png")], "Portraits" + ), + {"updated": 1, "unchanged": 1}, + ) + # An unchanged file keeps its updated_at. + self.assertEqual( + {(r[0], r[1]): r[3] for r in self.rows()}[("/out", "a.png")], + stamps[("/out", "a.png")], + ) + self.assertEqual( + self.db.set_image_categories([("/out", "a.png")], "Scans"), + {"updated": 1, "unchanged": 0}, + ) + self.assertEqual(self.db.image_category("/out", "a.png"), "Scans") + self.assertEqual(self.db.image_category("/other", "a.png"), "Portraits") + self.assertIsNone(self.db.image_category("/out", "missing.png")) + self.assertEqual( + self.db.own_categories_of(["/out", "/out", "/none"]), + {("/out", "a.png"): "Scans", ("/out", "b.png"): "Portraits"}, + ) + + def test_clear_counts_removed_rows(self): + self.db.set_image_categories([("/out", "a.png"), ("/b", "c.png")], "Scans") + self.assertEqual( + self.db.clear_image_categories( + [("/out", "a.png"), ("/b", "c.png"), ("/out", "gone.png")] + ), + 2, + ) + self.assertEqual(self.rows(), []) + self.assertEqual(self.db.clear_image_categories([("/out", "a.png")]), 0) + + def test_a_failed_write_writes_nothing(self): + self.db.set_image_categories([("/out", "a.png")], "Scans") + keys = [("/out", "b.png"), ("/out", None)] # NOT NULL fails the second + with self.assertRaises(sqlite3.IntegrityError): + self.db.set_image_categories(keys, "Portraits") + self.assertEqual([r[:3] for r in self.rows()], [("/out", "a.png", "Scans")]) + + def test_store_never_touches_prompt_categories(self): + conn = self.db.model.get_connection() + with conn: + conn.execute( + "INSERT INTO prompts (text, category, hash) VALUES ('p', 'L', 'h')" + ) + before = category_checksum(conn) + seen = set() + conn.set_authorizer( + lambda action, arg1, *_: seen.add(arg1) or sqlite3.SQLITE_OK + ) + try: + self.db.set_image_categories([("/out", "a.png")], "Scans") + self.db.own_categories_of(["/out"]) + self.db.image_category("/out", "a.png") + self.db.clear_image_categories([("/out", "a.png")]) + finally: + conn.set_authorizer(None) + self.assertFalse(seen & {"prompts", "prompt_pairs", "generated_images"}, seen) + self.assertEqual(category_checksum(conn), before) + + def test_copy_keeps_kept_copy_category(self): + self.db.set_image_categories([("/out", "removed.png")], "Scans") + self.db.set_image_categories([("/out", "kept.png")], "Portraits") + with self.db._pair_transaction() as conn: + inserted = copy_image_category( + conn, ("/out", "removed.png"), ("/out", "kept.png") + ) + self.assertEqual(inserted, 0) + self.assertEqual(self.db.image_category("/out", "kept.png"), "Portraits") + with self.db._pair_transaction() as conn: + inserted = copy_image_category( + conn, ("/out", "removed.png"), ("/out", "new.png") + ) + self.assertEqual(inserted, 1) + self.assertEqual(self.db.image_category("/out", "new.png"), "Scans") + # The removed copy's row is never deleted. + self.assertEqual(self.db.image_category("/out", "removed.png"), "Scans") + with self.db._pair_transaction() as conn: + self.assertEqual( + copy_image_category(conn, ("/out", "none.png"), ("/out", "x.png")), 0 + ) + self.assertEqual( + copy_image_category( + conn, ("/out", "removed.png"), ("/out", "removed.png") + ), + 0, + ) + + def test_statements_never_scan_the_table(self): + conn = self.db.model.get_connection() + for sql in STATEMENTS: + with self.subTest(sql): + params = (None,) * sql.count("?") + plan = " ".join( + row[-1] for row in conn.execute(f"EXPLAIN QUERY PLAN {sql}", params) + ) + self.assertNotRegex(plan, rf"SCAN {TABLE}\b", plan) + plan = " ".join( + row[-1] + for row in conn.execute( + f"EXPLAIN QUERY PLAN SELECT rowid FROM {TABLE} WHERE category=? " + "ORDER BY rowid", + ("x",), + ) + ) + self.assertIn(INDEX, plan) + self.assertIn(COPY_SQL, STATEMENTS) + + +class TestImageCategoryKey(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.base = Path(os.path.realpath(self._tmp.name)) + self.outer = self.base / "output" + self.inner = self.outer / "nested" + (self.inner / "sub").mkdir(parents=True) + + def tearDown(self): + self._tmp.cleanup() + + def test_equals_the_favorite_key_without_nesting(self): + from py.api.images import _favorite_relative_key, _favorite_root_key + + key = image_category_key([self.outer], self.outer / "sub" / "a.png") + self.assertEqual( + key, (_favorite_root_key(self.outer), _favorite_relative_key("sub/a.png")) + ) + + def test_deepest_source_wins_in_any_order(self): + from py.api.images import _favorite_root_key + + path = self.inner / "sub" / "a.png" + for sources in ([self.outer, self.inner], [self.inner, self.outer]): + with self.subTest(sources=sources): + self.assertEqual( + image_category_key(sources, path), + (_favorite_root_key(self.inner), "sub/a.png"), + ) + # A file of the outer source only is keyed by it. + self.assertEqual( + image_category_key([self.outer, self.inner], self.outer / "b.png"), + (_favorite_root_key(self.outer), "b.png"), + ) + + def test_outside_every_source_has_no_key(self): + self.assertIsNone(image_category_key([self.inner], self.outer / "a.png")) + self.assertIsNone(image_category_key([], self.outer / "a.png")) + # A sibling sharing a name prefix is not inside. + self.assertIsNone( + image_category_key([self.outer], self.base / "output2" / "a.png") + ) + self.assertIsNone(image_category_key([self.outer], self.outer)) + + def test_symlinked_subfolder_is_not_resolved(self): + target = self.base / "elsewhere" + target.mkdir() + link = self.outer / "linked" + try: + os.symlink(target, link, target_is_directory=True) + except (OSError, NotImplementedError): + self.skipTest("symlinks are not available") + key = image_category_key([self.outer], link / "a.png") + self.assertEqual(key[1], "linked/a.png") + + def test_windows_folds_case_of_the_relative_key(self): + with patch("py.api.images._path_flavor", ntpath): + upper = image_category_key([self.outer], self.outer / "Sub" / "A.PNG") + lower = image_category_key([self.outer], self.outer / "sub" / "a.png") + self.assertEqual(upper[1], "sub/a.png") + self.assertEqual(upper, lower) + + def test_posix_keeps_case_distinct(self): + if os.path.normcase("A") == "a": + self.skipTest("this OS folds case") + upper = image_category_key([self.outer], self.outer / "A.png") + lower = image_category_key([self.outer], self.outer / "a.png") + self.assertNotEqual(upper, lower) + + +if __name__ == "__main__": + unittest.main()