From 72cf15455a5baf74677e1461d5496dfcabbad295 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 17 Oct 2026 14:28:01 -0700 Subject: [PATCH] fix(library): preserve full name order and batch previews (#208) --- database/library.py | 26 +++++++++- py/api/library.py | 85 ++++++++++++++++++++++++++++--- tests/test_api_library.py | 104 ++++++++++++++++++++++++++++++++++++++ tests/test_library.py | 53 +++++++++++++++++++ 4 files changed, 259 insertions(+), 9 deletions(-) diff --git a/database/library.py b/database/library.py index 41a3ca1..a33c02d 100644 --- a/database/library.py +++ b/database/library.py @@ -6,7 +6,7 @@ _ELIGIBLE = """(EXISTS(SELECT 1 FROM prompt_pairs ep WHERE ep.positive_prompt_id _ORDERS = { "newest": ("COALESCE(julianday(p.created_at),0)", "DESC"), "oldest": ("COALESCE(julianday(p.created_at),0)", "ASC"), - "name": ("LOWER(SUBSTR(p.text,1,160))", "ASC"), + "name": ("LOWER(p.text)", "ASC"), "rating": ("COALESCE(p.rating,0)", "DESC"), } @@ -249,6 +249,30 @@ class LibraryOperations: ) return {"categories": [row["category"] for row in rows]} + def get_library_previews(self, positive_ids: list[int]) -> dict[int, list[dict]]: + """Read at most two links per parent in one bounded page query.""" + if ( + not isinstance(positive_ids, list) + or len(positive_ids) > 200 + or any(type(value) is not int or value <= 0 for value in positive_ids) + ): + raise ValueError("Invalid preview parent IDs") + if not positive_ids: + return {} + result = {value: [] for value in positive_ids} + marks = ",".join("?" for _ in positive_ids) + rows = self.model.get_connection().execute( + "WITH ranked AS (SELECT id,prompt_id,image_path,ROW_NUMBER() OVER " + "(PARTITION BY prompt_id ORDER BY id DESC) AS position " + f"FROM generated_images WHERE prompt_id IN ({marks})) " + "SELECT id,prompt_id,image_path FROM ranked WHERE position<=2 " + "ORDER BY prompt_id,id DESC", + positive_ids, + ) + for row in rows: + result[row["prompt_id"]].append(dict(row)) + return result + def list_library_images( self, positive_id: int, diff --git a/py/api/library.py b/py/api/library.py index 2310670..7c328dd 100644 --- a/py/api/library.py +++ b/py/api/library.py @@ -6,6 +6,9 @@ import hashlib import hmac import json import os +import time +from collections import OrderedDict +from collections.abc import Callable from pathlib import Path from aiohttp import web @@ -16,16 +19,73 @@ from .pairs import pair_response _SECRET = os.urandom(32) +class _NameCursorStore: + """Full lexical keys stay off URLs, bounded by age, count and UTF-8 bytes.""" + + def __init__( + self, + ttl: int = 3600, + max_entries: int = 128, + max_bytes: int = 8 * 1024 * 1024, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self.clock = clock + self.ttl = ttl + self.max_entries = max_entries + self.max_bytes = max_bytes + self.entries = OrderedDict() + self.size = 0 + + def _remove(self, reference: str) -> None: + _expires, _scope, raw = self.entries.pop(reference) + self.size -= len(raw) + + def _prune(self) -> None: + now = self.clock() + for reference, (expires, _scope, _raw) in list(self.entries.items()): + if expires <= now: + self._remove(reference) + + def put(self, key: object, scope: str) -> str: + raw = json.dumps(key, ensure_ascii=False, allow_nan=False).encode() + if len(raw) > min(512 * 1024, self.max_bytes): + raise ValueError("Prompt name exceeds the cursor limit") + self._prune() + while self.entries and ( + len(self.entries) >= self.max_entries + or self.size + len(raw) > self.max_bytes + ): + self._remove(next(iter(self.entries))) + reference = base64.urlsafe_b64encode(os.urandom(24)).decode() + self.entries[reference] = (self.clock() + self.ttl, scope, raw) + self.size += len(raw) + return reference + + def get(self, reference: str, scope: str) -> object: + self._prune() + entry = self.entries.get(reference) + if entry is None or entry[1] != scope: + raise ValueError("Name cursor expired; refresh the library") + self.entries.move_to_end(reference) + return json.loads(entry[2]) + + +_NAME_CURSORS = _NameCursorStore() + + def _scope(value: dict) -> str: return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest() -def _encode(key: object, scope: str) -> str | None: +def _encode(key: object, scope: str, *, name: bool = False) -> str | None: if key is None: return None - payload = json.dumps( - {"scope": scope, "key": key}, separators=(",", ":"), allow_nan=False - ).encode() + value = ( + {"scope": scope, "reference": _NAME_CURSORS.put(key, scope)} + if name + else {"scope": scope, "key": key} + ) + payload = json.dumps(value, separators=(",", ":"), allow_nan=False).encode() signature = hmac.new(_SECRET, payload, hashlib.sha256).digest() return base64.urlsafe_b64encode(signature + payload).decode().rstrip("=") @@ -45,8 +105,13 @@ def _decode(token: str | None, scope: str) -> object: ): raise ValueError value = json.loads(payload) - if set(value) != {"scope", "key"} or value["scope"] != scope: + if ( + set(value) not in ({"scope", "key"}, {"scope", "reference"}) + or value["scope"] != scope + ): raise ValueError + if "reference" in value: + return _NAME_CURSORS.get(value["reference"], scope) return value["key"] except (ValueError, TypeError, binascii.Error, UnicodeError): raise ValueError("Invalid library cursor; refresh the library") from None @@ -124,12 +189,16 @@ class LibraryRoutesMixin: roots = list(self._get_all_output_dirs()) def previews() -> None: + records = self.db.get_library_previews( + [parent["id"] for parent in result["items"]] + ) for parent in result["items"]: - images = self.db.list_library_images(parent["id"], limit=2) - parent["previews"] = self._library_media(images["images"], roots) + parent["previews"] = self._library_media(records[parent["id"]], roots) await self._run_in_executor(previews) - result["next_cursor"] = _encode(result.pop("next_key"), scope) + result["next_cursor"] = _encode( + result.pop("next_key"), scope, name=options["sort"] == "name" + ) return web.json_response({"success": True, **result}) @pair_response diff --git a/tests/test_api_library.py b/tests/test_api_library.py index eb88198..99d8c30 100644 --- a/tests/test_api_library.py +++ b/tests/test_api_library.py @@ -1,6 +1,7 @@ """Native library endpoints use bounded cursors and rooted media identities.""" import json +import unittest from unittest.mock import patch try: @@ -178,3 +179,106 @@ class TestLibraryAPI(ImageAPITestCase): self.assertEqual( [item["text"] for item in literal["items"]], ["literal % and _"] ) + + async def test_maximum_parent_page_uses_bounded_sql_queries_for_previews(self): + for index in range(200): + self._save_prompt(f"Parent {index}") + statements = [] + original = self.api.db.model.get_connection + + def trace_connection(): + connection = original() + connection.set_trace_callback(statements.append) + return connection + + with patch.object( + self.api.db.model, "get_connection", side_effect=trace_connection + ): + response, result = await self.get(limit=200) + self.assertEqual(response.status, 200) + self.assertEqual(len(result["items"]), 200) + selects = [ + query + for query in statements + if query.lstrip().upper().startswith(("SELECT", "WITH")) + ] + self.assertLessEqual(len(selects), 5) + + async def test_long_name_cursor_stays_small_and_preserves_full_sorting(self): + prefix = "猫" * 10_000 + self._save_prompt(prefix + "Z") + alpha = self._save_prompt(prefix + "A") + _, first = await self.get(sort="name", limit=1) + self.assertEqual(first["items"][0]["id"], alpha) + self.assertLess(len(first["next_cursor"]), 512) + with self.api.db.model.get_connection() as connection: + connection.execute("DELETE FROM prompts WHERE id=?", (alpha,)) + _, second = await self.get(sort="name", cursor=first["next_cursor"]) + self.assertEqual([item["text"] for item in second["items"]], [prefix + "Z"]) + + async def test_name_cursor_expiry_and_eviction_require_refresh_without_large_urls( + self, + ): + from py.api import library as module + + for text in ("Alpha", "Beta", "Gamma"): + self._save_prompt(text) + now = [100] + store = module._NameCursorStore(max_entries=1, clock=lambda: now[0]) + with patch.object(module, "_NAME_CURSORS", store): + _, first = await self.get(sort="name", limit=1) + now[0] = 3699 + response, _ = await self.get(sort="name", cursor=first["next_cursor"]) + self.assertEqual(response.status, 200) + now[0] = 3700 + response, result = await self.get(sort="name", cursor=first["next_cursor"]) + self.assertEqual(response.status, 400) + self.assertIn("refresh", result["error"]) + response, fresh = await self.get(sort="name", limit=1) + self.assertEqual(response.status, 200) + _, replacement = await self.get(sort="name", limit=1) + self.assertNotEqual(fresh["next_cursor"], replacement["next_cursor"]) + response, result = await self.get(sort="name", cursor=fresh["next_cursor"]) + self.assertEqual(response.status, 400) + self.assertIn("refresh", result["error"]) + response, _ = await self.get( + sort="name", cursor=replacement["next_cursor"], text="Alpha" + ) + self.assertEqual(response.status, 400) + + async def test_batch_preview_hydration_keeps_two_real_images_per_parent(self): + for parent_index in range(2): + positive = self._save_prompt(f"Parent {parent_index}") + for index in range(3): + path = make_png(self.output_dir / f"{parent_index}-{index}.png") + self.api.db.link_image_to_prompt(positive, str(path)) + _, result = await self.get() + self.assertEqual(len(result["items"]), 2) + for item in result["items"]: + self.assertEqual(len(item["previews"]), 2) + self.assertTrue(item["previews"][0]["path"].endswith("-2.png")) + self.assertTrue(item["previews"][1]["path"].endswith("-1.png")) + self.assertTrue(all(image["root_index"] == 0 for image in item["previews"])) + + +class TestNameCursorBudget(unittest.TestCase): + def test_lru_byte_budget_scope_and_oversize_keys_are_bounded(self): + from py.api.library import _NameCursorStore + + store = _NameCursorStore(max_entries=2, max_bytes=40) + one = store.put("one", "scope") + two = store.put("two", "scope") + self.assertEqual(store.get(one, "scope"), "one") + three = store.put("three", "scope") + with self.assertRaises(ValueError): + store.get(two, "scope") + self.assertEqual(store.get(three, "scope"), "three") + with self.assertRaises(ValueError): + store.get(one, "other") + self.assertLessEqual(store.size, store.max_bytes) + larger = store.put("x" * 35, "scope") + self.assertEqual(len(store.entries), 1) + self.assertEqual(store.get(larger, "scope"), "x" * 35) + with self.assertRaises(ValueError): + store.put("x" * 41, "scope") + self.assertEqual(len(store.entries), 1) diff --git a/tests/test_library.py b/tests/test_library.py index 5efcb17..6d241a4 100644 --- a/tests/test_library.py +++ b/tests/test_library.py @@ -173,6 +173,7 @@ class TestLibrary(unittest.TestCase): self.db.list_library() self.db.get_library_positive(b) self.db.get_library_options() + self.db.get_library_previews([b]) self.db.list_library_images(b) finally: connection.set_authorizer(None) @@ -244,3 +245,55 @@ class TestLibrary(unittest.TestCase): self.assertTrue( all(self.db.list_library(category=value)["total"] == 1 for value in options) ) + + def test_name_sort_uses_full_long_text_and_keeps_deleted_anchor_key(self): + prefix = "shared prefix " * 30 + self.save(prefix + "Zulu") + alpha = self.save(prefix + "Alpha") + self.save(prefix + "Middle") + first = self.db.list_library(sort="name", limit=1) + self.assertEqual(first["items"][0]["id"], alpha) + self.assertEqual(first["next_key"]["value"], (prefix + "Alpha").lower()) + with self.db.model.get_connection() as connection: + connection.execute("DELETE FROM prompts WHERE id=?", (alpha,)) + rest = self.db.list_library(sort="name", after=first["next_key"]) + self.assertEqual( + [item["text"] for item in rest["items"]], + [prefix + "Middle", prefix + "Zulu"], + ) + + def test_batch_previews_return_two_newest_links_per_parent_in_one_select(self): + parents = [self.save(f"Parent {index}") for index in range(200)] + expected = {} + for parent in parents[:3]: + ids = [ + self.db.link_image_to_prompt( + parent, str(Path(self.temp.name) / f"{parent}-{index}.png") + ) + for index in range(3) + ] + expected[parent] = ids[-1:0:-1] + connection = self.db.model.get_connection() + statements = [] + connection.set_trace_callback(statements.append) + try: + previews = self.db.get_library_previews(parents) + finally: + connection.set_trace_callback(None) + selects = [ + query + for query in statements + if query.lstrip().upper().startswith(("SELECT", "WITH")) + ] + self.assertEqual(len(selects), 1) + self.assertEqual( + { + parent: [image["id"] for image in images] + for parent, images in previews.items() + if images + }, + expected, + ) + self.assertEqual(self.db.get_library_previews([]), {}) + with self.assertRaises(ValueError): + self.db.get_library_previews(parents + [9999])