zulip/zerver/lib/bulk_create.py

import random
from collections import defaultdict
from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union

from django.db.models import Model

from zerver.lib.create_user import create_user_profile, get_display_email_address
from zerver.lib.initial_password import initial_password
from zerver.lib.streams import render_stream_description
from zerver.models import (
    Message,
    Reaction,
    Realm,
    RealmAuditLog,
    Recipient,
    Stream,
    Subscription,
    UserMessage,
    UserProfile,
)


def bulk_create_users(realm: Realm,
                      users_raw: Set[Tuple[str, str, bool]],
                      bot_type: Optional[int]=None,
                      bot_owner: Optional[UserProfile]=None,
                      tos_version: Optional[str]=None,
                      timezone: str="") -> None:
    """
    Creates and saves a UserProfile with the given email.
    Has some code based off of UserManage.create_user, but doesn't .save()
    """
    existing_users = frozenset(UserProfile.objects.filter(
        realm=realm).values_list('email', flat=True))
    users = sorted(user_raw for user_raw in users_raw if user_raw[0] not in existing_users)

    # Now create user_profiles
    profiles_to_create: List[UserProfile] = []
    for (email, full_name, active) in users:
        profile = create_user_profile(realm, email,
                                      initial_password(email), active, bot_type,
                                      full_name, bot_owner, False, tos_version,
                                      timezone, tutorial_status=UserProfile.TUTORIAL_FINISHED,
                                      enter_sends=True)
        profiles_to_create.append(profile)

    if realm.email_address_visibility == Realm.EMAIL_ADDRESS_VISIBILITY_EVERYONE:
        UserProfile.objects.bulk_create(profiles_to_create)
    else:
        for user_profile in profiles_to_create:
            user_profile.email = user_profile.delivery_email

        UserProfile.objects.bulk_create(profiles_to_create)

        for user_profile in profiles_to_create:
            user_profile.email = get_display_email_address(user_profile, realm)
        UserProfile.objects.bulk_update(profiles_to_create, ['email'])

    user_ids = {user.id for user in profiles_to_create}

    RealmAuditLog.objects.bulk_create(
        RealmAuditLog(realm=realm, modified_user=profile_,
                      event_type=RealmAuditLog.USER_CREATED, event_time=profile_.date_joined)
        for profile_ in profiles_to_create)

    recipients_to_create: List[Recipient] = []
    for user_id in user_ids:
        recipient = Recipient(type_id=user_id, type=Recipient.PERSONAL)
        recipients_to_create.append(recipient)

    Recipient.objects.bulk_create(recipients_to_create)

    bulk_set_users_or_streams_recipient_fields(UserProfile, profiles_to_create, recipients_to_create)

    recipients_by_user_id: Dict[int, Recipient] = {}
    for recipient in recipients_to_create:
        recipients_by_user_id[recipient.type_id] = recipient

    subscriptions_to_create: List[Subscription] = []
    for user_id in user_ids:
        recipient = recipients_by_user_id[user_id]
        subscription = Subscription(user_profile_id=user_id, recipient=recipient)
        subscriptions_to_create.append(subscription)

    Subscription.objects.bulk_create(subscriptions_to_create)

def bulk_set_users_or_streams_recipient_fields(model: Model,
                                               objects: Union[Iterable[UserProfile], Iterable[Stream]],
                                               recipients: Optional[Iterable[Recipient]]=None) -> None:
    assert model in [UserProfile, Stream]
    for obj in objects:
        assert isinstance(obj, model)

    if model == UserProfile:
        recipient_type = Recipient.PERSONAL
    elif model == Stream:
        recipient_type = Recipient.STREAM

    if recipients is None:
        object_ids = [obj.id for obj in objects]
        recipients = Recipient.objects.filter(type=recipient_type, type_id__in=object_ids)

    objects_dict = {obj.id: obj for obj in objects}

    objects_to_update = set()
    for recipient in recipients:
        assert recipient.type == recipient_type
        result = objects_dict.get(recipient.type_id)
        if result is not None:
            result.recipient = recipient
            objects_to_update.add(result)
    model.objects.bulk_update(objects_to_update, ['recipient'])

