test(search): pin prompt-only tags and multi-chunk tag lookups (#379)

Reads the tag names and the locators in the same annotation read helper.
This commit is contained in:
Vito Sansevero
2026-10-27 12:39:58 -07:00
parent d769623754
commit 136c984a03
3 changed files with 39 additions and 8 deletions
+6 -6
View File
@@ -454,12 +454,12 @@ class AnnotationOperations:
already folded). One read of the names that tag an image.
"""
wanted = set(wanted)
conn = self.model.get_connection()
return [
tag_id
for tag_id, name in conn.execute(_IMAGE_TAG_NAMES)
if name in wanted or (query and query in name.casefold())
]
with self._annotation_read() as conn:
return [
tag_id
for tag_id, name in conn.execute(_IMAGE_TAG_NAMES)
if name in wanted or (query and query in name.casefold())
]
def image_tag_locators(self, tag_ids) -> set[str]:
"""Locators with any byte record (current or not) carrying one of
+12 -2
View File
@@ -553,8 +553,8 @@ class CurrentImageTagsTests(AnnotationTestCase):
self.assertEqual(current, {locators[0]: ["z", "y"], locators[502]: ["x"]})
lookups = [s for s in statements if "media_annotation_files" in s]
self.assertEqual(len(lookups), 2)
# The tag search finds every record's file by tag id, the empty set's
# file never, and each set keeps its position order.
# The tag search finds every record's file by tag id, and the empty
# set's file never.
self.assertEqual(
self.db.image_tag_locators(self.db.image_tag_ids(["z", "x"], "")),
{locators[0], locators[502]},
@@ -564,6 +564,16 @@ class CurrentImageTagsTests(AnnotationTestCase):
{locators[250]: []},
)
def test_tag_locators_cover_more_ids_than_one_chunk(self):
# A broad text search can match more than 500 names: every id's file
# comes back across the chunks.
locators = [self.locator(f"t{i:04}.png") for i in range(503)]
for i, loc in enumerate(locators):
self.replace(loc, tags=[f"name_{i:04}"])
ids = self.db.image_tag_ids([], "name_")
self.assertEqual(len(ids), 503)
self.assertEqual(self.db.image_tag_locators(ids), set(locators))
class AnnotationIsolationTests(AnnotationTestCase):
def test_cascade_from_vocabulary_delete_keeps_byte_records(self):
+21
View File
@@ -2163,6 +2163,27 @@ class TestWorkspaceImageTags(ImageAPITestCase):
"tag_id IN" in sql or "EXISTS" in sql or "byte_id IN" in sql, sql
)
async def test_the_image_tag_lookup_never_reads_prompt_tags(self):
self.seed_differential()
linked = self.media("linked.png", mtime=105)
prompt = self._save_prompt(
"Prompt with its own tag", tags=["prompt_only_beach"]
)
self.api.db.link_image_to_prompt(prompt, str(linked))
db = self.api.db
# The tag is only on a prompt: the image-tag lookup never yields it,
# under the Tags filter or the text search.
self.assertEqual(db.image_tag_ids(["prompt_only_beach"], ""), [])
self.assertEqual(db.image_tag_ids([], "prompt_only"), [])
self.assertEqual(
db.image_tag_locators(db.image_tag_ids([], "beach")),
reference_candidates(db, [], "beach"),
)
# The search still matches through the linked positive's prompt tags
# (the design keeps that, #379 "Behavior that does not change").
_, body = await self.get(q="prompt_only")
self.assertEqual(self.names(body), ["linked.png"])
@unittest.skipUnless(os.environ.get("PM_PERF") == "1", "timing run (PM_PERF=1)")
async def test_tag_search_timing(self):
"""Linux timing for the PR body: 23,600 records, 10 tags each."""