"""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), )