book_library_server/fastapi_book_server/app/services/common.py

import abc
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Generic, TypeVar, Union

import aioredis
from databases import Database
from fastapi_pagination.api import resolve_params
from fastapi_pagination.bases import AbstractParams, RawParams
import meilisearch
import orjson
from ormar import Model, QuerySet
from sqlalchemy import Table

from app.utils.pagination import Page, CustomPage
from core.config import env_config


T = TypeVar("T", bound=Model)


class BaseSearchService(Generic[T], abc.ABC):
    MODEL_CLASS: Optional[T] = None
    SELECT_RELATED: Optional[Union[list[str], str]] = None
    PREFETCH_RELATED: Optional[Union[list[str], str]] = None
    CUSTOM_CACHE_PREFIX: Optional[str] = None
    CACHE_TTL = 60 * 60

    @classmethod
    def get_params(cls) -> AbstractParams:
        return resolve_params()

    @classmethod
    def get_raw_params(cls) -> RawParams:
        return resolve_params().to_raw_params()

    @classmethod
    @property
    def model(cls) -> T:
        assert cls.MODEL_CLASS is not None, f"MODEL in {cls.__name__} don't set!"
        return cls.MODEL_CLASS

    @classmethod
    @property
    def table(cls) -> Table:
        return cls.model.Meta.table

    @classmethod
    @property
    def database(cls) -> Database:
        return cls.model.Meta.database

    @classmethod
    @property
    def cache_prefix(cls) -> str:
        return cls.CUSTOM_CACHE_PREFIX or cls.model.Meta.tablename

    @classmethod
    async def _get_object_ids(
        cls, query_data: str, allowed_langs: list[str]
    ) -> list[int]:
        ...

    @classmethod
    def get_cache_key(cls, query_data: str, allowed_langs: list[str]) -> str:
        model_class_name = cls.cache_prefix
        allowed_langs_part = ",".join(sorted(allowed_langs))
        return f"{model_class_name}_{query_data}_{allowed_langs_part}"

    @classmethod
    async def get_cached_ids(
        cls,
        query_data: str,
        allowed_langs: list[str],
        redis: aioredis.Redis,
    ) -> Optional[list[int]]:
        try:
            key = cls.get_cache_key(query_data, allowed_langs)
            data = await redis.get(key)

            if data is None:
                return None

            return orjson.loads(data)
        except aioredis.RedisError as e:
            print(e)
            return None

    @classmethod
    async def cache_object_ids(
        cls,
        query_data: str,
        allowed_langs: list[str],
        object_ids: list[int],
        redis: aioredis.Redis,
    ):
        try:
            key = cls.get_cache_key(query_data, allowed_langs)
            await redis.set(key, orjson.dumps(object_ids), ex=cls.CACHE_TTL)
        except aioredis.RedisError as e:
            print(e)

    @classmethod
    async def get_objects(
        cls,
        query_data: str,
        redis: aioredis.Redis,
        allowed_langs: list[str],
    ) -> tuple[int, list[T]]:
        params = cls.get_raw_params()

        cached_object_ids = await cls.get_cached_ids(query_data, allowed_langs, redis)

        if cached_object_ids is None:
            object_ids = await cls._get_object_ids(query_data, allowed_langs)
            await cls.cache_object_ids(query_data, allowed_langs, object_ids, redis)
        else:
            object_ids = cached_object_ids

        limited_object_ids = object_ids[params.offset : params.offset + params.limit]

        queryset: QuerySet[T] = cls.model.objects  # type: ignore

        if cls.PREFETCH_RELATED is not None:
            queryset = queryset.prefetch_related(cls.PREFETCH_RELATED)

        if cls.SELECT_RELATED:
            queryset = queryset.select_related(cls.SELECT_RELATED)

        return len(object_ids), await queryset.filter(id__in=limited_object_ids).all()

    @classmethod
    async def get(
        cls, query: str, redis: aioredis.Redis, allowed_langs: list[str]
    ) -> Page[T]:
        params = cls.get_params()

        total, objects = await cls.get_objects(query, redis, allowed_langs)

        return CustomPage.create(items=objects, total=total, params=params)