# This is only sed in populate_db, so doesn't really need tests
def bulk_create_streams(realm: Realm,
                        stream_dict: Dict[str, Dict[str, Any]]) -> None:  # nocoverage
    existing_streams = {
        name.lower()
        for name in Stream.objects.filter(realm=realm).values_list('name', flat=True)
    }
    streams_to_create: List[Stream] = []
    for name, options in stream_dict.items():
        if 'history_public_to_subscribers' not in options:
            options['history_public_to_subscribers'] = (
                not options.get("invite_only", False) and not realm.is_zephyr_mirror_realm)
        if name.lower() not in existing_streams:
            streams_to_create.append(
                Stream(
                    realm=realm,
                    name=name,
                    description=options["description"],
                    rendered_description=render_stream_description(options["description"]),
                    invite_only=options.get("invite_only", False),
                    stream_post_policy=options.get("stream_post_policy",
                                                   Stream.STREAM_POST_POLICY_EVERYONE),
                    history_public_to_subscribers=options["history_public_to_subscribers"],
                    is_web_public=options.get("is_web_public", False),
                    is_in_zephyr_realm=realm.is_zephyr_mirror_realm,
                ),
            )
    # Sort streams by name before creating them so that we can have a
    # reliable ordering of `stream_id` across different python versions.
    # This is required for test fixtures which contain `stream_id`. Prior
    # to python 3.3 hashes were not randomized but after a security fix
    # hash randomization was enabled in python 3.3 which made iteration
    # of dictionaries and sets completely unpredictable. Here the order
    # of elements while iterating `stream_dict` will be completely random
    # for python 3.3 and later versions.
    streams_to_create.sort(key=lambda x: x.name)
    Stream.objects.bulk_create(streams_to_create)

    recipients_to_create: List[Recipient] = []
    for stream in Stream.objects.filter(realm=realm).values('id', 'name'):
        if stream['name'].lower() not in existing_streams:
            recipients_to_create.append(Recipient(type_id=stream['id'],
                                                  type=Recipient.STREAM))
    Recipient.objects.bulk_create(recipients_to_create)

    bulk_set_users_or_streams_recipient_fields(Stream, streams_to_create, recipients_to_create)

DEFAULT_EMOJIS = [
    ('+1', '1f44d'),
    ('smiley', '1f603'),
    ('eyes', '1f440'),
    ('crying_cat_face', '1f63f'),
    ('arrow_up', '2b06'),
    ('confetti_ball', '1f38a'),
    ('hundred_points', '1f4af'),
]

def bulk_create_reactions(
        messages: Iterable[Message],
        users: Optional[List[UserProfile]] = None,
        emojis: Optional[List[Tuple[str, str]]] = None
) -> None:
    messages = list(messages)
    if not emojis:
        emojis = DEFAULT_EMOJIS
    emojis = list(emojis)

    reactions: List[Reaction] = []
    for message in messages:
        reactions.extend(_add_random_reactions_to_message(
            message, emojis, users))
    Reaction.objects.bulk_create(reactions)

def _add_random_reactions_to_message(
        message: Message,
        emojis: List[Tuple[str, str]],
        users: Optional[List[UserProfile]] = None,
        prob_reaction: float = 0.075,
        prob_upvote: float = 0.5,
        prob_repeat: float = 0.5) -> List[Reaction]:
    '''Randomly add emoji reactions to each message from a list.

    Algorithm:

    Give the message at least one reaction with probability `prob_reaction`.
    Once the first reaction is added, have another user upvote it with probability
    `prob_upvote`, provided there is another recipient of the message left to upvote.
    Repeat the process for a different emoji with probability `prob_repeat`.

    If the number of emojis or users is small, there is a chance the above process
    will produce multiple reactions with the same user and emoji, so group the
    reactions by emoji code and user profile and then return one reaction from
    each group.
    '''
    for p in (prob_reaction, prob_repeat, prob_upvote):
        # Prevent p=1 since for prob_repeat and prob_upvote, this will
        # lead to an infinite loop.
        if p >= 1 or p < 0:
            raise ValueError('Probability argument must be between 0 and 1.')

    # Avoid performing database queries if there will be no reactions.
    compute_next_reaction: bool = random.random() < prob_reaction
    if not compute_next_reaction:
        return []

    if users is None:
        users = []
    user_ids: Sequence[int] = [user.id for user in users]
    if not user_ids:
        user_ids = UserMessage.objects.filter(message=message) \
            .values_list("user_profile_id", flat=True)
        if not user_ids:
            return []

    emojis = list(emojis)

    reactions = []
    while compute_next_reaction:
        # We do this O(users) operation only if we've decided to do a
        # reaction, to avoid performance issues with large numbers of
        # users.
        users_available = set(user_ids)

        (emoji_name, emoji_code) = random.choice(emojis)
        while True:
            # Handle corner case where all the users have reacted.
            if not users_available:
                break

            user_id = random.choice(list(users_available))
            reactions.append(Reaction(
                user_profile_id=user_id,
                message=message,
                emoji_name=emoji_name,
                emoji_code=emoji_code,
                reaction_type=Reaction.UNICODE_EMOJI
            ))
            users_available.remove(user_id)

            # Add an upvote with the defined probability.
            if not random.random() < prob_upvote:
                break

        # Repeat with a possibly different random emoji with the
        # defined probability.
        compute_next_reaction = random.random() < prob_repeat

    # Avoid returning duplicate reactions by deduplicating on
    # (user_profile_id, emoji_code).
    grouped_reactions = defaultdict(list)
    for reaction in reactions:
        k = (str(reaction.user_profile_id), str(reaction.emoji_code))
        grouped_reactions[k].append(reaction)
    return [reactions[0] for reactions in grouped_reactions.values()]