nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(images): search and filter workspace images by linked prompts

Browse Source 
The workspace listing now takes category (exact), tags (comma list, all
required) and sort=rating, and q also matches the linked positive's text,
Library title, pair title and tag names. One bounded, parameterized SELECT
(get_image_link_matches / get_image_paths_matching) returns normalized
paths of unambiguously linked images; the listing is intersected with it.
Unlinked and ambiguous files are excluded by a category or tag filter and
sort last by rating, with a newest-first tie-break. The cursor scope covers
the new filters and sort; unknown values are still rejected.

The derived listing values and sort orders are cached per gallery scan so a
50,000-file listing with 20,000 links pages well under 150 ms at p95.
This commit is contained in:
Vito Sansevero
2026-10-19 23:05:34 -07:00
parent b5c2925468
commit 7c469510d9
4 changed files with 691 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
database/operations.py
+96 -1
View File
@@ -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:
py/api/__init__.py
+4
View File
@@ -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.
py/api/images.py
+251 -72
View File
@@ -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."""
tests/test_api_workspace_images.py
+340
View File
@@ -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):