From 7bdbc3635a58152ea8f73383fc0d64027f1c9c57 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Mon, 19 Oct 2026 22:41:50 -0700 Subject: [PATCH] feat(images): persist image favorites Favorites are stored in image_favorites keyed by the gallery root label and the relative path, set and cleared through PUT/DELETE /prompt_manager/images/favorite with a {root_index, path} body validated against the rooted listing. Workspace image cards carry favorite: bool. The table is created by migration, travels with backups, is optional in old backups and required after a restore's migration. --- database/image_favorites.py | 136 +++++++++++++++ database/models.py | 10 ++ database/operations.py | 2 + py/api/images.py | 113 +++++++++++- tests/test_api_workspace_images.py | 16 +- tests/test_database_backup.py | 66 +++++++ tests/test_image_favorites.py | 267 +++++++++++++++++++++++++++++ 7 files changed, 604 insertions(+), 6 deletions(-) create mode 100644 database/image_favorites.py create mode 100644 tests/test_image_favorites.py diff --git a/database/image_favorites.py b/database/image_favorites.py new file mode 100644 index 0000000..09abf68 --- /dev/null +++ b/database/image_favorites.py @@ -0,0 +1,136 @@ +"""Persisted image favorites, keyed by gallery root label + relative path. + +The root label is the root's public form (``_public_path``), so a favorite +follows its folder when the configured roots are reordered. Optional in old +backups, required once a database is current (see ``library_schema.py``). +""" + +import re +import sqlite3 +from typing import Iterable + +TABLE = "image_favorites" +MAX_KEY_LENGTH = 4096 +_LOOKUP_CHUNK = 500 +_DDL = ( + f"CREATE TABLE {TABLE} (root_label TEXT NOT NULL, relative_path TEXT NOT NULL, " + "created_at TEXT NOT NULL, PRIMARY KEY(root_label, relative_path))" +) +_AUTOINDEX = f"sqlite_autoindex_{TABLE}_1" +_COLUMNS = ( + ("root_label", "TEXT", 1, 1), + ("relative_path", "TEXT", 1, 2), + ("created_at", "TEXT", 1, 0), +) + + +def _compact(sql: str) -> str: + return re.sub(r"\s+", "", sql).lower() + + +def image_favorites_schema_error( + conn: sqlite3.Connection, *, allow_missing: bool = True +) -> str | None: + """Reject unsupported favorites schemas and invalid persisted values. + + Reads only. A missing table is acceptable for backups that predate it + unless ``allow_missing`` is false. + """ + objects = conn.execute( + "SELECT type,name,sql FROM sqlite_master WHERE name=? " + "OR (type IN ('index','trigger') AND tbl_name=?)", + (TABLE, TABLE), + ).fetchall() + if not objects: + return None if allow_missing else "Missing image favorites schema" + tables = [sql for kind, _name, sql in objects if kind == "table"] + others = {(kind, name, sql) for kind, name, sql in objects if kind != "table"} + if len(tables) != 1 or others != {("index", _AUTOINDEX, None)}: + return "Image favorites name is reserved by an unsupported schema object" + if _compact(tables[0] or "") != _compact(_DDL): + return "Image favorites has an invalid table definition" + columns = tuple( + (row[1], row[2].upper(), row[3], row[5]) + for row in conn.execute(f"PRAGMA table_info({TABLE})") + ) + if columns != _COLUMNS: + return "Image favorites has invalid columns" + # Bytes bound characters from above (UTF-8 is at most 4 bytes each). + limit = MAX_KEY_LENGTH * 4 + if conn.execute( + f"SELECT 1 FROM {TABLE} WHERE typeof(root_label) != 'text' " + "OR typeof(relative_path) != 'text' OR typeof(created_at) != 'text' " + f"OR length(CAST(root_label AS BLOB)) > {limit} " + f"OR length(CAST(relative_path AS BLOB)) > {limit} " + "OR length(CAST(created_at AS BLOB)) > 64 LIMIT 1" + ).fetchone(): + return "Image favorites contains invalid values" + for label, path in conn.execute( + f"SELECT CAST(root_label AS BLOB), CAST(relative_path AS BLOB) FROM {TABLE} " + f"WHERE length(CAST(root_label AS BLOB)) > {MAX_KEY_LENGTH} " + f"OR length(CAST(relative_path AS BLOB)) > {MAX_KEY_LENGTH}" + ): + for raw in (label, path): + if len(bytes(raw).decode("utf-8", "replace")) > MAX_KEY_LENGTH: + return "Image favorites contains invalid values" + return None + + +def migrate_image_favorites(conn: sqlite3.Connection) -> None: + """Install the table atomically; validate before and after creating it.""" + conn.execute("SAVEPOINT image_favorites_schema") + try: + error = image_favorites_schema_error(conn) + if error: + raise sqlite3.DatabaseError(error) + conn.execute(_DDL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1)) + error = image_favorites_schema_error(conn, allow_missing=False) + if error: + raise sqlite3.DatabaseError(error) + conn.execute("RELEASE image_favorites_schema") + except BaseException: + conn.execute("ROLLBACK TO image_favorites_schema") + conn.execute("RELEASE image_favorites_schema") + raise + + +class ImageFavoriteOperations: + """PromptDatabase mixin: set, clear and look up image favorites.""" + + def set_image_favorite(self, root_label: str, relative_path: str) -> bool: + """Mark one file as a favorite; repeating it is a no-op. Returns True.""" + from .models import utc_now_iso # models imports this module's schema + + with self.model.get_connection() as conn: + conn.execute( + "INSERT OR IGNORE INTO image_favorites VALUES (?, ?, ?)", + (root_label, relative_path, utc_now_iso()), + ) + return True + + def clear_image_favorite(self, root_label: str, relative_path: str) -> bool: + """Remove one favorite; True when a row was removed.""" + with self.model.get_connection() as conn: + cursor = conn.execute( + "DELETE FROM image_favorites WHERE root_label=? AND relative_path=?", + (root_label, relative_path), + ) + return cursor.rowcount > 0 + + def get_image_favorites( + self, root_label: str, relative_paths: Iterable[str] + ) -> set[str]: + """The subset of *relative_paths* favorited under *root_label*.""" + wanted = list(dict.fromkeys(relative_paths)) + found: set[str] = set() + conn = self.model.get_connection() + for start in range(0, len(wanted), _LOOKUP_CHUNK): + chunk = wanted[start : start + _LOOKUP_CHUNK] + marks = ",".join("?" * len(chunk)) + rows = conn.execute( # only "?" placeholders are interpolated + "SELECT relative_path FROM image_favorites WHERE root_label=? " + f"AND relative_path IN ({marks})", # nosec B608 + (root_label, *chunk), + ) + found.update(path for (path,) in rows) + return found diff --git a/database/models.py b/database/models.py index 448fd3a..85f1f23 100644 --- a/database/models.py +++ b/database/models.py @@ -16,6 +16,7 @@ from pathlib import Path from typing import Optional, Set, Tuple from .library_schema import library_schema_error, migrate_library_state +from .image_favorites import image_favorites_schema_error, migrate_image_favorites from .execution_history_schema import ( history_schema_error, migrate_execution_history, @@ -298,6 +299,7 @@ class PromptModel: self._migrate_prompt_pairs(conn) migrate_execution_history(conn) migrate_library_state(conn) + migrate_image_favorites(conn) def _migrate_prompt_pairs(self, conn: sqlite3.Connection) -> None: conn.execute("SAVEPOINT prompt_pair_schema") @@ -1016,6 +1018,9 @@ class PromptModel: Args: path: File to inspect + require_library_state: Also require the tables every current + database has (Library state and image favorites), as a + migrated restore stage must; old backups may lack them. Returns: (ok, reason): reason is "ok" when the file is acceptable, otherwise @@ -1070,6 +1075,11 @@ class PromptModel: ) if library_error: return False, library_error + favorites_error = image_favorites_schema_error( + conn, allow_missing=not require_library_state + ) + if favorites_error: + return False, favorites_error return True, "ok" except sqlite3.Error as e: return False, f"Invalid SQLite database: {e}" diff --git a/database/operations.py b/database/operations.py index 6acb86b..3fb9f18 100644 --- a/database/operations.py +++ b/database/operations.py @@ -11,6 +11,7 @@ from typing import Optional, List, Dict, Any, Set, Tuple from .models import PromptModel, normalize_image_path, utc_now_iso from .pairs import PairDependencyError, PairOperations from .library import LibraryOperations, titled +from .image_favorites import ImageFavoriteOperations from .library_commands import LibraryCommandOperations from .vocabulary import VocabularyOperations from .execution_history import ExecutionHistoryOperations @@ -301,6 +302,7 @@ class PromptDatabase( ExecutionContextOperations, ExecutionOutputOperations, HistoryReadOperations, + ImageFavoriteOperations, ): """Database operations class for managing prompts.""" diff --git a/py/api/images.py b/py/api/images.py index 87a2795..c340163 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -440,6 +440,44 @@ async def _json_body_or_error(request): return None, _json_error("Invalid JSON body", 400) +MAX_FAVORITE_PATH = 4096 + + +def _favorite_target(data: Mapping, roots: list[Path]): + """``((index, relative), None)`` for a favorite body, else ``(None, error)``. + + Rejects unknown roots, traversal, absolute or drive paths and non-media + names before any filesystem or database access. + """ + index, relative = data.get("root_index"), data.get("path") + if type(index) is not int or not 0 <= index < len(roots): + return None, bad_request("Invalid gallery root") + if ( + not isinstance(relative, str) + or not relative + or len(relative) > MAX_FAVORITE_PATH + ): + return None, bad_request("path must be a relative media path") + if ( + _has_traversal(relative) + or "\0" in relative + or re.match(r"^[a-zA-Z]:", relative) + ): + return None, _forbidden() + if Path(relative).suffix.lower() not in MEDIA_EXTENSIONS: + return None, _forbidden("Only media files can be favorites") + return (index, relative), None + + +def _listed_relative(files: list, root: Path, relative: str) -> str | None: + """The listing's own POSIX spelling of *relative* under *root*, or None.""" + target = os.path.normcase(str(Path(root) / relative)) + for path, _mtime in files: + if os.path.normcase(str(path)) == target: + return path.relative_to(root).as_posix() + return None + + def _iter_media_files(output_path, max_files): """Yield media files under *output_path* with bounded effort. @@ -716,6 +754,14 @@ class ImageRoutesMixin: async def get_gallery_subfolders_route(request): return await self.get_gallery_subfolders(request) + @routes.put("/prompt_manager/images/favorite") + async def set_image_favorite_route(request): + return await self.update_image_favorite(request, favorite=True) + + @routes.delete("/prompt_manager/images/favorite") + async def clear_image_favorite_route(request): + return await self.update_image_favorite(request, favorite=False) + # POST: drops the server-side listing cache and rewalks the output roots. @routes.post("/prompt_manager/gallery/rescan") async def rescan_gallery_route(request): @@ -1968,6 +2014,66 @@ class ImageRoutesMixin: {"success": False, "error": _safe_error(e)}, status=500 ) + async def update_image_favorite(self, request, *, favorite: bool): + """PUT/DELETE one favorite, named by ``{root_index, path}``. + + A favorite is set only on a file in the root's listing; clearing also + accepts a well-formed path that is no longer listed, so a deleted + file's favorite can still be removed. Both are idempotent. + """ + data, error = await _json_body_or_error(request) + if error is not None: + return error + roots = list(self._get_all_output_dirs()) + target, error = _favorite_target(data, roots) + if error is not None: + return error + index, relative = target + root = roots[index] + try: + files = await self._get_gallery_files(root) + listed = await self._run_in_executor( + _listed_relative, files, root, relative + ) + if listed is None and favorite: + return web.json_response( + {"success": False, "error": "Image file not found"}, status=404 + ) + key = (_public_path(root), listed or Path(relative).as_posix()) + update = ( + self.db.set_image_favorite if favorite else self.db.clear_image_favorite + ) + await self._run_in_executor(update, *key) + except Exception: + self.logger.exception("Favorite update failed") + return web.json_response( + {"success": False, "error": "Favorite update failed"}, status=500 + ) + return web.json_response( + { + "success": True, + "favorite": favorite, + "root_index": index, + "path": key[1], + } + ) + + def _workspace_favorites(self, images: list[dict]) -> set[tuple[str, str]]: + """``(root label, relative path)`` favorites among *images*.""" + by_root: dict[str, list[str]] = {} + for image in images: + by_root.setdefault(image["root_dir"], []).append(image["path"]) + try: + return { + (label, path) + for label, paths in by_root.items() + for path in self.db.get_image_favorites(label, paths) + } + except Exception: + # The listing stays usable; stars show as not favorited. + self.logger.warning("Image favorites could not be read") + return set() + async def inspect_image(self, request): """Read prompt evidence for one explicitly rooted physical media file.""" roots = list(self._get_all_output_dirs()) @@ -2208,7 +2314,7 @@ class ImageRoutesMixin: def _format_workspace_page( self, page: list[tuple], *, linked: bool = True ) -> list[dict]: - """Workspace descriptors; ``linked`` adds card facts (link summary, seed).""" + """Workspace descriptors; ``linked`` adds card facts (link, seed, favorite).""" images = self._format_output_page(page, scoped=True) paths = { (entry[3], entry[0].relative_to(entry[2]).as_posix()): entry[0] @@ -2242,6 +2348,11 @@ class ImageRoutesMixin: image["seed"] = card_seed(params) summary = summaries.get(str(path)) or _unavailable_link() image["linked"] = _card_link(summary, parsed_positive(params)) + if linked: + # Image cards carry the star; Library previews show none. + favorites = self._workspace_favorites(images) + for image in images: + image["favorite"] = (image["root_dir"], image["path"]) in favorites return images def _format_output_page(self, page, *, scoped=False): diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index 7d155c2..909e14d 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -560,8 +560,10 @@ class TestWorkspaceImages(ImageAPITestCase): for query in statements if query.lstrip().upper().startswith(("SELECT", "WITH")) ] - # The page's link summary, and the positive count's one link read. - self.assertEqual(len(selects), 2, selects) + # The page's link summary, the positive count's one link read and + # the page's one favorites read. + self.assertEqual(len(selects), 3, selects) + self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects) self.assertEqual(data["positive_count"], 1) # six files, one prompt async def test_link_summary_runs_in_executor_and_survives_database_errors(self): @@ -893,7 +895,8 @@ class TestWorkspaceImages(ImageAPITestCase): for query in statements if query.lstrip().upper().startswith(("SELECT", "WITH")) ] - self.assertEqual(len(selects), 2, selects) + self.assertEqual(len(selects), 3, selects) + self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects) _, first = await self.get(limit=2) statements.clear() with patch.object(db.model, "get_connection", side_effect=trace_connection): @@ -904,7 +907,9 @@ class TestWorkspaceImages(ImageAPITestCase): for query in statements if query.lstrip().upper().startswith(("SELECT", "WITH")) ] - self.assertEqual(len(selects), 1, selects) # the page's link summary only + # The page's link summary and its favorites read only. + self.assertEqual(len(selects), 2, selects) + self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects) self.assertNotIn("gi.file_path, p.rating", selects[0]) # Without filters the listing outlives a database failure; the count # is then unknown rather than wrong. @@ -1012,7 +1017,8 @@ class TestWorkspaceImages(ImageAPITestCase): for query in statements if query.lstrip().upper().startswith(("SELECT", "WITH")) ] - self.assertEqual(len(selects), 2, selects) + self.assertEqual(len(selects), 3, selects) + self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects) async def test_rating_sort_orders_linked_then_unlinked(self): self._filter_fixtures() diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 68c68f5..922ba31 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -428,6 +428,72 @@ class TestRestoreFromFile(BackupTestCase): PromptModel.verify_database_file(dst, require_library_state=True)[0] ) + def test_restore_old_backup_without_favorites(self): + self.db.set_image_favorite("output", "live.png") + src = self._make_source_db(["predates favorites"]) + with closing(sqlite3.connect(src)) as conn: + conn.execute("DROP TABLE image_favorites") + conn.commit() + self.assertTrue(PromptModel.verify_database_file(src)[0]) + + self.db.model.restore_from_file(src) + + self.assertEqual(_texts(self.path), ["predates favorites"]) + self.assertEqual(self.db.get_image_favorites("output", ["live.png"]), set()) + self.assertTrue(self.db.set_image_favorite("output", "new.png")) + ok, reason = PromptModel.verify_database_file( + self.path, require_library_state=True + ) + self.assertTrue(ok, reason) + with closing(sqlite3.connect(src)) as conn: # the source is untouched + self.assertIsNone( + conn.execute( + "SELECT 1 FROM sqlite_master WHERE name='image_favorites'" + ).fetchone() + ) + + def test_restore_preserves_favorites_and_closes_handles(self): + self.db.set_image_favorite("output", "live.png") + with tempfile.TemporaryDirectory() as directory: + src_path = os.path.join(directory, "source.db") + src = PromptDatabase(src_path) + try: + self._save(src, "with favorites") + src.set_image_favorite("output", "sub/kept.png") + src.set_image_favorite("extra/root", "kept.png") + finally: + src.close_all() + # The favorites still sit in the WAL; a backup must carry them. + backup = os.path.join(directory, "backup.db") + reopened = PromptDatabase(src_path) + try: + self.assertTrue(reopened.model.backup_database(backup)) + finally: + reopened.close_all() + + self.db.model.restore_from_file(backup) + + self.assertEqual( + self.db.get_image_favorites("output", ["sub/kept.png", "live.png"]), + {"sub/kept.png"}, + ) + self.assertEqual( + self.db.get_image_favorites("extra/root", ["kept.png"]), {"kept.png"} + ) + PromptModel.close_all_instances() + # Windows refuses to delete a file a restore left open. + assert_closed(self, directory) + backups = list(Path(self.tmpdir).glob("prompts.db.backup_*")) + self.assertEqual(len(backups), 1) + with closing(sqlite3.connect(str(backups[0]))) as conn: + self.assertEqual( + conn.execute( + "SELECT root_label, relative_path FROM image_favorites" + ).fetchall(), + [("output", "live.png")], + ) + assert_closed(self, str(backups[0])) + 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_favorites.py b/tests/test_image_favorites.py new file mode 100644 index 0000000..5368a0f --- /dev/null +++ b/tests/test_image_favorites.py @@ -0,0 +1,267 @@ +"""Persisted image favorites (spec D7): store, schema and rooted API boundary.""" + +import os +import sqlite3 +import tempfile +import unittest +from contextlib import closing + +try: + from test_api_images import ImageAPITestCase, make_png + from open_handles import assert_closed +except ImportError: + from tests.test_api_images import ImageAPITestCase, make_png + from tests.open_handles import assert_closed + +from database.image_favorites import ( + TABLE, + image_favorites_schema_error, + migrate_image_favorites, +) +from database.models import PromptModel +from database.operations import PromptDatabase + +FAVORITE = "/prompt_manager/images/favorite" + + +class TestImageFavoriteApi(ImageAPITestCase): + async def tearDownAsync(self): + self.api.db.close_all() + assert_closed(self, self._temp_db.name, self._output_tmp.name) + await super().tearDownAsync() + + async def send(self, method, body): + response = await self.client.request(method, FAVORITE, json=body) + return response, await response.json() + + async def workspace(self, **params): + response = await self.client.get( + "/prompt_manager/images/output", params={"mode": "workspace", **params} + ) + self.assertEqual(response.status, 200) + data = await response.json() + return {image["path"]: image for image in data["images"]} + + async def test_favorite_round_trip_and_idempotent(self): + make_png(self.output_dir / "sub" / "a.png") + make_png(self.output_dir / "b.png") + items = await self.workspace() + self.assertIs(items["sub/a.png"]["favorite"], False) + self.assertIs(items["b.png"]["favorite"], False) + + body = {"root_index": 0, "path": "sub/a.png"} + for _ in range(2): # a repeated PUT is a no-op, not an error + response, data = await self.send("PUT", body) + self.assertEqual(response.status, 200) + self.assertEqual( + data, + { + "success": True, + "favorite": True, + "root_index": 0, + "path": "sub/a.png", + }, + ) + items = await self.workspace() + self.assertIs(items["sub/a.png"]["favorite"], True) + self.assertIs(items["b.png"]["favorite"], False) + count = self.api.db.model.get_connection().execute( + f"SELECT COUNT(*) FROM {TABLE}" + ) + self.assertEqual(count.fetchone()[0], 1) + + for _ in range(2): + response, data = await self.send("DELETE", body) + self.assertEqual(response.status, 200) + self.assertTrue(data["success"]) + self.assertIs(data["favorite"], False) + items = await self.workspace() + self.assertIs(items["sub/a.png"]["favorite"], False) + + async def test_favorite_is_keyed_by_root_label_not_index(self): + second = self.comfy_dir / "second" + make_png(second / "a.png") + make_png(self.output_dir / "a.png") + self.api._get_all_output_dirs = lambda: [self.output_dir, second] + response, _ = await self.send("PUT", {"root_index": 1, "path": "a.png"}) + self.assertEqual(response.status, 200) + row = self.api.db.model.get_connection().execute( + f"SELECT root_label, relative_path FROM {TABLE}" + ) + self.assertEqual([tuple(r) for r in row], [("second", "a.png")]) + # The same root at another index keeps its favorite; the other root + # with the same relative path does not inherit it. + self.api._get_all_output_dirs = lambda: [second, self.output_dir] + self.api.invalidate_gallery_cache() + self.assertIs((await self.workspace(root=0))["a.png"]["favorite"], True) + self.assertIs((await self.workspace(root=1))["a.png"]["favorite"], False) + + async def test_favorite_rejects_traversal_and_unknown_root(self): + make_png(self.output_dir / "a.png") + make_png(self.comfy_dir / "outside.png") + (self.output_dir / "notes.txt").write_text("not media") + make_png(self.output_dir / "thumbnails" / "a_thumb.png") + cases = [ + ({"root_index": 0, "path": "../outside.png"}, 403), + ({"root_index": 0, "path": "sub/../../outside.png"}, 403), + ({"root_index": 0, "path": str(self.comfy_dir / "outside.png")}, 403), + ({"root_index": 0, "path": "/a.png"}, 403), + ({"root_index": 0, "path": "C:/a.png"}, 403), + ({"root_index": 0, "path": "a.png\0.png"}, 403), + ({"root_index": 0, "path": "notes.txt"}, 403), + ({"root_index": 1, "path": "a.png"}, 400), + ({"root_index": -1, "path": "a.png"}, 400), + ({"root_index": "0", "path": "a.png"}, 400), + ({"root_index": True, "path": "a.png"}, 400), + ({"root_index": 0}, 400), + ({"root_index": 0, "path": ""}, 400), + ({"root_index": 0, "path": 7}, 400), + ({"root_index": 0, "path": "x" * 5000 + ".png"}, 400), + ({"root_index": 0, "path": "missing.png"}, 404), + ({"root_index": 0, "path": "thumbnails/a_thumb.png"}, 404), + ] + for method in ("PUT", "DELETE"): + for body, status in cases: + if method == "DELETE" and status == 404: + continue # unlisting a stale favorite is allowed (below) + with self.subTest(method=method, body=str(body)[:60]): + response, data = await self.send(method, body) + self.assertEqual(response.status, status) + self.assertIs(data["success"], False) + self.assertNotIn(str(self.comfy_dir), data["error"]) + response = await self.client.put(FAVORITE, data=b"[1]") + self.assertEqual(response.status, 400) + response = await self.client.put(FAVORITE, data=b"{not json") + self.assertEqual(response.status, 400) + count = self.api.db.model.get_connection().execute( + f"SELECT COUNT(*) FROM {TABLE}" + ) + self.assertEqual(count.fetchone()[0], 0) + + async def test_delete_clears_a_favorite_whose_file_is_gone(self): + target = make_png(self.output_dir / "gone.png") + await self.send("PUT", {"root_index": 0, "path": "gone.png"}) + target.unlink() + self.api.invalidate_gallery_cache() + response, data = await self.send( + "DELETE", {"root_index": 0, "path": "gone.png"} + ) + self.assertEqual(response.status, 200) + self.assertIs(data["favorite"], False) + count = self.api.db.model.get_connection().execute( + f"SELECT COUNT(*) FROM {TABLE}" + ) + self.assertEqual(count.fetchone()[0], 0) + + async def test_legacy_listing_is_unchanged_by_favorites(self): + make_png(self.output_dir / "a.png") + await self.send("PUT", {"root_index": 0, "path": "a.png"}) + response = await self.client.get("/prompt_manager/images/output") + data = await response.json() + self.assertEqual(len(data["images"]), 1) + self.assertNotIn("favorite", data["images"][0]) + + async def test_favorite_store_failure_keeps_the_listing(self): + make_png(self.output_dir / "a.png") + + def broken(*_args, **_kwargs): + raise sqlite3.OperationalError("database is locked") + + self.api.db.get_image_favorites = broken + items = await self.workspace() + self.assertIs(items["a.png"]["favorite"], False) + self.api.db.set_image_favorite = broken + response, data = await self.send("PUT", {"root_index": 0, "path": "a.png"}) + self.assertEqual(response.status, 500) + self.assertEqual(data, {"success": False, "error": "Favorite update failed"}) + + +class TestImageFavoriteSchema(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + self.path = os.path.join(self.tmpdir, "favorites.db") + + def tearDown(self): + PromptModel.close_all_instances() + assert_closed(self, self.tmpdir) + for name in os.listdir(self.tmpdir): + os.unlink(os.path.join(self.tmpdir, name)) + os.rmdir(self.tmpdir) + + def test_store_round_trip_and_bounded_lookup(self): + db = PromptDatabase(self.path) + try: + self.assertTrue(db.set_image_favorite("output", "a.png")) + self.assertTrue(db.set_image_favorite("output", "a.png")) + db.set_image_favorite("other", "a.png") + self.assertEqual( + db.get_image_favorites("output", ["a.png", "b.png"]), {"a.png"} + ) + self.assertEqual(db.get_image_favorites("output", []), set()) + many = [f"{n}.png" for n in range(2000)] + ["a.png"] + self.assertEqual(db.get_image_favorites("output", many), {"a.png"}) + self.assertTrue(db.clear_image_favorite("output", "a.png")) + self.assertFalse(db.clear_image_favorite("output", "a.png")) + self.assertEqual(db.get_image_favorites("output", ["a.png"]), set()) + created = db.model.get_connection().execute( + f"SELECT created_at FROM {TABLE}" + ) + self.assertRegex(created.fetchone()[0], r"^\d{4}-\d\d-\d\dT") + finally: + db.close_all() + + def test_schema_rejects_unsupported_definitions_and_values(self): + variants = { + "missing": None, + "wrong table": f"CREATE TABLE {TABLE} (x TEXT)", + "view": f"CREATE VIEW {TABLE} AS SELECT 1", + "nullable key": ( + f"CREATE TABLE {TABLE} (root_label TEXT, relative_path TEXT, " + "created_at TEXT, PRIMARY KEY(root_label, relative_path))" + ), + } + for label, ddl in variants.items(): + with self.subTest(label): + with closing(sqlite3.connect(":memory:")) as conn: + if ddl: + conn.execute(ddl) + error = image_favorites_schema_error(conn) + if ddl is None: + self.assertIsNone(error) + self.assertIsNotNone( + image_favorites_schema_error(conn, allow_missing=False) + ) + else: + self.assertIsNotNone(error) + with self.assertRaises(sqlite3.DatabaseError): + migrate_image_favorites(conn) + with closing(sqlite3.connect(":memory:")) as conn: + migrate_image_favorites(conn) + migrate_image_favorites(conn) # idempotent + self.assertIsNone(image_favorites_schema_error(conn, allow_missing=False)) + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', 'a.png', 'now')") + self.assertIsNone(image_favorites_schema_error(conn)) + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', x'00ff', 'now')") + self.assertIsNotNone(image_favorites_schema_error(conn)) + conn.execute(f"DELETE FROM {TABLE} WHERE typeof(relative_path)='blob'") + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', ?, 'now')", ("x" * 5000,)) + self.assertIsNotNone(image_favorites_schema_error(conn)) + conn.execute(f"DELETE FROM {TABLE}") + conn.execute(f"CREATE INDEX extra ON {TABLE}(created_at)") + self.assertIsNotNone(image_favorites_schema_error(conn)) + + def test_verify_rejects_a_backup_with_an_invalid_favorites_table(self): + db = PromptDatabase(self.path) + db.close_all() + PromptModel.close_all_instances() + with closing(sqlite3.connect(self.path)) as conn: + conn.execute(f"DROP TABLE {TABLE}") + conn.execute(f"CREATE TABLE {TABLE} (anything TEXT)") + conn.commit() + ok, reason = PromptModel.verify_database_file(self.path) + self.assertFalse(ok) + self.assertIn("favorites", reason.lower()) + + +if __name__ == "__main__": + unittest.main()