feat(collections): store collections in the library database

This commit is contained in:
Vito Sansevero
2026-10-27 17:43:08 -07:00
parent d11b354ed6
commit 2455ab21dd
8 changed files with 825 additions and 7 deletions
+231
View File
@@ -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
+11 -1
View File
@@ -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
)
+8 -4
View File
@@ -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
+2 -2
View File
@@ -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_<date and time>' },
});
+44
View File
@@ -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()
+485
View File
@@ -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()
+1
View File
@@ -49,6 +49,7 @@ SUMMARY_KEYS = {
"image_links",
"runs",
"saved_views",
"collections",
}
+43
View File
@@ -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."""