diff --git a/database/image_collections.py b/database/image_collections.py new file mode 100644 index 00000000..674b3937 --- /dev/null +++ b/database/image_collections.py @@ -0,0 +1,231 @@ +"""Collections (#331): named, hand-picked groups of image files. + +A member names its file exactly as a favorite does (``image_favorites.py``): +the source folder's real, case-normalized path (``root_key``) and the path +below it (``relative_path``), both server-side only. A member whose file is +moved, renamed or deleted outside PromptManager is kept and simply not listed; +nothing follows a file to a new place. + +Names are trimmed with each run of whitespace collapsed to one space, 1 to 100 +characters without control characters; ``name_key`` (NFC of the case-folded +NFC name) is UNIQUE, so names are unique ignoring case and Unicode form. + +The two tables install like ``image_favorites`` and ``saved_views``: optional +in old backups, required once a database is current, created empty with no +backfill. Startup checks only the structure, so a value saved under older +rules never stops the library from opening; restore checks the values too. +""" + +import re +import sqlite3 +import unicodedata +from typing import Optional + +from .image_favorites import MAX_KEY_LENGTH + +COLLECTIONS_TABLE = "image_collections" +ITEMS_TABLE = "image_collection_items" +FILE_INDEX = "idx_image_collection_items_file" +TABLES = (COLLECTIONS_TABLE, ITEMS_TABLE) +MAX_NAME_LENGTH = 100 +_TIMESTAMP_BYTES = 64 +_CONTROL = re.compile(r"[\x00-\x1f\x7f]") +_WHITESPACE = re.compile(r"\s+") + +_TABLE_DDL = { + COLLECTIONS_TABLE: ( + f"CREATE TABLE {COLLECTIONS_TABLE} (id INTEGER PRIMARY KEY AUTOINCREMENT, " + "name TEXT NOT NULL, name_key TEXT NOT NULL UNIQUE, " + "created_at TEXT NOT NULL)" + ), + ITEMS_TABLE: ( + f"CREATE TABLE {ITEMS_TABLE} (collection_id INTEGER NOT NULL REFERENCES " + f"{COLLECTIONS_TABLE}(id) ON DELETE CASCADE, root_key TEXT NOT NULL, " + "relative_path TEXT NOT NULL, added_at TEXT NOT NULL, " + "PRIMARY KEY (collection_id, root_key, relative_path))" + ), +} +_INDEX_DDL = f"CREATE INDEX {FILE_INDEX} ON {ITEMS_TABLE}(root_key, relative_path)" +# (name, declared type, not null, default, primary key position) +_COLUMNS = { + COLLECTIONS_TABLE: ( + ("id", "INTEGER", 0, None, 1), + ("name", "TEXT", 1, None, 0), + ("name_key", "TEXT", 1, None, 0), + ("created_at", "TEXT", 1, None, 0), + ), + ITEMS_TABLE: ( + ("collection_id", "INTEGER", 1, None, 1), + ("root_key", "TEXT", 1, None, 2), + ("relative_path", "TEXT", 1, None, 3), + ("added_at", "TEXT", 1, None, 0), + ), +} +# (column, parent table, parent column, on delete) +_FOREIGN_KEYS = { + COLLECTIONS_TABLE: set(), + ITEMS_TABLE: {("collection_id", COLLECTIONS_TABLE, "id", "CASCADE")}, +} + + +class CollectionInvalid(ValueError): + """A name, id or target breaks the write rules (a 400).""" + + +def normalize_collection_name(name: object) -> str: + """The stored spelling of *name*: trimmed, whitespace collapsed, checked. + + Raises ``CollectionInvalid`` unless the result is 1 to 100 characters + without control characters. + """ + if not isinstance(name, str): + raise CollectionInvalid("Name must be a string") + collapsed = _WHITESPACE.sub(" ", name).strip() + if not 1 <= len(collapsed) <= MAX_NAME_LENGTH: + raise CollectionInvalid(f"Name must be 1 to {MAX_NAME_LENGTH} characters") + if _CONTROL.search(collapsed): + raise CollectionInvalid("Name contains control characters") + return collapsed + + +def collection_name_key(name: str) -> str: + """The uniqueness key of a name: NFC of the case-folded NFC spelling.""" + return unicodedata.normalize("NFC", unicodedata.normalize("NFC", name).casefold()) + + +def _compact(sql: str) -> str: + return re.sub(r"\s+", "", sql).lower() + + +def _expected_catalog() -> set: + objects = {("table", name, _compact(sql)) for name, sql in _TABLE_DDL.items()} + objects.add(("index", FILE_INDEX, _compact(_INDEX_DDL))) + objects |= {("index", f"sqlite_autoindex_{table}_1", "") for table in TABLES} + return objects + + +def _catalog(conn: sqlite3.Connection) -> list: + # sqlite_sequence (AUTOINCREMENT) is shared with other tables and ignored. + return conn.execute( + "SELECT type,name,sql FROM sqlite_master " + "WHERE name IN (?,?,?) OR tbl_name IN (?,?)", + (*TABLES, FILE_INDEX, *TABLES), + ).fetchall() + + +def _structure_error(conn: sqlite3.Connection) -> Optional[str]: + for table, expected in _COLUMNS.items(): + columns = tuple( + (row[1], row[2].upper(), row[3], row[4], row[5]) + for row in conn.execute(f"PRAGMA table_info({table})") + ) + if columns != expected: + return f"Image collections table {table} has invalid columns" + foreign_keys = { + (row[3], row[2], row[4], row[6]) + for row in conn.execute(f"PRAGMA foreign_key_list({table})") + } + if foreign_keys != _FOREIGN_KEYS[table]: + return f"Image collections table {table} has invalid foreign keys" + return None + + +def _decoded_too_long(conn, sql: str, limit: int) -> bool: + """Whether a byte-bounded candidate row decodes to over *limit* characters.""" + for row in conn.execute(sql): + for raw in row: + if len(bytes(raw).decode("utf-8", "replace")) > limit: + return True + return False + + +def _value_error(conn: sqlite3.Connection) -> Optional[str]: + # Only this module's constants are interpolated into the SQL below. Bytes + # bound characters from above (UTF-8 is at most 4 bytes each); the few + # candidates over the character bound are measured decoded in Python. + if conn.execute( + f"SELECT 1 FROM {COLLECTIONS_TABLE} WHERE " # nosec B608 + "typeof(id) != 'integer' OR typeof(name) != 'text' " + "OR typeof(name_key) != 'text' OR typeof(created_at) != 'text' " + f"OR length(CAST(name AS BLOB)) > {MAX_NAME_LENGTH * 4} " + f"OR length(CAST(created_at AS BLOB)) > {_TIMESTAMP_BYTES} LIMIT 1" + ).fetchone(): + return "Image collections contains invalid values" + for name, name_key in conn.execute( + f"SELECT name, name_key FROM {COLLECTIONS_TABLE}" # nosec B608 + ): + if ( + not 1 <= len(name) <= MAX_NAME_LENGTH + or _CONTROL.search(name) + or name != name.strip() + or name_key != collection_name_key(name) + ): + return "Image collections contains invalid names" + limit = MAX_KEY_LENGTH * 4 + if conn.execute( + f"SELECT 1 FROM {ITEMS_TABLE} WHERE " # nosec B608 + "typeof(collection_id) != 'integer' OR typeof(root_key) != 'text' " + "OR typeof(relative_path) != 'text' OR typeof(added_at) != 'text' " + f"OR length(CAST(root_key AS BLOB)) > {limit} " + f"OR length(CAST(relative_path AS BLOB)) > {limit} " + f"OR length(CAST(added_at AS BLOB)) > {_TIMESTAMP_BYTES} LIMIT 1" + ).fetchone() or _decoded_too_long( + conn, + f"SELECT CAST(root_key AS BLOB), CAST(relative_path AS BLOB) " # nosec B608 + f"FROM {ITEMS_TABLE} WHERE length(CAST(root_key AS BLOB)) > {MAX_KEY_LENGTH} " + f"OR length(CAST(relative_path AS BLOB)) > {MAX_KEY_LENGTH}", + MAX_KEY_LENGTH, + ): + return "Image collections contains invalid members" + if conn.execute(f"PRAGMA foreign_key_check({ITEMS_TABLE})").fetchone(): + return "Image collections contains members of a missing collection" + return None + + +def image_collections_schema_error( + conn: sqlite3.Connection, + *, + allow_missing: bool = True, + check_values: bool = True, +) -> Optional[str]: + """Reject partial or unsupported collections schemas and invalid rows. + + Reads only: nothing is ever repaired. A database without either table is + acceptable for backups that predate them unless ``allow_missing`` is + false. ``check_values=False`` (live startup) checks the structure only. + """ + objects = _catalog(conn) + if not objects: + return None if allow_missing else "Missing image collections schema" + present = {name for kind, name, _sql in objects if kind == "table"} + if present != set(TABLES): + return "Incomplete image collections schema" + actual = {(kind, name, _compact(sql or "")) for kind, name, sql in objects} + if actual != _expected_catalog(): + return "Image collections has an invalid schema definition" + error = _structure_error(conn) + if error or not check_values: + return error + return _value_error(conn) + + +def migrate_image_collections(conn: sqlite3.Connection) -> None: + """Install both tables atomically; validate before and after creating them.""" + conn.execute("SAVEPOINT image_collections_schema") + try: + error = image_collections_schema_error(conn, check_values=False) + if error: + raise sqlite3.DatabaseError(error) + for sql in _TABLE_DDL.values(): + conn.execute(sql.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1)) + conn.execute(_INDEX_DDL.replace("INDEX", "INDEX IF NOT EXISTS", 1)) + error = image_collections_schema_error( + conn, allow_missing=False, check_values=False + ) + if error: + raise sqlite3.DatabaseError(error) + conn.execute("RELEASE image_collections_schema") + except BaseException: + conn.execute("ROLLBACK TO image_collections_schema") + conn.execute("RELEASE image_collections_schema") + raise diff --git a/database/models.py b/database/models.py index dad9cf07..fe30685e 100644 --- a/database/models.py +++ b/database/models.py @@ -22,6 +22,10 @@ from .managed_connection import ( from . import restore_stage from .library_schema import library_schema_error, migrate_library_state from .image_favorites import image_favorites_schema_error, migrate_image_favorites +from .image_collections import ( + image_collections_schema_error, + migrate_image_collections, +) from .annotation_schema import annotation_schema_error, migrate_annotation_state from .autotag_schema import ( autotag_schema_error, @@ -337,6 +341,7 @@ class PromptModel: migrate_library_state(conn) migrate_image_favorites(conn) migrate_saved_views(conn) + migrate_image_collections(conn) # A review group without its annotation group is refused before the # annotation install could make it look valid (catalog only; the rows # are validated once, in migrate_review_state). @@ -1156,7 +1161,7 @@ class PromptModel: 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 + collections, 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 @@ -1225,6 +1230,11 @@ class PromptModel: ) if views_error: return False, views_error + collections_error = image_collections_schema_error( + conn, allow_missing=not require_library_state + ) + if collections_error: + return False, collections_error annotation_error = annotation_schema_error( conn, allow_missing=not require_library_state ) diff --git a/database/restore_stage.py b/database/restore_stage.py index 75d8883a..704435bf 100644 --- a/database/restore_stage.py +++ b/database/restore_stage.py @@ -59,6 +59,7 @@ SUMMARY_KEYS = ( "image_links", "runs", "saved_views", + "collections", ) @@ -144,7 +145,8 @@ def library_summary(conn: sqlite3.Connection) -> Dict[str, int]: ``positives`` are the Library's prompts (``library_counts``'s ``positive_prompts``); ``pairs`` the active saved pairs and ``archived_pairs`` the archived ones; ``image_links`` the indexed image - rows; ``runs`` the recorded executions. + rows; ``runs`` the recorded executions; ``collections`` the image + collections (#331). """ from .library import _ELIGIBLE @@ -169,6 +171,7 @@ def library_summary(conn: sqlite3.Connection) -> Dict[str, int]: "image_links": _count(conn, "generated_images"), "runs": _count(conn, "execution_runs"), "saved_views": _count(conn, "saved_views"), + "collections": _count(conn, "image_collections"), } @@ -180,6 +183,7 @@ _FINGERPRINT_SQL = ( ("prompt_pairs", "SELECT MAX(id), MAX(updated_at) FROM prompt_pairs"), ("execution_runs", "SELECT COUNT(*), MAX(rowid) FROM execution_runs"), ("saved_views", "SELECT MAX(rowid) FROM saved_views"), + ("image_collections", "SELECT MAX(rowid) FROM image_collections"), ) @@ -190,9 +194,9 @@ def live_fingerprint(conn: sqlite3.Connection): tables a restore replaces, or when a prompt, pair, run or saved view is the newest edited (counts plus the newest row and edit of each table); the apply refuses with ``plan-changed`` then. It does not see every - edit (a tag or favorite change on an older row); those are kept in the - restore's safety backup. A read transaction the caller already has is - used and left open; one opened here is ended here. + edit (a tag, favorite or collection membership change on an older row); + those are kept in the restore's safety backup. A read transaction the + caller already has is used and left open; one opened here is ended here. """ import json diff --git a/tests/js/workspace-system-backups.test.js b/tests/js/workspace-system-backups.test.js index 05899351..3a1454ae 100644 --- a/tests/js/workspace-system-backups.test.js +++ b/tests/js/workspace-system-backups.test.js @@ -63,9 +63,9 @@ test('status line shows only server-reported backups (D103)', async () => { const PREVIEW = Object.freeze({ success: true, token: 'r'.repeat(43), expires_at: '2026-10-25T12:15:00+00:00', - summary: { prompts: 1234, positives: 900, pairs: 3, archived_pairs: 0, image_links: 56, runs: 1, saved_views: 2 }, + summary: { prompts: 1234, positives: 900, pairs: 3, archived_pairs: 0, image_links: 56, runs: 1, saved_views: 2, collections: 1 }, schema: { migrated: true, library_state: true }, integrity: 'ok', documents: { repaired: 0, cleared: 0, images: 0 }, - live: { prompts: 1, positives: 1, pairs: 0, archived_pairs: 0, image_links: 0, runs: 0, saved_views: 0 }, size: 4096, + live: { prompts: 1, positives: 1, pairs: 0, archived_pairs: 0, image_links: 0, runs: 0, saved_views: 0, collections: 0 }, size: 4096, safety_backup: { destination_label: 'prompts.db.backup_' }, }); diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 2437220b..8e08b588 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -614,6 +614,50 @@ class TestRestoreFromFile(BackupTestCase): ) assert_closed(self, str(backups[0])) + def test_backup_and_restore_keep_collections(self): + with self.db.model.get_connection() as conn: + conn.execute( + "INSERT INTO image_collections (id, name, name_key, created_at) " + "VALUES (4, 'LoRa testing', 'lora testing', '2026-10-27T00:00:00Z')" + ) + conn.execute( + "INSERT INTO image_collection_items VALUES " + "(4, 'root', 'grid.png', '2026-10-27T00:00:00Z')" + ) + with tempfile.TemporaryDirectory() as directory: + # The rows still sit in the WAL; a backup must carry them. + backup = os.path.join(directory, "backup.db") + self.assertTrue(self.db.model.backup_database(backup)) + self.assertEqual( + PromptModel.verify_database_file(backup, require_library_state=True), + (True, "ok"), + ) + with self.db.model.get_connection() as conn: + conn.execute("DELETE FROM image_collections") + + self.db.model.restore_from_file(backup) + + conn = self.db.model.get_connection() + self.assertEqual( + [ + tuple(r) + for r in conn.execute("SELECT id, name FROM image_collections") + ], + [(4, "LoRa testing")], + ) + self.assertEqual( + [ + tuple(r) + for r in conn.execute( + "SELECT collection_id, root_key, relative_path " + "FROM image_collection_items" + ) + ], + [(4, "root", "grid.png")], + ) + PromptModel.close_all_instances() + assert_closed(self, directory) + def test_restore_without_a_live_database_returns_empty_backup_path(self): src = self._make_source_db(["fresh"]) self.db.close_all() diff --git a/tests/test_image_collections_schema.py b/tests/test_image_collections_schema.py new file mode 100644 index 00000000..fb06ba2d --- /dev/null +++ b/tests/test_image_collections_schema.py @@ -0,0 +1,485 @@ +"""Collections schema (#331): install, validation and restore safety.""" + +import re +import sqlite3 +import tempfile +import unittest +import unicodedata +from contextlib import closing +from pathlib import Path +from unittest.mock import patch + +try: + from tests.open_handles import assert_closed + from tests.test_annotation_schema import LEGACY_TABLES, table_rows +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + from test_annotation_schema import LEGACY_TABLES, table_rows + +from database.models import PromptModel + +try: # guarded so a missing module fails assertions, not test loading + from database import image_collections + from database.image_collections import ( + image_collections_schema_error, + migrate_image_collections, + ) +except ImportError: + image_collections = None + image_collections_schema_error = None + migrate_image_collections = None + +COLLECTIONS = "image_collections" +ITEMS = "image_collection_items" +INDEX = "idx_image_collection_items_file" +TABLES = (COLLECTIONS, ITEMS) +DDL = { + COLLECTIONS: ( + f"CREATE TABLE {COLLECTIONS} (id INTEGER PRIMARY KEY AUTOINCREMENT, " + "name TEXT NOT NULL, name_key TEXT NOT NULL UNIQUE, " + "created_at TEXT NOT NULL)" + ), + ITEMS: ( + f"CREATE TABLE {ITEMS} (collection_id INTEGER NOT NULL REFERENCES " + f"{COLLECTIONS}(id) ON DELETE CASCADE, root_key TEXT NOT NULL, " + "relative_path TEXT NOT NULL, added_at TEXT NOT NULL, " + "PRIMARY KEY (collection_id, root_key, relative_path))" + ), +} +INDEX_DDL = f"CREATE INDEX {INDEX} ON {ITEMS}(root_key, relative_path)" +NOW = "2026-10-27T00:00:00.000000+00:00" +IMAGES = 23_000 + + +def _compact(sql): + return re.sub(r"\s+", "", sql or "").lower() + + +def expected_catalog(): + return { + ("table", COLLECTIONS, _compact(DDL[COLLECTIONS])), + ("table", ITEMS, _compact(DDL[ITEMS])), + ("index", INDEX, _compact(INDEX_DDL)), + ("index", f"sqlite_autoindex_{COLLECTIONS}_1", ""), + ("index", f"sqlite_autoindex_{ITEMS}_1", ""), + } + + +def catalog(conn): + return { + (kind, name, _compact(sql)) + for kind, name, sql in conn.execute( + "SELECT type,name,sql FROM sqlite_master " + "WHERE name IN (?,?,?) OR tbl_name IN (?,?)", + (*TABLES, INDEX, *TABLES), + ) + } + + +def key(name): + return unicodedata.normalize("NFC", unicodedata.normalize("NFC", name).casefold()) + + +def schema(conn, exclude=()): + return sorted( + tuple(row) + for row in conn.execute( + "SELECT type,name,tbl_name,sql FROM sqlite_master ORDER BY name" + ) + if row[1] not in exclude and row[2] not in exclude + ) + + +def row_counts(conn): + tables = [ + name + for (name,) in conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name" + ) + ] + return { + table: conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] # nosec + for table in tables + } + + +class ImageCollectionsSchemaTests(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( + image_collections_schema_error, + "database.image_collections is not implemented", + ) + + def closed_current(self, name="current.db", *statements): + """A closed current database, altered afterwards by *statements*.""" + path = str(Path(self.temp.name) / name) + model = self.model(path) + model.close_all() + PromptModel.reset_schema_cache() + with closing(sqlite3.connect(path)) as raw, raw: + raw.execute("PRAGMA foreign_keys=OFF") + for statement in statements: + if isinstance(statement, tuple): + raw.execute(*statement) + else: + raw.execute(statement) + return path + + @staticmethod + def drop_group(path): + with closing(sqlite3.connect(path)) as raw, raw: + raw.execute(f"DROP TABLE {ITEMS}") + raw.execute(f"DROP TABLE {COLLECTIONS}") + + def test_migration_creates_tables_and_index(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + self.assertEqual(catalog(conn), expected_catalog()) + self.assertIsNone(image_collections_schema_error(conn, allow_missing=False)) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {COLLECTIONS}").fetchone()[0], 0 + ) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {ITEMS}").fetchone()[0], 0 + ) + self.assertEqual( + PromptModel.verify_database_file(self.path, require_library_state=True), + (True, "ok"), + ) + + def test_migration_is_idempotent(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + conn.execute( + f"INSERT INTO {COLLECTIONS} (name, name_key, created_at) " + "VALUES ('Kept', 'kept', ?)", + (NOW,), + ) + conn.execute(f"INSERT INTO {ITEMS} VALUES (1, 'root', 'a.png', ?)", (NOW,)) + before = schema(conn) + model.close_all() + for _ in range(2): + PromptModel.reset_schema_cache() + model = self.model() + with model.get_connection() as conn: + migrate_image_collections(conn) + self.assertEqual(schema(conn), before) + self.assertEqual( + [tuple(r) for r in conn.execute(f"SELECT * FROM {ITEMS}")], + [(1, "root", "a.png", NOW)], + ) + model.close_all() + + def test_migration_reads_only_its_own_tables(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + conn.execute("INSERT INTO prompts(id,text,hash) VALUES(1,'positive','p')") + conn.executemany( + "INSERT INTO generated_images(prompt_id,image_path,filename," + "file_path) VALUES(1,?,?,?)", + ((f"/out/{n}.png", f"{n}.png", f"/out/{n}.png") for n in range(IMAGES)), + ) + conn.execute("INSERT INTO image_favorites VALUES('root','a.png',?)", (NOW,)) + model.close_all() + self.drop_group(self.path) + with closing(sqlite3.connect(self.path)) as raw: + raw.execute("PRAGMA foreign_keys=ON") + counts_before = row_counts(raw) + schema_before = schema(raw) + legacy_before = table_rows(raw, LEGACY_TABLES) + self.assertEqual(counts_before["generated_images"], IMAGES) + touched = [] + + def record(action, arg1, arg2, _db, _trigger): + touched.append((action, arg1, arg2)) + return sqlite3.SQLITE_OK + + raw.set_authorizer(record) + try: + migrate_image_collections(raw) + finally: + raw.set_authorizer(None) + raw.commit() + reads = { + arg1 for action, arg1, _arg2 in touched if action == sqlite3.SQLITE_READ + } + self.assertTrue(reads) + self.assertLessEqual(reads, {"sqlite_master", COLLECTIONS, ITEMS}) + writes = { + arg1 + for action, arg1, _arg2 in touched + if action + in (sqlite3.SQLITE_INSERT, sqlite3.SQLITE_UPDATE, sqlite3.SQLITE_DELETE) + } + self.assertLessEqual(writes, {"sqlite_master"}) + self.assertEqual(catalog(raw), expected_catalog()) + counts_after = row_counts(raw) + self.assertEqual(counts_after.pop(COLLECTIONS), 0) + self.assertEqual(counts_after.pop(ITEMS), 0) + self.assertEqual(counts_after, counts_before) + self.assertEqual(schema(raw, exclude=TABLES), schema_before) + self.assertEqual(table_rows(raw, LEGACY_TABLES), legacy_before) + # A second run is a no-op. + after_first = schema(raw) + changes = raw.total_changes + migrate_image_collections(raw) + raw.commit() + self.assertEqual(schema(raw), after_first) + self.assertEqual(raw.total_changes, changes) + + def test_migration_refuses_a_foreign_object_of_the_same_name(self): + self.require_module() + for statements in ( + (f"CREATE TABLE {COLLECTIONS} (x)",), + (f"CREATE VIEW {COLLECTIONS} AS SELECT 1",), + (f"CREATE INDEX {INDEX} ON prompts(text)",), + ): + with self.subTest(statements=statements): + path = str(Path(self.temp.name) / f"foreign{len(self.models)}.db") + self.model(path).close_all() + PromptModel.reset_schema_cache() + self.drop_group(path) + with closing(sqlite3.connect(path)) as raw, raw: + for statement in statements: + raw.execute(statement) + with closing(sqlite3.connect(path)) as raw: + before = schema(raw) + self.assertIsNotNone(image_collections_schema_error(raw)) + with self.assertRaises(sqlite3.DatabaseError): + self.model(path).get_connection() + PromptModel.close_all_instances() + PromptModel.reset_schema_cache() + with closing(sqlite3.connect(path)) as raw: + self.assertEqual(schema(raw), before) + self.assertIsNone( + raw.execute( + "SELECT 1 FROM sqlite_master WHERE name=?", (ITEMS,) + ).fetchone() + ) + self.assertFalse(PromptModel.verify_database_file(path)[0]) + + def test_migration_rolls_back_on_error(self): + self.require_module() + self.model().close_all() + PromptModel.reset_schema_cache() + self.drop_group(self.path) + with closing(sqlite3.connect(self.path)) as raw: + with patch.object( + image_collections, + "image_collections_schema_error", + side_effect=[None, "Image collections went wrong"], + ): + with self.assertRaisesRegex(sqlite3.DatabaseError, "went wrong"): + migrate_image_collections(raw) + raw.commit() + self.assertEqual(catalog(raw), set()) + + def test_startup_accepts_a_value_saved_under_older_rules(self): + self.require_module() + path = self.closed_current( + "older.db", + ( + f"INSERT INTO {COLLECTIONS} (name, name_key, created_at) " + "VALUES (?, 'not the key', ?)", + (" " + "x" * 150 + " ", NOW), + ), + ) + model = self.model(path) + with model.get_connection() as conn: + self.assertIsNone(image_collections_schema_error(conn, check_values=False)) + self.assertIsNotNone(image_collections_schema_error(conn)) + self.assertEqual( + conn.execute(f"SELECT COUNT(*) FROM {COLLECTIONS}").fetchone()[0], 1 + ) + model.close_all() + self.assertFalse( + PromptModel.verify_database_file(path, require_library_state=True)[0] + ) + + def test_restore_accepts_an_old_backup_without_collections(self): + self.require_module() + path = self.closed_current("old.db") + self.drop_group(path) + with closing(sqlite3.connect(path)) as raw: + self.assertIsNone(image_collections_schema_error(raw)) + self.assertEqual( + image_collections_schema_error(raw, allow_missing=False), + "Missing image collections schema", + ) + self.assertEqual(PromptModel.verify_database_file(path), (True, "ok")) + model = self.model(path) + with model.get_connection() as conn: + self.assertEqual(catalog(conn), expected_catalog()) + model.close_all() + self.assertEqual( + PromptModel.verify_database_file(path, require_library_state=True), + (True, "ok"), + ) + + def test_restore_requires_collections_on_a_current_stage(self): + self.require_module() + path = self.closed_current("stage.db") + self.drop_group(path) + self.assertEqual( + PromptModel.verify_database_file(path, require_library_state=True), + (False, "Missing image collections schema"), + ) + # One table without the other is never acceptable. + partial = self.closed_current("partial.db", f"DROP TABLE {ITEMS}") + self.assertFalse(PromptModel.verify_database_file(partial)[0]) + + def assert_refused(self, name, *statements): + self.require_module() + path = self.closed_current( + f"{name}.db", + ( + f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) " + "VALUES (1, 'Good', 'good', ?)", + (NOW,), + ), + (f"INSERT INTO {ITEMS} VALUES (1, 'root', 'a.png', ?)", (NOW,)), + *statements, + ) + with closing(sqlite3.connect(path)) as raw: + self.assertIsNone(image_collections_schema_error(raw, check_values=False)) + error = image_collections_schema_error(raw) + self.assertIsNotNone(error) + self.assertIn("collections", error.lower()) + ok, reason = PromptModel.verify_database_file(path) + self.assertFalse(ok) + self.assertEqual(reason, error) + self.assertFalse( + PromptModel.verify_database_file(path, require_library_state=True)[0] + ) + + def test_a_valid_current_database_is_accepted(self): + self.require_module() + path = self.closed_current( + "valid.db", + ( + f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) " + "VALUES (1, ?, ?, ?)", + ("LoRa testing", key("LoRa testing"), NOW), + ), + (f"INSERT INTO {ITEMS} VALUES (1, 'root', ?, ?)", ("é" * 4096, NOW)), + ) + self.assertEqual( + PromptModel.verify_database_file(path, require_library_state=True), + (True, "ok"), + ) + + def test_restore_refuses_a_non_text_name(self): + self.assert_refused("int-name", f"UPDATE {COLLECTIONS} SET name = 5") + + def test_restore_refuses_an_empty_name(self): + self.assert_refused("empty", f"UPDATE {COLLECTIONS} SET name = '', name_key=''") + + def test_restore_refuses_a_name_over_100_characters(self): + name = "é" * 101 + self.assert_refused( + "long", + (f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", (name, key(name))), + ) + + def test_restore_refuses_a_control_character(self): + name = "a\x07b" + self.assert_refused( + "control", + (f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", (name, key(name))), + ) + + def test_restore_refuses_a_leading_or_trailing_space(self): + for name in (" Good", "Good "): + with self.subTest(name=name): + self.assert_refused( + f"padded{len(name)}{name[0] == ' '}", + ( + f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", + (name, key(name)), + ), + ) + + def test_restore_refuses_a_wrong_name_key(self): + self.assert_refused("key", f"UPDATE {COLLECTIONS} SET name_key = 'Good'") + + def test_restore_refuses_a_non_text_key(self): + # Column affinity turns a number into text; a blob stays a blob. + self.assert_refused("blob-key", f"UPDATE {ITEMS} SET relative_path = x'61'") + + def test_restore_refuses_a_key_over_4096_characters(self): + self.assert_refused( + "long-key", (f"UPDATE {ITEMS} SET root_key = ?", ("é" * 4097,)) + ) + + def test_restore_refuses_a_timestamp_over_64_bytes(self): + self.assert_refused( + "stamp", (f"UPDATE {COLLECTIONS} SET created_at = ?", ("t" * 65,)) + ) + self.assert_refused( + "item-stamp", (f"UPDATE {ITEMS} SET added_at = ?", ("t" * 65,)) + ) + + def test_restore_refuses_an_item_of_a_missing_collection(self): + self.assert_refused("orphan", f"INSERT INTO {ITEMS} VALUES (9, 'r', 'b', 'n')") + + def test_restore_refuses_a_non_integer_collection_id(self): + self.assert_refused("text-id", f"UPDATE {ITEMS} SET collection_id = 'one'") + + def test_restore_refuses_a_trigger_or_extra_index(self): + self.require_module() + for name, statement in ( + ( + "trigger", + f"CREATE TRIGGER t AFTER INSERT ON {ITEMS} BEGIN SELECT 1; END", + ), + ("extra", f"CREATE INDEX extra ON {COLLECTIONS}(name)"), + ): + with self.subTest(name=name): + path = self.closed_current(f"{name}.db", statement) + with closing(sqlite3.connect(path)) as raw: + self.assertIsNotNone(image_collections_schema_error(raw)) + self.assertFalse(PromptModel.verify_database_file(path)[0]) + + def test_deleting_a_collection_deletes_its_members(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + conn.execute( + f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) " + "VALUES (1, 'A', 'a', ?), (2, 'B', 'b', ?)", + (NOW, NOW), + ) + conn.executemany( + f"INSERT INTO {ITEMS} VALUES (?, 'root', ?, ?)", + ((1, "a.png", NOW), (1, "b.png", NOW), (2, "a.png", NOW)), + ) + conn.execute(f"DELETE FROM {COLLECTIONS} WHERE id = 1") + self.assertEqual( + [tuple(r) for r in conn.execute(f"SELECT collection_id FROM {ITEMS}")], + [(2,)], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_restore_plans.py b/tests/test_restore_plans.py index 5ad20eaa..f5825169 100644 --- a/tests/test_restore_plans.py +++ b/tests/test_restore_plans.py @@ -49,6 +49,7 @@ SUMMARY_KEYS = { "image_links", "runs", "saved_views", + "collections", } diff --git a/tests/test_restore_stage.py b/tests/test_restore_stage.py index 3c863799..72ae806e 100644 --- a/tests/test_restore_stage.py +++ b/tests/test_restore_stage.py @@ -226,6 +226,7 @@ class TestStageAndPublish(StageTestCase): for key in ("pairs", "archived_pairs", "image_links", "runs"): self.assertEqual(staged.summary[key], 0, key) self.assertEqual(staged.summary["saved_views"], 0) + self.assertEqual(staged.summary["collections"], 0) self.assertEqual(staged.schema, {"migrated": False, "library_state": True}) self.assertEqual( staged.digest, @@ -239,6 +240,48 @@ class TestStageAndPublish(StageTestCase): staged.discard() +class TestCollectionsSummary(StageTestCase): + """#331: collections join the review's summary and the live fingerprint.""" + + def add_collection(self, db, name): + with db.model.get_connection() as conn: + conn.execute( + "INSERT INTO image_collections (name, name_key, created_at) " + "VALUES (?, ?, '2026-10-27T00:00:00+00:00')", + (name, name.casefold()), + ) + + def test_summary_counts_collections(self): + src_path = os.path.join(self.uploads, "with-collections.db") + src = PromptDatabase(src_path) + try: + self.save(src, "kept") + self.add_collection(src, "First") + self.add_collection(src, "Second") + finally: + src.close_all() + staged = self.db.model.stage_restore(src_path) + try: + self.assertEqual(staged.summary["collections"], 2) + finally: + staged.discard() + self.assertIn("collections", restore_stage.SUMMARY_KEYS) + + def test_adding_a_collection_changes_the_fingerprint(self): + conn = self.db.model.get_connection() + summary, before = restore_stage.live_fingerprint(conn) + self.assertEqual(summary["collections"], 0) + self.add_collection(self.db, "After the review") + summary, after = restore_stage.live_fingerprint(conn) + self.assertEqual(summary["collections"], 1) + self.assertNotEqual(after, before) + # A delete and a re-add with the same count still changes it (rowid). + with self.db.model.get_connection() as write: + write.execute("DELETE FROM image_collections") + self.add_collection(self.db, "Replacement") + self.assertNotEqual(restore_stage.live_fingerprint(conn)[1], after) + + class TestOlderBackups(StageTestCase): """Rollback-journal and older-schema backups restore exactly as before."""