diff --git a/database/image_favorites.py b/database/image_favorites.py index 09abf68..4334926 100644 --- a/database/image_favorites.py +++ b/database/image_favorites.py @@ -1,8 +1,11 @@ -"""Persisted image favorites, keyed by gallery root label + relative path. +"""Persisted image favorites, keyed by gallery root + 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``). +The root key is the root's real, case-normalized path and the relative key is +case-normalized the same way (see ``py/api/images.py``); both stay +server-side. Folder names alone are not unique across roots, so a favorite +follows the directory itself when the configured roots are reordered. +Optional in old backups, required once a database is current (see +``library_schema.py``). """ import re @@ -13,12 +16,12 @@ 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))" + f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT NOT NULL, " + "created_at TEXT NOT NULL, PRIMARY KEY(root_key, relative_path))" ) _AUTOINDEX = f"sqlite_autoindex_{TABLE}_1" _COLUMNS = ( - ("root_label", "TEXT", 1, 1), + ("root_key", "TEXT", 1, 1), ("relative_path", "TEXT", 1, 2), ("created_at", "TEXT", 1, 0), ) @@ -58,19 +61,19 @@ def image_favorites_schema_error( # 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' " + f"SELECT 1 FROM {TABLE} WHERE typeof(root_key) != 'text' " "OR typeof(relative_path) != 'text' OR typeof(created_at) != 'text' " - f"OR length(CAST(root_label AS BLOB)) > {limit} " + f"OR length(CAST(root_key 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} " + for root, path in conn.execute( + f"SELECT CAST(root_key AS BLOB), CAST(relative_path AS BLOB) FROM {TABLE} " + f"WHERE length(CAST(root_key AS BLOB)) > {MAX_KEY_LENGTH} " f"OR length(CAST(relative_path AS BLOB)) > {MAX_KEY_LENGTH}" ): - for raw in (label, path): + for raw in (root, path): if len(bytes(raw).decode("utf-8", "replace")) > MAX_KEY_LENGTH: return "Image favorites contains invalid values" return None @@ -97,30 +100,30 @@ def migrate_image_favorites(conn: sqlite3.Connection) -> None: class ImageFavoriteOperations: """PromptDatabase mixin: set, clear and look up image favorites.""" - def set_image_favorite(self, root_label: str, relative_path: str) -> bool: + def set_image_favorite(self, root_key: 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()), + (root_key, relative_path, utc_now_iso()), ) return True - def clear_image_favorite(self, root_label: str, relative_path: str) -> bool: + def clear_image_favorite(self, root_key: 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), + "DELETE FROM image_favorites WHERE root_key=? AND relative_path=?", + (root_key, relative_path), ) return cursor.rowcount > 0 def get_image_favorites( - self, root_label: str, relative_paths: Iterable[str] + self, root_key: str, relative_paths: Iterable[str] ) -> set[str]: - """The subset of *relative_paths* favorited under *root_label*.""" + """The subset of *relative_paths* favorited under *root_key*.""" wanted = list(dict.fromkeys(relative_paths)) found: set[str] = set() conn = self.model.get_connection() @@ -128,9 +131,9 @@ class ImageFavoriteOperations: 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=? " + "SELECT relative_path FROM image_favorites WHERE root_key=? " f"AND relative_path IN ({marks})", # nosec B608 - (root_label, *chunk), + (root_key, *chunk), ) found.update(path for (path,) in rows) return found diff --git a/py/api/__init__.py b/py/api/__init__.py index 68c2ce5..c2ee06f 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -287,6 +287,8 @@ class PromptManagerAPI( # (py/api/images.py _workspace_listed/_workspace_orders). self._workspace_listed_cache = {} self._workspace_order_cache = {} + # Favorite keys of each root's cached scan (images._favorite_listing). + self._favorite_listing_cache = {} # Single-flight locks for long jobs, keyed by job family; see # _acquire_long_job(). Created lazily so no lock binds to an event # loop before the server's loop exists. @@ -364,6 +366,7 @@ class PromptManagerAPI( """Drop every cached listing and derived workspace value (roots changed).""" self.invalidate_gallery_cache() self._workspace_listed_cache = {} + self._favorite_listing_cache = {} def add_routes(self, routes): """Register all API routes with the ComfyUI server. diff --git a/py/api/images.py b/py/api/images.py index c340163..aa75d55 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -469,13 +469,26 @@ def _favorite_target(data: Mapping, roots: list[Path]): 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 +# The path module whose case rules key favorites (tests swap in ntpath). +_path_flavor = os.path + + +def _favorite_relative_key(relative: str) -> str: + """Case-normalized POSIX key of a root-relative path. + + Lowercased with "/" separators where the platform folds case (Windows), + identity on POSIX, so a case-only rename keeps its favorite on Windows + and distinct Linux files stay distinct. + """ + return _path_flavor.normcase(relative).replace(_path_flavor.sep, "/") + + +def _favorite_root_key(root) -> str: + """Server-side identity of a gallery root: its real, case-normalized path. + + Never sent in a response; folder names alone are not unique across roots. + """ + return _favorite_relative_key(os.path.realpath(str(root))) def _iter_media_files(output_path, max_files): @@ -2030,20 +2043,20 @@ class ImageRoutesMixin: return error index, relative = target root = roots[index] + relative = Path(relative).as_posix() try: files = await self._get_gallery_files(root) - listed = await self._run_in_executor( - _listed_relative, files, root, relative - ) + listing = await self._run_in_executor(self._favorite_listing, root, files) + key = _favorite_relative_key(relative) + listed = listing.get(key) 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) + await self._run_in_executor(update, _favorite_root_key(root), key) except Exception: self.logger.exception("Favorite update failed") return web.json_response( @@ -2054,21 +2067,41 @@ class ImageRoutesMixin: "success": True, "favorite": favorite, "root_index": index, - "path": key[1], + "path": listed or relative, } ) - 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"]) + def _favorite_listing(self, root: Path, files: list) -> dict[str, str]: + """``{relative key: listing POSIX spelling}`` for one root's scan. + + Built once per scan (the gallery cache hands out the same list until + it rescans), so a toggle is a dictionary lookup. + """ + cache_key = str(root) + cached = self._favorite_listing_cache.get(cache_key) + if cached is not None and cached[0] is files: + return cached[1] + listing = {} + for path, _mtime in files: + posix = path.relative_to(root).as_posix() + listing.setdefault(_favorite_relative_key(posix), posix) + self._favorite_listing_cache[cache_key] = (files, listing) + return listing + + def _workspace_favorites(self, page: list[tuple]) -> set[tuple[int, str]]: + """``(root index, listing path)`` of the favorites among *page* entries.""" + by_root: dict[tuple[int, Path], dict[str, str]] = {} + for path, _mtime, root, index in page: + posix = path.relative_to(root).as_posix() + by_root.setdefault((index, root), {})[_favorite_relative_key(posix)] = posix try: - return { - (label, path) - for label, paths in by_root.items() - for path in self.db.get_image_favorites(label, paths) - } + found = set() + for (index, root), keys in by_root.items(): + for key in self.db.get_image_favorites( + _favorite_root_key(root), list(keys) + ): + found.add((index, keys[key])) + return found except Exception: # The listing stays usable; stars show as not favorited. self.logger.warning("Image favorites could not be read") @@ -2350,9 +2383,9 @@ class ImageRoutesMixin: image["linked"] = _card_link(summary, parsed_positive(params)) if linked: # Image cards carry the star; Library previews show none. - favorites = self._workspace_favorites(images) + favorites = self._workspace_favorites(page) for image in images: - image["favorite"] = (image["root_dir"], image["path"]) in favorites + image["favorite"] = (image["root_index"], image["path"]) in favorites return images def _format_output_page(self, page, *, scoped=False): diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 922ba31..dc7436a 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -488,7 +488,7 @@ class TestRestoreFromFile(BackupTestCase): with closing(sqlite3.connect(str(backups[0]))) as conn: self.assertEqual( conn.execute( - "SELECT root_label, relative_path FROM image_favorites" + "SELECT root_key, relative_path FROM image_favorites" ).fetchall(), [("output", "live.png")], ) diff --git a/tests/test_image_favorites.py b/tests/test_image_favorites.py index 5368a0f..1bf58a5 100644 --- a/tests/test_image_favorites.py +++ b/tests/test_image_favorites.py @@ -1,10 +1,14 @@ """Persisted image favorites (spec D7): store, schema and rooted API boundary.""" +import json +import ntpath import os import sqlite3 import tempfile import unittest from contextlib import closing +from pathlib import Path +from unittest.mock import patch try: from test_api_images import ImageAPITestCase, make_png @@ -78,17 +82,24 @@ class TestImageFavoriteApi(ImageAPITestCase): items = await self.workspace() self.assertIs(items["sub/a.png"]["favorite"], False) - async def test_favorite_is_keyed_by_root_label_not_index(self): + def rows(self): + return [ + tuple(row) + for row in self.api.db.model.get_connection().execute( + f"SELECT root_key, relative_path FROM {TABLE} ORDER BY 1, 2" + ) + ] + + async def test_favorite_is_keyed_by_root_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"}) + response, data = 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")]) + self.assertNotIn(str(self.comfy_dir), json.dumps(data)) + root_key = Path(os.path.normcase(os.path.realpath(second))).as_posix() + self.assertEqual(self.rows(), [(root_key, "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] @@ -96,6 +107,71 @@ class TestImageFavoriteApi(ImageAPITestCase): 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_roots_sharing_a_folder_name_keep_separate_favorites(self): + # Two roots under different parents, both shown as "output". + with tempfile.TemporaryDirectory() as other_parent: + other = Path(os.path.realpath(other_parent)) / "output" + make_png(other / "a.png") + make_png(self.output_dir / "a.png") + self.api._get_all_output_dirs = lambda: [self.output_dir, other] + response = await self.client.get( + "/prompt_manager/images/output", params={"mode": "workspace"} + ) + labels = { + item["root_index"]: item["root_dir"] + for item in (await response.json())["images"] + } + self.assertEqual(labels, {0: "output", 1: "output"}) + + await self.send("PUT", {"root_index": 0, "path": "a.png"}) + self.assertIs((await self.workspace(root=0))["a.png"]["favorite"], True) + self.assertIs((await self.workspace(root=1))["a.png"]["favorite"], False) + + await self.send("PUT", {"root_index": 1, "path": "a.png"}) + response, data = await self.send( + "DELETE", {"root_index": 0, "path": "a.png"} + ) + self.assertEqual(response.status, 200) + self.assertNotIn(other_parent, json.dumps(data)) + self.assertIs((await self.workspace(root=0))["a.png"]["favorite"], False) + self.assertIs((await self.workspace(root=1))["a.png"]["favorite"], True) + listing = await self.client.get( + "/prompt_manager/images/output", params={"mode": "workspace"} + ) + self.assertNotIn(other_parent, await listing.text()) + + async def test_windows_case_only_rename_keeps_the_star(self): + make_png(self.output_dir / "Sub" / "A.png") + with patch("py.api.images._path_flavor", ntpath): + response, data = await self.send( + "PUT", {"root_index": 0, "path": "Sub/A.png"} + ) + self.assertEqual(response.status, 200) + self.assertEqual(data["path"], "Sub/A.png") # the listing's spelling + (self.output_dir / "Sub" / "A.png").rename( + self.output_dir / "Sub" / "a.png" + ) + (self.output_dir / "Sub").rename(self.output_dir / "sub") + self.api.invalidate_gallery_cache() + self.assertIs((await self.workspace())["sub/a.png"]["favorite"], True) + response, data = await self.send( + "DELETE", {"root_index": 0, "path": "SUB/A.PNG"} + ) + self.assertEqual(response.status, 200) + self.assertEqual(data["path"], "sub/a.png") + self.assertIs((await self.workspace())["sub/a.png"]["favorite"], False) + self.assertEqual(self.rows(), []) + + async def test_posix_case_distinct_files_keep_independent_stars(self): + make_png(self.output_dir / "A.png") + make_png(self.output_dir / "a.png") + if len(list(self.output_dir.iterdir())) != 2: + self.skipTest("case-insensitive filesystem") + await self.send("PUT", {"root_index": 0, "path": "A.png"}) + items = await self.workspace() + self.assertIs(items["A.png"]["favorite"], True) + self.assertIs(items["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") @@ -161,6 +237,22 @@ class TestImageFavoriteApi(ImageAPITestCase): self.assertEqual(len(data["images"]), 1) self.assertNotIn("favorite", data["images"][0]) + async def test_toggle_lookup_is_built_once_per_scan(self): + make_png(self.output_dir / "a.png") + files = await self.api._get_gallery_files(self.output_dir) + first = self.api._favorite_listing(self.output_dir, files) + self.assertEqual(first, {"a.png": "a.png"}) + self.assertIs(self.api._favorite_listing(self.output_dir, files), first) + make_png(self.output_dir / "b.png") + self.api.invalidate_gallery_cache() + rescanned = await self.api._get_gallery_files(self.output_dir) + self.assertEqual( + self.api._favorite_listing(self.output_dir, rescanned), + {"a.png": "a.png", "b.png": "b.png"}, + ) + self.api.reset_gallery_roots_cache() + self.assertEqual(self.api._favorite_listing_cache, {}) + async def test_favorite_store_failure_keeps_the_listing(self): make_png(self.output_dir / "a.png") @@ -216,8 +308,8 @@ class TestImageFavoriteSchema(unittest.TestCase): "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))" + f"CREATE TABLE {TABLE} (root_key TEXT, relative_path TEXT, " + "created_at TEXT, PRIMARY KEY(root_key, relative_path))" ), } for label, ddl in variants.items():