class TRGMSearchService(BaseSearchService[T]):
    GET_OBJECT_IDS_QUERY: Optional[str] = None

    @classmethod
    @property
    def object_ids_query(cls) -> str:
        assert (
            cls.GET_OBJECT_IDS_QUERY is not None
        ), f"GET_OBJECT_IDS_QUERY in {cls.__name__} don't set!"
        return cls.GET_OBJECT_IDS_QUERY

    @classmethod
    async def _get_object_ids(
        cls, query_data: str, allowed_langs: list[str]
    ) -> list[int]:
        row = await cls.database.fetch_one(
            cls.object_ids_query, {"query": query_data, "langs": allowed_langs}
        )

        if row is None:
            raise ValueError("Something is wrong!")

        return row["array"]


class MeiliSearchService(BaseSearchService[T]):
    MS_INDEX_NAME: Optional[str] = None
    MS_INDEX_LANG_KEY: Optional[str] = None

    _executor = ThreadPoolExecutor(4)

    @classmethod
    @property
    def lang_key(cls) -> str:
        assert cls.MS_INDEX_LANG_KEY is not None, f"MODEL in {cls.__name__} don't set!"
        return cls.MS_INDEX_LANG_KEY

    @classmethod
    @property
    def index_name(cls) -> str:
        assert cls.MS_INDEX_NAME is not None, f"MODEL in {cls.__name__} don't set!"
        return cls.MS_INDEX_NAME

    @classmethod
    def get_allowed_langs_filter(cls, allowed_langs: list[str]) -> list[list[str]]:
        return [[f"{cls.lang_key} = {lang}" for lang in allowed_langs]]

    @classmethod
    def make_request(
        cls, query: str, allowed_langs_filter: list[list[str]], offset: int
    ):
        client = meilisearch.Client(env_config.MEILI_HOST, env_config.MEILI_MASTER_KEY)
        index = client.index(cls.index_name)

        result = index.search(
            query,
            {
                "filter": allowed_langs_filter,
                "offset": offset,
                "limit": 630,
                "attributesToRetrieve": ["id"],
            },
        )

        total: int = result["nbHits"]
        ids: list[int] = [r["id"] for r in result["hits"][:total]]

        return ids

    @classmethod
    async def _get_object_ids(
        cls, query_data: str, allowed_langs: list[str]
    ) -> list[int]:
        params = cls.get_raw_params()

        allowed_langs_filter = cls.get_allowed_langs_filter(allowed_langs)

        return await asyncio.get_event_loop().run_in_executor(
            cls._executor,
            cls.make_request,
            query_data,
            allowed_langs_filter,
            params.offset,
        )


class GetRandomService(Generic[T]):
    MODEL_CLASS: Optional[T] = None
    GET_RANDOM_OBJECT_ID_QUERY: Optional[str] = None

    @classmethod
    @property
    def model(cls) -> T:
        assert cls.MODEL_CLASS is not None, f"MODEL in {cls.__name__} don't set!"
        return cls.MODEL_CLASS

    @classmethod
    @property
    def database(cls) -> Database:
        return cls.model.Meta.database

    @classmethod
    @property
    def random_object_id_query(cls) -> str:
        assert (
            cls.GET_RANDOM_OBJECT_ID_QUERY is not None
        ), f"GET_OBJECT_IDS_QUERY in {cls.__name__} don't set!"
        return cls.GET_RANDOM_OBJECT_ID_QUERY

    @classmethod
    async def get_random_id(cls, allowed_langs: list[str]) -> int:
        return await cls.database.fetch_val(
            cls.random_object_id_query, {"langs": allowed_langs}
        )