feat(ebook-search): implement async phrase matching for chunks and add ChunkPhraseHit model

python/ebook_search/protected_phrases/matching.py

@@ -6,10 +6,11 @@ import logging
 from collections import defaultdict
 from typing import TYPE_CHECKING
 
-from sqlalchemy import and_, delete, or_, select, union
+from sqlalchemy import and_, delete, func, or_, select, union
 
 from python.ebook_search.protected_phrases.config import get_ignored_phrases
 from python.ebook_search.protected_phrases.models import (
+    ChunkPhraseHit,
     PhraseLookup,
     PhraseMatch,
 )
@@ -355,3 +356,50 @@ def generate_query_ngrams(
             if phrase_norm in get_ignored_phrases():
                 continue
             yield phrase_norm, start, end
+
+
+async def phrase_hits_for_chunks(
+    session: AsyncSession,
+    *,
+    chunk_ids: Sequence[int],
+    phrase_ids: Sequence[int],
+) -> dict[int, tuple[ChunkPhraseHit, ...]]:
+    """Return matched protected phrases with mention counts by chunk id.
+
+    Args:
+        session (AsyncSession): Active database session.
+        chunk_ids (Sequence[int]): Chunk ids to look up mentions for.
+        phrase_ids (Sequence[int]): Protected phrase ids to restrict the results to.
+
+    Returns:
+        dict[int, tuple[ChunkPhraseHit, ...]]: Phrase hits per chunk id, ordered by mention count.
+    """
+    if not chunk_ids or not phrase_ids:
+        return {}
+
+    mention_count = func.count(EbookChunkPhraseMention.phrase_id).label("mention_count")
+    statement = (
+        select(
+            EbookChunkPhraseMention.chunk_id,
+            EbookProtectedPhrase.id.label("phrase_id"),
+            EbookProtectedPhrase.phrase_text,
+            mention_count,
+        )
+        .join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookChunkPhraseMention.phrase_id)
+        .where(
+            EbookChunkPhraseMention.chunk_id.in_(chunk_ids),
+            EbookChunkPhraseMention.phrase_id.in_(phrase_ids),
+        )
+        .group_by(EbookChunkPhraseMention.chunk_id, EbookProtectedPhrase.id, EbookProtectedPhrase.phrase_text)
+        .order_by(EbookChunkPhraseMention.chunk_id, mention_count.desc(), EbookProtectedPhrase.phrase_text)
+    )
+    hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list)
+    for row in await session.execute(statement):
+        hits[row.chunk_id].append(
+            ChunkPhraseHit(
+                phrase_id=row.phrase_id,
+                phrase_text=row.phrase_text,
+                mention_count=row.mention_count,
+            )
+        )
+    return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()}