diff --git a/database/operations.py b/database/operations.py index 56622e2..218ca54 100644 --- a/database/operations.py +++ b/database/operations.py @@ -6,7 +6,7 @@ import sqlite3 import json import datetime import os -from typing import Optional, List, Dict, Any +from typing import Optional, List, Dict, Any, Set, Tuple from .models import PromptModel, normalize_image_path, utc_now_iso from .pairs import PairDependencyError, PairOperations @@ -196,6 +196,46 @@ LEFT JOIN prompt_pairs pp ON pp.id = ( LEFT JOIN prompts pt ON pt.id = pp.positive_prompt_id LEFT JOIN prompt_library_state pls ON pls.prompt_id = pp.positive_prompt_id WHERE gi.image_path IN ({marks}) OR gi.file_path IN ({marks}){fallback}""" +# The workspace filters read at most this many linked rows in one SELECT: more +# than the listing's own ceiling for several output roots (MAX_GALLERY_FILES +# per root), so a real gallery never reaches it. +MAX_LINK_FILTER_ROWS = 500_000 +# One image row describes a card only when no other row claims the same file +# (as in get_image_link_summaries, several rows are ambiguous). The pair scope +# and the linked positive follow _LINK_SUMMARY_SQL; category, rating, tags and +# the searchable text and titles come from that positive (spec D12/D13). Every +# writer since the 3.2.4 migration stores file_path, so it is the path key. +# The prompt-level conditions are uncorrelated IN sets SQLite builds once, and +# the CASE evaluates the per-row ambiguity probe only for admitted rows. +_LINK_FILTER_SQL = """SELECT gi.file_path, p.rating, {hit} +FROM generated_images gi +LEFT JOIN prompt_pairs pp ON pp.id = ( + SELECT CASE WHEN COUNT(DISTINCT ip.pair_id) = 1 THEN MIN(ip.pair_id) END + FROM image_prompt_pairs ip WHERE ip.image_id = gi.id) +JOIN prompts p ON p.id = COALESCE(pp.positive_prompt_id, gi.prompt_id) +WHERE gi.file_path IS NOT NULL AND CASE WHEN {admit} THEN NOT EXISTS ( + SELECT 1 FROM generated_images other + WHERE other.file_path = gi.file_path AND other.id <> gi.id) ELSE 0 END +LIMIT :limit""" +# The positive's text, its Library title or one of its tag names, or the pair +# title, contains :pattern (LIKE, so ASCII case is ignored as in Library search). +_LINK_FILTER_HIT = """(p.id IN ( + SELECT hp.id FROM prompts hp + LEFT JOIN prompt_library_state hls ON hls.prompt_id = hp.id + WHERE hp.text LIKE :pattern ESCAPE '\\' OR hls.title LIKE :pattern ESCAPE '\\' + OR hp.id IN (SELECT ht.prompt_id FROM prompt_tags ht JOIN tags hn + ON hn.id = ht.tag_id WHERE hn.name LIKE :pattern ESCAPE '\\')) + OR pp.title LIKE :pattern ESCAPE '\\')""" +_LINK_FILTER_TAG = ( + "p.id IN (SELECT ft.prompt_id FROM prompt_tags ft JOIN tags fn " + "ON fn.id = ft.tag_id WHERE fn.name = :tag{index})" +) + + +class LinkFilterTooBroad(ValueError): + """The linked rows a workspace filter reads exceed MAX_LINK_FILTER_ROWS.""" + + _ASCII_FOLD = str.maketrans("ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz") @@ -2265,6 +2305,61 @@ class PromptDatabase( ) return summaries + def get_image_link_matches( + self, + q: str, + category: Optional[str], + tags: List[str], + *, + every_link: bool = False, + ) -> Dict[str, Tuple[bool, Optional[int]]]: + """Map normalized paths of unambiguously linked images to (hit, rating). + + Rows are limited to images whose linked positive has *category* and + every tag in *tags* (exact names). ``hit`` says whether *q* occurs in + the positive's text, Library title or a tag name, or in the pair title. + Without category or tags only hits are returned, unless *every_link* + asks for every linked image (the rating sort). One SELECT, bounded by + MAX_LINK_FILTER_ROWS: LinkFilterTooBroad is raised rather than + returning a truncated set. + """ + params: Dict[str, Any] = {"limit": MAX_LINK_FILTER_ROWS + 1} + hit = "1" + if q: + hit = _LINK_FILTER_HIT + params["pattern"] = f"%{escape_like(q)}%" + admitted = [] + if category is not None: + admitted.append("p.category = :category") + params["category"] = category + for index, tag in enumerate(tags): + admitted.append(_LINK_FILTER_TAG.format(index=index)) + params[f"tag{index}"] = tag + if q and not admitted and not every_link: + admitted.append(hit) + sql = _LINK_FILTER_SQL.format(hit=hit, admit=" AND ".join(admitted) or "1") + with self.model.get_connection() as conn: + rows = conn.execute(sql, params).fetchall() + if len(rows) > MAX_LINK_FILTER_ROWS: + raise LinkFilterTooBroad("Too many linked images match; narrow the filter") + return {path: (bool(matched), rating) for path, rating, matched in rows} + + def get_image_paths_matching( + self, q: str, category: Optional[str], tags: List[str] + ) -> Set[str]: + """Normalized paths whose linked positive matches *q*, *category* and *tags*. + + See get_image_link_matches; an empty *q* matches every linked image + the category and tags admit. + """ + return { + path + for path, (hit, _rating) in self.get_image_link_matches( + q, category, tags + ).items() + if hit or not q + } + def get_prompt_id_for_image(self, image_path: str) -> Optional[int]: """Return the prompt_id linked to an image, or None.""" with self.model.get_connection() as conn: diff --git a/py/api/__init__.py b/py/api/__init__.py index ed24041..4078ca9 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -283,6 +283,10 @@ class PromptManagerAPI( self._html_cache = {} # Cached HTML file contents keyed by path self._gallery_cache = {} # dict: path_str -> (files, timestamp) self._gallery_cache_ttl = 30 # Cache TTL in seconds + # Workspace listing values and sort orders derived from a cached scan + # (py/api/images.py _workspace_listed/_workspace_orders). + self._workspace_listed_cache = {} + self._workspace_order_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. diff --git a/py/api/images.py b/py/api/images.py index 77e7abb..bc50d4d 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -11,8 +11,10 @@ import re import threading import time as _time import urllib.parse +from bisect import bisect_left, bisect_right from collections.abc import Mapping from pathlib import Path +from typing import NamedTuple from aiohttp import web from PIL import Image, UnidentifiedImageError @@ -30,11 +32,15 @@ from .inspection_service import ( try: from ...database.library import display_title + from ...database.models import normalize_image_path + from ...database.operations import LinkFilterTooBroad from ...utils.card_seed import card_seed, header_params, parsed_positive from ...utils.parallel import map_parallel from ...utils.inspection_response import encode_inspection_response except ImportError: from database.library import display_title + from database.models import normalize_image_path + from database.operations import LinkFilterTooBroad from utils.card_seed import card_seed, header_params, parsed_positive from utils.parallel import map_parallel from utils.inspection_response import encode_inspection_response @@ -292,14 +298,66 @@ def _workspace_root(query: dict, roots: list[Path]) -> int | None: return index +WORKSPACE_ORDERS = ("newest", "oldest", "name", "name-desc", "rating") +# Library-compatible bounds on the linked-positive filters (database/library.py). +MAX_FILTER_CATEGORY = 200 +MAX_FILTER_TAGS = 30 +MAX_FILTER_TAG = 200 + + +class _Listed(NamedTuple): + """One listed file plus the derived values every filter and sort reads. + + Built once per gallery scan (see ``_workspace_orders``); the first four + fields are the ``(path, mtime, root, index)`` entry the page formatter takes. + """ + + path: Path + mtime: float + root: Path + index: int + folded: str # relative POSIX path, case-folded (filename search) + folder: str # str(relative.parent), native separators (subfolder filter) + name: str # case-folded basename (name sorts) + tie: str # sha256 of the relative POSIX path (bounded tie-break) + video: bool + normalized: str # normalize_image_path, compared with linked rows + + +def _listed(path: Path, mtime: float, root: Path, index: int) -> _Listed: + text, prefix = str(path), os.path.join(str(root), "") + relative = ( + text[len(prefix) :] if text.startswith(prefix) else str(path.relative_to(root)) + ) + posix = relative if os.sep == "/" else relative.replace(os.sep, "/") + name = relative.rpartition(os.sep)[2] + return _Listed( + path, + mtime, + root, + index, + posix.casefold(), + os.path.dirname(relative) or ".", + name.casefold(), + hashlib.sha256(posix.encode("utf-8")).hexdigest(), + os.path.splitext(name)[1].lower() in VIDEO_EXTENSIONS, + normalize_image_path(text), + ) + + def _workspace_key(entry: tuple, order: str) -> tuple: - path, mtime, root, index = entry - relative = path.relative_to(root).as_posix() - tie = hashlib.sha256(relative.encode("utf-8")).hexdigest() - name = path.name.casefold() + """The keyset sort key of *entry* (a _Listed or a 4-tuple entry).""" + if not isinstance(entry, _Listed): + entry = _listed(*entry) if order in ("name", "name-desc"): - return (name, tie, index, mtime) - return ((-mtime if order == "newest" else mtime), name, tie, index) + return (entry.name, entry.tie, entry.index, entry.mtime) + mtime = -entry.mtime if order in ("newest", "rating") else entry.mtime + return (mtime, entry.name, entry.tie, entry.index) + + +def _rated_key(entry: _Listed, rank: int) -> tuple: + """Highest rated: rating 5..1, linked unrated (0), unlinked (-1); then newest.""" + return (-rank, *_workspace_key(entry, "rating")) def _workspace_cursor(key: tuple, scope: str) -> str: @@ -330,6 +388,8 @@ def _workspace_cursor_key(token: str, scope: str, order: str) -> tuple: if order in ("name", "name-desc") else ((float, int), str, str, int) ) + if order == "rating": + types = (int, *types) if len(key) != len(types) or not all( type(value) in (kind if isinstance(kind, tuple) else (kind,)) for value, kind in zip(key, types) @@ -461,6 +521,79 @@ def _output_image_entry( return entry +def _workspace_filters(query: Mapping) -> dict: + """Validate the workspace listing's filter and sort parameters. + + ``q`` (at most 500 characters) matches the relative path, case-folded, and + the linked positive's text, title and tag names. ``category`` is exact, + ``tags`` is a comma list that must all be present, ``sort=rating`` ranks by + the linked positive. Raises ValueError for any unknown or oversized value. + """ + text = query.get("q", "").strip() + folded = text.casefold() + media = query.get("media", "all") + order = query.get("sort", "newest") + folder = query.get("subfolder", "").strip().replace("\\", "/") + category = query.get("category") or None + tags = sorted( + {tag.strip() for tag in query.get("tags", "").split(",") if tag.strip()} + ) + if ( + len(text) > 500 + or len(folded) > 500 + or media not in ("all", "image", "video") + or order not in WORKSPACE_ORDERS + or (category is not None and len(category) > MAX_FILTER_CATEGORY) + or len(tags) > MAX_FILTER_TAGS + or any(len(tag) > MAX_FILTER_TAG for tag in tags) + ): + raise ValueError("Invalid gallery filters") + if folder and (_has_traversal(folder) or len(folder) > 1024): + raise ValueError("Invalid gallery folder") + return { + "text": text, + "query": folded, + "media": media, + "order": order, + "folder": folder, + "category": category, + "tags": tags, + } + + +def _filtered_page( + entries: list, start: int, size: int, filters: dict, links: dict | None +) -> tuple[int, list]: + """Count the *entries* the filters admit and take *size* from *start* on. + + *links* maps normalized paths to (hit, rating) for linked images (see + PromptDatabase.get_image_link_matches); a category or tag filter admits + only those, and ``q`` admits a filename match or a linked-positive hit. + """ + query, media, folder = filters["query"], filters["media"], filters["folder"] + members_only = filters["category"] is not None or bool(filters["tags"]) + native_folder = folder.replace("/", os.sep) + nested = native_folder + os.sep + total, page = 0, [] + for position, entry in enumerate(entries): + if folder and not ( + entry.folder == native_folder or entry.folder.startswith(nested) + ): + continue + if media != "all" and (media == "video") != entry.video: + continue + if links is not None and (members_only or query): + match = links.get(entry.normalized) + if members_only and match is None: + continue + if query and query not in entry.folded and not (match and match[0]): + continue + total += 1 + if position >= start and len(page) < size: + page.append(entry) + return total, page + + def _in_subfolder(rel_dir, subfolder): return rel_dir == subfolder or rel_dir.startswith(subfolder + os.sep) @@ -810,32 +943,26 @@ class ImageRoutesMixin: return bad_request("limit and offset must be integers") try: root_index = _workspace_root(request.query, output_dirs) - query = request.query.get("q", "").strip().casefold() - media = request.query.get("media", "all") - order = request.query.get("sort", "newest") - folder = request.query.get("subfolder", "").strip().replace("\\", "/") - if ( - len(query) > 500 - or media not in ("all", "image", "video") - or order not in ("newest", "oldest", "name", "name-desc") - ): - raise ValueError("Invalid gallery filters") - if folder and (_has_traversal(folder) or len(folder) > 1024): - raise ValueError("Invalid gallery folder") + filters = _workspace_filters(request.query) scope = hashlib.sha256( json.dumps( [ - query, - media, - order, - folder, + filters["query"], + filters["text"], + filters["media"], + filters["order"], + filters["folder"], root_index, [str(root) for root in output_dirs], + filters["category"], + filters["tags"], ] ).encode() ).hexdigest() token = request.query.get("cursor") - after = _workspace_cursor_key(token, scope, order) if token else None + after = ( + _workspace_cursor_key(token, scope, filters["order"]) if token else None + ) except ValueError as exc: return bad_request(str(exc)) @@ -844,17 +971,12 @@ class ImageRoutesMixin: for index, root in enumerate(output_dirs) if root_index is None or index == root_index ] - page, total, has_more, next_cursor = await self._run_in_executor( - self._select_workspace_page, - listings, - folder, - query, - media, - order, - after, - limit, - scope, - ) + try: + page, total, has_more, next_cursor = await self._run_in_executor( + self._select_workspace_page, listings, filters, after, limit, scope + ) + except LinkFilterTooBroad as exc: + return bad_request(str(exc)) images = await self._run_in_executor(self._format_workspace_page, page) return web.json_response( { @@ -867,52 +989,109 @@ class ImageRoutesMixin: } ) - @staticmethod + def _workspace_listed(self, index: int, root: Path, files: list) -> list: + """The _Listed values of one root's scan, built once per scan. + + The gallery cache hands out the same list object until it rescans, so + the values are reused while the scan is still that same object. + """ + key = (index, str(root)) + cached = self._workspace_listed_cache.get(key) + if cached is not None and cached[0] is files: + return cached[1] + listed = [_listed(path, mtime, root, index) for path, mtime in files] + self._workspace_listed_cache[key] = (files, listed) + return listed + + def _workspace_orders(self, listings: list[tuple]) -> dict: + """The selected roots' files in newest order, plus lazily added orders.""" + parts = tuple( + self._workspace_listed(index, root, files) + for index, root, files in listings + ) + key = tuple((index, str(root)) for index, root, _files in listings) + cached = self._workspace_order_cache.get(key) + if cached is not None and all( + old is new for old, new in zip(cached[0], parts, strict=True) + ): + return cached[1] + newest = sorted( + (entry for part in parts for entry in part), + key=lambda entry: _workspace_key(entry, "newest"), + ) + orders = {"newest": newest} + self._workspace_order_cache[key] = (parts, orders) + return orders + + def _workspace_order(self, listings: list[tuple], order: str) -> list: + """The listing sorted by a filesystem order (built lazily, cached).""" + orders = self._workspace_orders(listings) + if order not in orders: + if order == "name-desc": + orders[order] = self._workspace_order(listings, "name")[::-1] + else: + orders[order] = sorted( + orders["newest"], key=lambda entry: _workspace_key(entry, order) + ) + return orders[order] + def _select_workspace_page( + self, listings: list[tuple], - folder: str, - query: str, - media: str, - order: str, + filters: dict, after: tuple | None, limit: int, scope: str, ) -> tuple: - entries = [] - for index, root, files in listings: - for path, mtime in files: - relative = path.relative_to(root) - if folder and not _in_subfolder( - str(relative.parent), folder.replace("/", os.sep) - ): - continue - if query and query not in relative.as_posix().casefold(): - continue - kind = "video" if path.suffix.lower() in VIDEO_EXTENSIONS else "image" - if media != "all" and media != kind: - continue - entries.append((path, mtime, root, index)) - total = len(entries) - reverse = order == "name-desc" - entries.sort(key=lambda entry: _workspace_key(entry, order), reverse=reverse) + """One keyset page of the filtered listing (blocking; run in executor). + + At most one SELECT (get_image_link_matches) when ``q``, a category, + tags or the rating sort need the linked positives. + """ + order, category, tags = filters["order"], filters["category"], filters["tags"] + links = None + if filters["query"] or category is not None or tags or order == "rating": + links = self.db.get_image_link_matches( + filters["text"], category, tags, every_link=order == "rating" + ) + if order == "rating": + + def rank(entry): + match = links.get(entry.normalized) + return -1 if match is None else (match[1] or 0) + + # Stable: within one rank the newest order (the tie-break) stays. + entries = sorted( + self._workspace_order(listings, "newest"), + key=lambda entry: -rank(entry), + ) + + def key(entry): + return _rated_key(entry, rank(entry)) + + else: + entries = self._workspace_order(listings, order) + + def key(entry): + return _workspace_key(entry, order) + + start = 0 if after is not None: - entries = [ - entry - for entry in entries - if ( - _workspace_key(entry, order) < after - if reverse - else _workspace_key(entry, order) > after - ) - ] - page = entries[:limit] - has_more = len(entries) > limit - next_cursor = ( - _workspace_cursor(_workspace_key(page[-1], order), scope) - if has_more - else None + if order == "name-desc": + ascending = self._workspace_order(listings, "name") + start = len(entries) - bisect_left(ascending, after, key=key) + else: + start = bisect_right(entries, after, key=key) + total, page = _filtered_page(entries, start, limit + 1, filters, links) + has_more = len(page) > limit + page = page[:limit] + next_cursor = _workspace_cursor(key(page[-1]), scope) if has_more else None + return ( + [(entry.path, entry.mtime, entry.root, entry.index) for entry in page], + total, + has_more, + next_cursor, ) - return page, total, has_more, next_cursor async def serve_image(self, request): """Serve the actual image file using streamed FileResponse.""" diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index 843c451..da42b28 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -1,9 +1,11 @@ """Opt-in workspace gallery contracts against disposable SQLite and media.""" +import asyncio import json import os import shutil import subprocess # nosec B404 +import sys import unittest from pathlib import Path from unittest.mock import AsyncMock, patch @@ -759,6 +761,344 @@ class TestWorkspaceImages(ImageAPITestCase): (await self.inspected("shared.png", items["shared.png"]))["title"], title ) + # B3: search, category/tag filters and rating sort over linked positives. + + def _names(self, data): + return [item["filename"] for item in data["images"]] + + def _filter_fixtures(self): + """Linked, pair-linked, ambiguous and unlinked files for the B3 filters.""" + db = self.api.db + night_ab = self._save_prompt( + "A ZEBRA under neon", category="Night", tags=["a", "b"], rating=5 + ) + night_a = self._save_prompt("Quiet harbor", category="Night", tags=["a"]) + day_ab = self._save_prompt( + "Morning field", category="Day", tags=["a", "b", "zebra-tag"], rating=3 + ) + titled = self._save_prompt("Plain words", category="Day", rating=1) + with db.model.get_connection() as connection: + connection.execute( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + (titled, "Zebra stored title"), + ) + percent = self._save_prompt("100% wild_card", category="Day") + pair_positive = self._save_prompt("Pair positive", category="Night", rating=4) + pair = db.save_prompt_pair("Pair positive", "Noise", title="Zebra pair") + owners = ( + self._save_prompt("Zebra one", category="Night"), + self._save_prompt("Zebra two", category="Night"), + ) + files = { + name: self.media(f"{name}.png", mtime=100 + index) + for index, name in enumerate( + ( + "night_ab", + "night_a", + "day_ab", + "titled", + "percent", + "paired", + "ambiguous", + "zebra_unlinked", + "plain_unlinked", + ) + ) + } + db.link_image_to_prompt(night_ab, str(files["night_ab"])) + # A non-normalized spelling of the same file still compares equal. + db.link_image_to_prompt( + night_a, str(self.output_dir / "sub" / ".." / "night_a.png") + ) + db.link_image_to_prompt(day_ab, str(files["day_ab"])) + db.link_image_to_prompt(titled, str(files["titled"])) + db.link_image_to_prompt(percent, str(files["percent"])) + self.assertEqual(db.get_prompt_pair(pair)["positive_prompt_id"], pair_positive) + image = db.link_image_to_prompt(pair_positive, str(files["paired"])) + db.link_image_to_pair(pair, image, {"source": "fixture"}) + for owner in owners: + db.link_image_to_prompt(owner, str(files["ambiguous"])) + return files + + async def test_search_matches_filename_prompt_text_and_tags(self): + from database.models import normalize_image_path + + files = self._filter_fixtures() + _, data = await self.get(q="ZEBRA", sort="name") + # Text, stored title, tag, pair title and filename; never an ambiguous link. + self.assertEqual( + self._names(data), + [ + "day_ab.png", + "night_ab.png", + "paired.png", + "titled.png", + "zebra_unlinked.png", + ], + ) + self.assertEqual(data["total"], 5) + _, data = await self.get(q="0% w") + self.assertEqual(self._names(data), ["percent.png"]) + _, data = await self.get(q="%") + self.assertEqual(self._names(data), ["percent.png"]) + _, data = await self.get(q="harbor") + self.assertEqual(self._names(data), ["night_a.png"]) + self.assertEqual( + self.api.db.get_image_paths_matching("zebra", None, []), + { + normalize_image_path(str(files[name])) + for name in ("night_ab", "day_ab", "titled", "paired") + }, + ) + + async def test_category_and_tags_filter_linked_images_only(self): + self._filter_fixtures() + cases = ( + ({"category": "Night"}, ["night_ab.png", "night_a.png", "paired.png"]), + ({"category": "night"}, []), + ({"tags": "a,b"}, ["night_ab.png", "day_ab.png"]), + ({"tags": " a , b ,"}, ["night_ab.png", "day_ab.png"]), + ({"category": "Night", "tags": "a,b"}, ["night_ab.png"]), + ({"category": "Night", "q": "harbor"}, ["night_a.png"]), + ({"category": "Night", "q": "paired"}, ["paired.png"]), + ({"category": "Day", "q": "zebra"}, ["day_ab.png", "titled.png"]), + ({"category": ""}, None), + ) + _, everything = await self.get(sort="name") + for params, expected in cases: + with self.subTest(**params): + response, data = await self.get(sort="name", **params) + self.assertEqual(response.status, 200, data) + if expected is None: + self.assertEqual(data["images"], everything["images"]) + else: + self.assertEqual(self._names(data), sorted(expected)) + self.assertEqual(data["total"], len(expected)) + for params in ( + {"category": "x" * 201}, + {"tags": ",".join(f"t{i}" for i in range(31))}, + {"tags": "x" * 201}, + ): + response, data = await self.get(**params) + self.assertEqual(response.status, 400, params) + # The filters read one SELECT besides the page's link summary. + 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, _ = await self.get( + category="Night", tags="a", q="zebra", sort="rating" + ) + self.assertEqual(response.status, 200) + selects = [ + query + for query in statements + if query.lstrip().upper().startswith(("SELECT", "WITH")) + ] + self.assertEqual(len(selects), 2, selects) + + async def test_rating_sort_orders_linked_then_unlinked(self): + self._filter_fixtures() + response, data = await self.get(sort="rating") + self.assertEqual(response.status, 200, data) + self.assertEqual( + self._names(data), + [ + "night_ab.png", # 5 + "paired.png", # 4, the pair's positive + "day_ab.png", # 3 + "titled.png", # 1 + # linked without a rating, newest first + "percent.png", + "night_a.png", + # unlinked (and ambiguous) last, newest first + "plain_unlinked.png", + "zebra_unlinked.png", + "ambiguous.png", + ], + ) + _, first = await self.get(sort="rating", limit=4) + _, rest = await self.get(sort="rating", limit=10, cursor=first["next_cursor"]) + self.assertEqual(first["images"] + rest["images"], data["images"]) + _, night = await self.get(sort="rating", category="Night") + self.assertEqual( + self._names(night), ["night_ab.png", "paired.png", "night_a.png"] + ) + + async def test_cursor_rejects_changed_filters(self): + self._filter_fixtures() + _, first = await self.get(category="Night", limit=1) + token = first["next_cursor"] + self.assertIsInstance(token, str) + response, _ = await self.get(category="Night", limit=1, cursor=token) + self.assertEqual(response.status, 200) + for params in ( + {"category": "Day"}, + {}, + {"category": "Night", "tags": "a"}, + {"category": "Night", "q": "zebra"}, + {"category": "Night", "sort": "rating"}, + ): + with self.subTest(**params): + response, _ = await self.get(cursor=token, **params) + self.assertEqual(response.status, 400) + _, rated = await self.get(sort="rating", limit=1) + response, _ = await self.get(sort="rating", cursor=rated["next_cursor"]) + self.assertEqual(response.status, 200) + response, _ = await self.get(sort="newest", cursor=rated["next_cursor"]) + self.assertEqual(response.status, 400) + _, tagged = await self.get(tags="a,b", limit=1) + response, _ = await self.get(tags="b,a", cursor=tagged["next_cursor"]) + self.assertEqual(response.status, 200) + response, _ = await self.get(tags="a", cursor=tagged["next_cursor"]) + self.assertEqual(response.status, 400) + for params in ({"sort": "rated"}, {"sort": "rating-desc"}, {"media": "x"}): + response, _ = await self.get(**params) + self.assertEqual(response.status, 400, params) + + async def test_windows_style_link_paths_compare_normalized(self): + import ntpath + + def windows_normalize(path): + return ntpath.normcase(ntpath.normpath(path)) + + prompt = self._save_prompt("Case words", category="Night", rating=2) + listed = self.media("Sub/Night.PNG") + self.media("Sub/other.png") + stored = str(listed).upper().replace("/", "\\") + with ( + patch("database.operations.normalize_image_path", windows_normalize), + patch("py.api.images.normalize_image_path", windows_normalize), + ): + self.api.db.link_image_to_prompt(prompt, stored) + _, data = await self.get(category="Night") + self.assertEqual(self._names(data), ["Night.PNG"]) + _, data = await self.get(q="case words") + self.assertEqual(self._names(data), ["Night.PNG"]) + _, data = await self.get(sort="rating") + self.assertEqual(self._names(data), ["Night.PNG", "other.png"]) + + async def test_too_broad_link_filter_is_a_safe_bad_request(self): + self._filter_fixtures() + with patch("database.operations.MAX_LINK_FILTER_ROWS", 1): + response, data = await self.get(category="Night") + self.assertEqual(response.status, 400) + self.assertIn("narrow", data["error"]) + self.assertNotIn(str(self.comfy_dir), str(data)) + + def _budget_fixture(self, files, links): + """*files* listed paths and *links* linked rows, written in bulk.""" + from database.models import normalize_image_path + + listed = [ + (self.output_dir / f"day-{i % 50:02}" / f"img_{i:05}.png", 1_000.0 + i) + for i in range(files) + ] + prompts = 2_000 + with self.api.db.model.get_connection() as conn: + conn.executemany( + "INSERT INTO tags(id,name) VALUES(?,?)", + [(i + 1, f"tag-{i}") for i in range(20)], + ) + conn.executemany( + "INSERT INTO prompts(id,text,category,rating,hash) VALUES(?,?,?,?,?)", + [ + ( + i + 1, + f"prompt {i} " + "lantern " * (i % 7) + "harbor " * 20, + f"cat-{i % 10}", + (i % 6) or None, + f"hash-{i}", + ) + for i in range(prompts) + ], + ) + conn.executemany( + "INSERT OR IGNORE INTO prompt_tags(prompt_id,tag_id) VALUES(?,?)", + [(i + 1, t) for i in range(prompts) for t in (i % 20 + 1, 1)], + ) + conn.executemany( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + [(i + 1, f"Title {i}") for i in range(0, prompts, 3)], + ) + conn.executemany( + "INSERT INTO generated_images(prompt_id,image_path,filename,file_path)" + " VALUES(?,?,?,?)", + [ + ( + i % prompts + 1, + str(listed[i * 2][0]), + listed[i * 2][0].name, + normalize_image_path(str(listed[i * 2][0])), + ) + for i in range(links) + ], + ) + return listed + + async def test_filter_budget(self): + import math + import time + + # Built off the event loop: the bulk inserts take a few hundred ms. + listed = await asyncio.get_running_loop().run_in_executor( + None, self._budget_fixture, 50_000, 20_000 + ) + + def format_page(page): + return [{"filename": entry[0].name} for entry in page] + + requests = [ + {"q": "lantern"}, + {"q": "img_0123"}, + {"q": "Title 1"}, + {"category": "cat-3"}, + {"tags": "tag-0,tag-2"}, + {"sort": "rating"}, + {"sort": "rating", "category": "cat-4"}, + {"q": "harbor", "category": "cat-5", "tags": "tag-5"}, + {"sort": "name", "q": "tag-7"}, + {"subfolder": "day-07", "q": "lantern"}, + ] + timings = [] + with ( + patch.object( + self.api, "_get_gallery_files", AsyncMock(return_value=listed) + ), + patch.object(self.api, "_format_workspace_page", side_effect=format_page), + ): + # The first request after a scan builds the cached sort orders. + started = time.perf_counter() + response, _ = await self.get(sort="rating") + cold = time.perf_counter() - started + self.assertEqual(response.status, 200) + for index in range(20): + params = dict(requests[index % len(requests)], limit=100) + if index >= len(requests): + _, first = await self.get(**params) + if first["next_cursor"]: + params["cursor"] = first["next_cursor"] + started = time.perf_counter() + response, data = await self.get(**params) + timings.append(time.perf_counter() - started) + self.assertEqual(response.status, 200, (params, data)) + _, data = await self.get(category="cat-3") + timings.sort() + p95 = timings[math.ceil(0.95 * len(timings)) - 1] + print( + f"\nworkspace filter budget: p95={p95 * 1000:.1f} ms " + f"max={timings[-1] * 1000:.1f} ms cold={cold * 1000:.1f} ms", + file=sys.stderr, + ) + self.assertLess(p95, 0.150, [round(t * 1000, 1) for t in timings]) + class TestWorkspaceCursorEncoding(unittest.TestCase): def test_ntfs_unicode_basename_cursor_is_bounded_and_round_trips(self):