zulip/zerver/lib/display_recipient.py

from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, TypedDict

from zerver.lib.cache import (
    bulk_cached_fetch,
    cache_with_key,
    display_recipient_bulk_get_users_by_id_cache_key,
    display_recipient_cache_key,
    transformed_bulk_cached_fetch,
)
from zerver.lib.types import DisplayRecipientT, UserDisplayRecipient
from zerver.models import Recipient, Stream, UserProfile, bulk_get_huddle_user_ids

if TYPE_CHECKING:
    from django.db.models.query import _QuerySet as ValuesQuerySet

display_recipient_fields = [
    "id",
    "email",
    "full_name",
    "is_mirror_dummy",
]


class TinyStreamResult(TypedDict):
    id: int
    name: str


def get_display_recipient_cache_key(
    recipient_id: int, recipient_type: int, recipient_type_id: Optional[int]
) -> str:
    return display_recipient_cache_key(recipient_id)


@cache_with_key(get_display_recipient_cache_key, timeout=3600 * 24 * 7)
def get_display_recipient_remote_cache(
    recipient_id: int, recipient_type: int, recipient_type_id: Optional[int]
) -> DisplayRecipientT:
    """
    returns: an appropriate object describing the recipient.  For a
    stream this will be the stream name as a string.  For a huddle or
    personal, it will be an array of dicts about each recipient.
    """
    if recipient_type == Recipient.STREAM:
        assert recipient_type_id is not None
        stream = Stream.objects.values("name").get(id=recipient_type_id)
        return stream["name"]

    # The main priority for ordering here is being deterministic.
    # Right now, we order by ID, which matches the ordering of user
    # names in the left sidebar.
    user_profile_list = (
        UserProfile.objects.filter(
            subscription__recipient_id=recipient_id,
        )
        .order_by("id")
        .values(*display_recipient_fields)
    )
    return list(user_profile_list)


def user_dict_id_fetcher(user_dict: UserDisplayRecipient) -> int:
    return user_dict["id"]


def bulk_get_user_profile_by_id(uids: List[int]) -> Dict[int, UserDisplayRecipient]:
    return bulk_cached_fetch(
        # Use a separate cache key to protect us from conflicts with
        # the get_user_profile_by_id cache.
        # (Since we fetch only several fields here)
        cache_key_function=display_recipient_bulk_get_users_by_id_cache_key,
        query_function=lambda ids: list(
            UserProfile.objects.filter(id__in=ids).values(*display_recipient_fields)
        ),
        object_ids=uids,
        id_fetcher=user_dict_id_fetcher,
    )


