feat(collections): store collections in the library database
This commit is contained in:
@@ -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
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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>' },
|
||||
});
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -49,6 +49,7 @@ SUMMARY_KEYS = {
|
||||
"image_links",
|
||||
"runs",
|
||||
"saved_views",
|
||||
"collections",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user