inbox/resolvers/search.py

import time
from typing import Any, Dict, List, Union

from resolvers.load import load_messages
from services.auth import login_required
from services.core import CacheStorage
from services.rediscache import redis
from services.schema import query


@query.field("search_recipients")
@login_required
async def search_recipients(_, info, text: str, limit: int = 50, offset: int = 0):
    result = set()

    # TODO: maybe redis scan?

    author_id = info.context["author_id"]

    existed_chats = await redis.execute("SMEMBERS", f"/chats_by_author/{author_id}")
    if isinstance(existed_chats, set):
        chats_list = list(existed_chats)
        for chat_id in chats_list[offset : (offset + limit)]:
            members_ids = await redis.execute("SMEMBERS", f"/chats/{chat_id}/members")
            if isinstance(members_ids, set):
                for member_id in members_ids:
                    author = CacheStorage.authors_by_id.get(str(member_id))
                    if author:
                        if author["name"].startswith(text):
                            result.add(author)

    more_amount = limit - len(result)
    if more_amount > 0:
        result.update(CacheStorage.authors[0:more_amount])
    return {"members": list(result), "error": None}


@query.field("search_messages")
@login_required
async def search_messages(
    _, info, by: Dict[str, Union[str, int]], limit: int, offset: int
) -> Dict[str, Union[List[Dict[str, Any]], None]]:
    messages_set = set()
    author_id = info.context["author_id"]
    lookup_chats = await redis.execute("SMEMBERS", f"chats_by_author/{author_id}")
    if isinstance(lookup_chats, set):
        # pre-filter lookup chats
        by_member = by.get("author")
        if by_member:
            lookup_chats = filter(
                lambda ca: by_member in ca["members"],
                list(lookup_chats),
            )

        # load the messages from lookup chats
        for c in lookup_chats:
            chat_id = c.decode()
            fltr = None
            now = int(time.time())
            if by_member:
                fltr = lambda mx: mx and mx["created_by"] == by_member  # noqa E731
            body_like = by.get("body") or ""
            if isinstance(body_like, str):
                fltr = lambda mx: mx and body_like in mx["body"]  # noqa E731
            days_ago = int(by.get("days") or "0")
            if days_ago:
                ts = days_ago * 24 * 60 * 60
                fltr = lambda mx: mx and now - mx["created_by"] < ts  # noqa E731
            if fltr:
                mmm = await load_messages(chat_id, limit, offset)
                if isinstance(mmm, list):
                    mmm = list(filter(fltr, mmm))
                    messages_set |= set(mmm)

    return {"messages": sorted(messages_set), "error": None}