def bulk_fetch_display_recipients(
    recipient_tuples: Set[Tuple[int, int, int]],
) -> Dict[int, DisplayRecipientT]:
    """
    Takes set of tuples of the form (recipient_id, recipient_type, recipient_type_id)
    Returns dict mapping recipient_id to corresponding display_recipient
    """

    # Build dict mapping recipient id to (type, type_id) of the corresponding recipient:
    recipient_id_to_type_pair_dict = {
        recipient[0]: (recipient[1], recipient[2]) for recipient in recipient_tuples
    }
    # And the inverse mapping:
    type_pair_to_recipient_id_dict = {
        (recipient[1], recipient[2]): recipient[0] for recipient in recipient_tuples
    }

    stream_recipients = {
        recipient for recipient in recipient_tuples if recipient[1] == Recipient.STREAM
    }
    personal_and_huddle_recipients = recipient_tuples - stream_recipients

    def stream_query_function(
        recipient_ids: List[int],
    ) -> "ValuesQuerySet[Stream, TinyStreamResult]":
        stream_ids = [
            recipient_id_to_type_pair_dict[recipient_id][1] for recipient_id in recipient_ids
        ]
        return Stream.objects.filter(id__in=stream_ids).values("name", "id")

    def stream_id_fetcher(stream: TinyStreamResult) -> int:
        return type_pair_to_recipient_id_dict[(Recipient.STREAM, stream["id"])]

    def stream_cache_transformer(stream: TinyStreamResult) -> str:
        return stream["name"]

    # ItemT = Stream, CacheItemT = str (name), ObjKT = int (recipient_id)
    stream_display_recipients: Dict[int, str] = transformed_bulk_cached_fetch(
        cache_key_function=display_recipient_cache_key,
        query_function=stream_query_function,
        object_ids=[recipient[0] for recipient in stream_recipients],
        id_fetcher=stream_id_fetcher,
        cache_transformer=stream_cache_transformer,
    )

    # Now we have to create display_recipients for personal and huddle messages.
    # We do this via generic_bulk_cached_fetch, supplying appropriate functions to it.

    def personal_and_huddle_query_function(
        recipient_ids: List[int],
    ) -> List[Tuple[int, List[UserDisplayRecipient]]]:
        """
        Return a list of tuples of the form (recipient_id, [list of UserProfiles])
        where [list of UserProfiles] has users corresponding to the recipient,
        so the receiving userin Recipient.PERSONAL case,
        or in Personal.HUDDLE case - users in the huddle.
        This is a pretty hacky return value, but it needs to be in this form,
        for this function to work as the query_function in generic_bulk_cached_fetch.
        """

        recipients = [
            Recipient(
                id=recipient_id,
                type=recipient_id_to_type_pair_dict[recipient_id][0],
                type_id=recipient_id_to_type_pair_dict[recipient_id][1],
            )
            for recipient_id in recipient_ids
        ]

        # Find all user ids whose UserProfiles we will need to fetch:
        user_ids_to_fetch: Set[int] = set()
        huddle_user_ids: Dict[int, List[int]] = {}
        huddle_user_ids = bulk_get_huddle_user_ids(
            [recipient for recipient in recipients if recipient.type == Recipient.HUDDLE]
        )
        for recipient in recipients:
            if recipient.type == Recipient.PERSONAL:
                user_ids_to_fetch.add(recipient.type_id)
            else:
                user_ids_to_fetch = user_ids_to_fetch.union(huddle_user_ids[recipient.id])

        # Fetch the needed UserProfiles:
        user_profiles: Dict[int, UserDisplayRecipient] = bulk_get_user_profile_by_id(
            list(user_ids_to_fetch)
        )

        # Build the return value:
        result: List[Tuple[int, List[UserDisplayRecipient]]] = []
        for recipient in recipients:
            if recipient.type == Recipient.PERSONAL:
                result.append((recipient.id, [user_profiles[recipient.type_id]]))
            else:
                result.append(
                    (
                        recipient.id,
                        [user_profiles[user_id] for user_id in huddle_user_ids[recipient.id]],
                    )
                )

        return result

    def personal_and_huddle_cache_transformer(
        db_object: Tuple[int, List[UserDisplayRecipient]],
    ) -> List[UserDisplayRecipient]:
        """
        Takes an element of the list returned by the query_function, maps it to the final
        display_recipient list.
        """
        user_profile_list = db_object[1]
        display_recipient = user_profile_list

        return display_recipient

    def personal_and_huddle_id_fetcher(db_object: Tuple[int, List[UserDisplayRecipient]]) -> int:
        # db_object is a tuple, with recipient_id in the first position
        return db_object[0]

    # ItemT = Tuple[int, List[UserDisplayRecipient]] (recipient_id, list of corresponding users)
    # CacheItemT = List[UserDisplayRecipient] (display_recipient list)
    # ObjKT = int (recipient_id)
    personal_and_huddle_display_recipients: Dict[
        int, List[UserDisplayRecipient]
    ] = transformed_bulk_cached_fetch(
        cache_key_function=display_recipient_cache_key,
        query_function=personal_and_huddle_query_function,
        object_ids=[recipient[0] for recipient in personal_and_huddle_recipients],
        id_fetcher=personal_and_huddle_id_fetcher,
        cache_transformer=personal_and_huddle_cache_transformer,
    )

    # Glue the dicts together and return:
    return {**stream_display_recipients, **personal_and_huddle_display_recipients}