"""Full text search module. Implements full text search by using Postgre's capabilities and creating temporary tables containing keywords as ts_vectors. """ from enum import Enum from typing import Tuple from ereuse_devicehub.db import db class Weight(Enum): """TS Rank weight as an Enum.""" A = 'A' B = 'B' C = 'C' D = 'D' class Search: """Methods for building queries with Postgre's Full text search. Based on `Rachid Belaid's post `_ and `Code for America's post `. """ LANG = 'english' @staticmethod def match(column: db.Column, search: str, lang=LANG): """Query that matches a TSVECTOR column with search words.""" return column.op('@@')(db.func.websearch_to_tsquery(lang, search)) @staticmethod def rank(column: db.Column, search: str, lang=LANG): """Query that ranks a TSVECTOR column with search words.""" return db.func.ts_rank(column, db.func.websearch_to_tsquery(lang, search)) @staticmethod def _vectorize(col: db.Column, weight: Weight = Weight.D, lang=LANG): return db.func.setweight(db.func.to_tsvector(lang, db.func.coalesce(col, '')), weight.name) @classmethod def vectorize(cls, *cols_with_weights: Tuple[db.Column, Weight], lang=LANG): """Produces a query that takes one ore more columns and their respective weights, and generates one big TSVECTOR. This method takes care of `null` column values. """ first, rest = cols_with_weights[0], cols_with_weights[1:] tokens = cls._vectorize(*first, lang=lang) for unit in rest: tokens = tokens.concat(cls._vectorize(*unit, lang=lang)) return tokens