Aurora-Core/app/services/search/postgresql_provider.py

"""PostgreSQL-backed search: full-text search ranking blended with trigram
similarity for typo tolerance, over primary titles/names and alternate
titles. Requires the `pg_trgm` extension and the functional GIN indexes
created in the initial migration.
"""

from __future__ import annotations

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession

from app.services.search.base import SearchMatch, SearchProvider, SearchResultIDs

_MOVIE_SQL = text(
    """
    WITH scored AS (
        SELECT m.id AS entity_id,
               GREATEST(
                   ts_rank_cd(to_tsvector('simple', coalesce(m.title, '') || ' ' || coalesce(m.original_title, '')), plainto_tsquery('simple', :q)),
                   similarity(m.title, :q)
               ) AS score
        FROM movies m
        WHERE to_tsvector('simple', coalesce(m.title, '') || ' ' || coalesce(m.original_title, '')) @@ plainto_tsquery('simple', :q)
           OR m.title % :q
        UNION ALL
        SELECT mat.movie_id AS entity_id, similarity(mat.title, :q) AS score
        FROM movie_alternate_titles mat
        WHERE mat.title % :q
    )
    SELECT entity_id, MAX(score) AS score
    FROM scored
    GROUP BY entity_id
    ORDER BY score DESC
    LIMIT :limit
    """
)

_SERIES_SQL = text(
    """
    WITH scored AS (
        SELECT s.id AS entity_id,
               GREATEST(
                   ts_rank_cd(to_tsvector('simple', coalesce(s.title, '') || ' ' || coalesce(s.original_title, '')), plainto_tsquery('simple', :q)),
                   similarity(s.title, :q)
               ) AS score
        FROM series s
        WHERE to_tsvector('simple', coalesce(s.title, '') || ' ' || coalesce(s.original_title, '')) @@ plainto_tsquery('simple', :q)
           OR s.title % :q
        UNION ALL
        SELECT sat.series_id AS entity_id, similarity(sat.title, :q) AS score
        FROM series_alternate_titles sat
        WHERE sat.title % :q
    )
    SELECT entity_id, MAX(score) AS score
    FROM scored
    GROUP BY entity_id
    ORDER BY score DESC
    LIMIT :limit
    """
)

_PEOPLE_SQL = text(
    """
    SELECT p.id AS entity_id,
           GREATEST(
               ts_rank_cd(to_tsvector('simple', p.name), plainto_tsquery('simple', :q)),
               similarity(p.name, :q)
           ) AS score
    FROM people p
    WHERE to_tsvector('simple', p.name) @@ plainto_tsquery('simple', :q)
       OR p.name % :q
    ORDER BY score DESC
    LIMIT :limit
    """
)

_COLLECTIONS_SQL = text(
    """
    SELECT c.id AS entity_id, similarity(c.name, :q) AS score
    FROM collections c
    WHERE c.name % :q OR to_tsvector('simple', c.name) @@ plainto_tsquery('simple', :q)
    ORDER BY score DESC
    LIMIT :limit
    """
)

_COMPANIES_SQL = text(
    """
    SELECT co.id AS entity_id, similarity(co.name, :q) AS score
    FROM companies co
    WHERE co.name % :q OR to_tsvector('simple', co.name) @@ plainto_tsquery('simple', :q)
    ORDER BY score DESC
    LIMIT :limit
    """
)


class PostgreSQLSearchProvider(SearchProvider):
    async def search(
        self, session: AsyncSession, query: str, *, limit_per_type: int = 10
    ) -> SearchResultIDs:
        params = {"q": query, "limit": limit_per_type}
        movies = await session.execute(_MOVIE_SQL, params)
        series = await session.execute(_SERIES_SQL, params)
        people = await session.execute(_PEOPLE_SQL, params)
        collections = await session.execute(_COLLECTIONS_SQL, params)
        companies = await session.execute(_COMPANIES_SQL, params)

        def to_matches(rows) -> list[SearchMatch]:
            return [
                SearchMatch(entity_id=row.entity_id, score=float(row.score or 0.0))
                for row in rows
            ]

        return SearchResultIDs(
            movies=to_matches(movies),
            series=to_matches(series),
            people=to_matches(people),
            collections=to_matches(collections),
            companies=to_matches(companies),
        )