From b81e00ff3d445023832ccd21bb7ea148c53dd105 Mon Sep 17 00:00:00 2001 From: Roland Date: Mon, 24 Oct 2022 15:34:13 -0600 Subject: [PATCH] Perserving ts_rank when searching. --- modules/videos/test/test_lib.py | 8 ++++---- modules/videos/video/api.py | 2 +- modules/videos/video/lib.py | 2 +- wrolpi/db.py | 9 ++++----- wrolpi/files/api.py | 4 ++-- wrolpi/files/lib.py | 14 +++++++++++--- wrolpi/files/models.py | 2 +- wrolpi/files/test/test_lib.py | 12 ++++++------ 8 files changed, 30 insertions(+), 23 deletions(-) diff --git a/modules/videos/test/test_lib.py b/modules/videos/test/test_lib.py index 6bcc7c85..3e059c42 100644 --- a/modules/videos/test/test_lib.py +++ b/modules/videos/test/test_lib.py @@ -6,7 +6,7 @@ import sqlalchemy from modules.videos.common import apply_info_json from modules.videos.lib import parse_video_file_name, validate_video, get_statistics from modules.videos.models import Video -from modules.videos.video.lib import video_search +from modules.videos.video.lib import search_videos from wrolpi.files import lib as files_lib from wrolpi.files.models import File @@ -27,19 +27,19 @@ def test_search_censored_videos(test_session, simple_channel, video_factory): # All source_ids are in the entries. set_entries(map(str, range(50))) - videos, total = video_search(filters=['censored'], order_by='id', limit=20) + videos, total = search_videos(filters=['censored'], order_by='id', limit=20) assert [i['video']['source_id'] for i in videos] == [] assert total == 0 # First 5 are censored. set_entries(map(str, range(5, 50))) - videos, total = video_search(filters=['censored'], order_by='id', limit=20) + videos, total = search_videos(filters=['censored'], order_by='id', limit=20) assert [i['video']['source_id'] for i in videos] == [str(i) for i in range(5)] assert total == 5 # First 25 are censored. set_entries(map(str, range(25, 50))) - videos, total = video_search(filters=['censored'], order_by='id', limit=20) + videos, total = search_videos(filters=['censored'], order_by='id', limit=20) assert [i['video']['source_id'] for i in videos] == [str(i) for i in range(20)] assert total == 25 diff --git a/modules/videos/video/api.py b/modules/videos/video/api.py index 8d9c9b14..f2100ad7 100644 --- a/modules/videos/video/api.py +++ b/modules/videos/video/api.py @@ -39,7 +39,7 @@ async def search(_: Request, body: schema.VideoSearchRequest): if body.order_by not in lib.VIDEO_ORDERS: raise InvalidOrderBy('Invalid order by') - files, videos_total = lib.video_search( + files, videos_total = lib.search_videos( body.search_str, body.offset, body.limit, diff --git a/modules/videos/video/lib.py b/modules/videos/video/lib.py index 652a1662..a785c9b1 100644 --- a/modules/videos/video/lib.py +++ b/modules/videos/video/lib.py @@ -76,7 +76,7 @@ DEFAULT_VIDEO_ORDER = 'rank' VIDEO_QUERY_LIMIT = 24 -def video_search( +def search_videos( search_str: str = None, offset: int = None, limit: int = VIDEO_QUERY_LIMIT, diff --git a/wrolpi/db.py b/wrolpi/db.py index 2259839c..35d672b9 100644 --- a/wrolpi/db.py +++ b/wrolpi/db.py @@ -1,4 +1,3 @@ -import logging from contextlib import contextmanager from functools import wraps from typing import ContextManager, Tuple, List, Union, Type @@ -160,10 +159,10 @@ def optional_session(commit: Union[callable, bool] = False): @optional_session -def get_ranked_models(ranked_ids: List, model: Type[Base], session: Session = None) -> List[Base]: - """Get all objects whose ids are in the `ranked_ids`, order them by their position in `ranked_ids`.""" +def get_ranked_models(ranked_primary_keys: List, model: Type[Base], session: Session = None) -> List[Base]: + """Get all objects whose primary keys are in the `ranked_primary_keys`, preserve their order.""" pkey = inspect(model).primary_key[0] pkey_name = pkey.name - results = list(session.query(model).filter(pkey.in_(ranked_ids)).all()) - results = sorted(results, key=lambda i: ranked_ids.index(getattr(i, pkey_name))) + results = list(session.query(model).filter(pkey.in_(ranked_primary_keys)).all()) + results = sorted(results, key=lambda i: ranked_primary_keys.index(getattr(i, pkey_name))) return results diff --git a/wrolpi/files/api.py b/wrolpi/files/api.py index 62dc71b0..65e256de 100644 --- a/wrolpi/files/api.py +++ b/wrolpi/files/api.py @@ -98,6 +98,6 @@ async def refresh_directory(_: Request, body: schema.DirectoryRefreshRequest): body=schema.FilesSearchRequest, ) @validate(schema.FilesSearchRequest) -async def search_files(_: Request, body: schema.FilesSearchRequest): - files, total = lib.file_search(body.search_str, body.limit, body.offset, body.mimetype, body.model) +async def post_search_files(_: Request, body: schema.FilesSearchRequest): + files, total = lib.search_files(body.search_str, body.limit, body.offset, body.mimetype, body.model) return json_response(dict(files=files, totals=dict(files=total))) diff --git a/wrolpi/files/lib.py b/wrolpi/files/lib.py index f5e2f9e1..0e974461 100644 --- a/wrolpi/files/lib.py +++ b/wrolpi/files/lib.py @@ -34,7 +34,7 @@ except ImportError: logger = logger.getChild(__name__) -__all__ = ['list_files', 'delete_file', 'split_path_stem_and_suffix', 'refresh_files', 'file_search', 'get_mimetype'] +__all__ = ['list_files', 'delete_file', 'split_path_stem_and_suffix', 'refresh_files', 'search_files', 'get_mimetype'] def filter_parent_directories(directories: List[Path]) -> List[Path]: @@ -276,8 +276,8 @@ async def refresh_directory_files_recursively(directory: Union[pathlib.Path, str refresh_logger.info(f'Done refreshing files in {directory}') -def file_search(search_str: str, limit: int, offset: int, mimetype: str = None, model: str = None) -> Tuple[ - List[dict], int]: +def search_files(search_str: str, limit: int, offset: int, mimetype: str = None, model: str = None) -> \ + Tuple[List[dict], int]: """ Search the Files table. Order the returned Files by their rank if `search_str` is provided. Return all files if `search_str` is empty. @@ -342,6 +342,11 @@ def handle_search_results(statement: str, params): return [], 0 total = results[0]['total'] if results else 0 ranked_paths = [pathlib.Path(i['path']) for i in results] + try: + ranks = [i['ts_rank'] for i in results] + except KeyError: + # No `ts_rank`, probably not searching `file.textsearch`. + ranks = [] with get_db_session() as session: results = get_ranked_models(ranked_paths, File, session=session) @@ -352,6 +357,9 @@ def handle_search_results(statement: str, params): if PYTEST: raise results = [i.__json__() for i in results] + # Preserve the ts_ranks, if any. + for idx, rank in enumerate(ranks): + results[idx]['ts_rank'] = rank return results, total diff --git a/wrolpi/files/models.py b/wrolpi/files/models.py index 1612e2c4..945ed407 100644 --- a/wrolpi/files/models.py +++ b/wrolpi/files/models.py @@ -49,7 +49,7 @@ class File(ModelHelper, Base): path = str(self.path.relative_to(get_media_directory())) return f'' - def __json__(self): + def __json__(self) -> dict: path = self.path.relative_to(get_media_directory()) d = dict( associated=self.associated, diff --git a/wrolpi/files/test/test_lib.py b/wrolpi/files/test/test_lib.py index 27ecdd19..8b842f36 100644 --- a/wrolpi/files/test/test_lib.py +++ b/wrolpi/files/test/test_lib.py @@ -326,17 +326,17 @@ async def test_files_indexer(test_session, make_files_structure): assert video_file.suffix == '.mp4' # File are indexed by their titles and contents. - files, total = lib.file_search('file', 10, 0) + files, total = lib.search_files('file', 10, 0) assert total == 7, 'All files contain "file" in their file name. The associated video file is hidden.' - files, total = lib.file_search('image', 10, 0) + files, total = lib.search_files('image', 10, 0) assert total == 1 and files[0]['title'] == 'an image file.jpeg', 'The image file title contains "image".' - files, total = lib.file_search('contents', 10, 0) + files, total = lib.search_files('contents', 10, 0) assert total == 1 and files[0]['title'] == 'a text file.txt', 'The text file contains "contents".' - files, total = lib.file_search('video', 10, 0) + files, total = lib.search_files('video', 10, 0) assert total == 1 and {i['title'] for i in files} == {'a video file.mp4'}, 'The video file contains "video".' - files, total = lib.file_search('yawn', 10, 0) + files, total = lib.search_files('yawn', 10, 0) assert total == 1 and files[0]['title'] == 'a video file.mp4', 'The video file captions contain "yawn".' - files, total = lib.file_search('bunny', 10, 0) + files, total = lib.search_files('bunny', 10, 0) assert total == 3 and {i['title'] for i in files} == {'a zip file.zip', 'a bzip file.bzip', 'a gzip file.gzip'}, \ 'The zip files contain a file with "bunny" in the title.'