upstream-zulip/zerver/lib/bulk_create.py

from collections.abc import Collection, Iterable
from typing import Any

from django.db.models import QuerySet
from django.utils.timezone import now as timezone_now

from zerver.lib.avatar import generate_and_upload_jdenticon_avatar
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 (
    get_default_values_for_stream_permission_group_settings,
    render_stream_description,
)
from zerver.models import (
    NamedUserGroup,
    Realm,
    RealmAuditLog,
    RealmUserDefault,
    Recipient,
    Stream,
    UserGroupMembership,
    UserProfile,
)
from zerver.models.groups import SystemGroups
from zerver.models.realm_audit_logs import AuditLogEventType
from zerver.models.streams import StreamTopicsPolicyEnum


def bulk_create_users(
    realm: Realm,
    users_raw: set[tuple[str, str, bool]],
    bot_type: int | None = None,
    bot_owner: UserProfile | None = None,
    tos_version: str | None = None,
    timezone: str = "",
    is_imported_stub: bool = False,
) -> 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)

    realm_user_default = RealmUserDefault.objects.get(realm=realm)
    if bot_type is None:
        email_address_visibility = realm_user_default.email_address_visibility
    else:
        # There is no privacy motivation for limiting access to bot email addresses,
        # so we hardcode them to EMAIL_ADDRESS_VISIBILITY_EVERYONE.
        email_address_visibility = UserProfile.EMAIL_ADDRESS_VISIBILITY_EVERYONE

    # 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,
            is_imported_stub=is_imported_stub,
            default_language=realm.default_language,
            email_address_visibility=email_address_visibility,
        )

        if bot_type is None:
            # This block simulates copy_default_settings from
            # zerver/lib/create_user.py.
            #
            # We cannot use 'copy_default_settings' directly here
            # because it calls '.save' after copying the settings, and
            # we are bulk creating the objects here instead.
            for settings_name in RealmUserDefault.property_types:
                if settings_name in ["default_language", "enable_login_emails"]:
                    continue
                value = getattr(realm_user_default, settings_name)
                setattr(profile, settings_name, value)
        profiles_to_create.append(profile)

    if email_address_visibility == UserProfile.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)
        UserProfile.objects.bulk_update(profiles_to_create, ["email"])

    for user_profile in profiles_to_create:
        generate_and_upload_jdenticon_avatar(user_profile, str(realm.uuid), future=False)

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

    full_members_system_group = NamedUserGroup.objects.get(
        name=SystemGroups.FULL_MEMBERS, realm_for_sharding=realm, is_system_group=True
    )
    members_system_group = NamedUserGroup.objects.get(
        name=SystemGroups.MEMBERS, realm_for_sharding=realm, is_system_group=True
    )
    group_memberships_to_create: list[UserGroupMembership] = []
    for user_profile in profiles_to_create:
        # All users are members since this function is only used to create bots
        # and test and development environment users.
        assert user_profile.role == UserProfile.ROLE_MEMBER
        group_memberships_to_create.append(
            UserGroupMembership(user_profile=user_profile, user_group=members_system_group)
        )
        if not user_profile.is_provisional_member:
            group_memberships_to_create.append(
                UserGroupMembership(user_profile=user_profile, user_group=full_members_system_group)
            )

    UserGroupMembership.objects.bulk_create(group_memberships_to_create)
    now = timezone_now()
    RealmAuditLog.objects.bulk_create(
        RealmAuditLog(
            realm=realm,
            modified_user=membership.user_profile,
            modified_user_group=membership.user_group.named_user_group,
            event_type=AuditLogEventType.USER_GROUP_DIRECT_USER_MEMBERSHIP_ADDED,
            event_time=now,
            acting_user=None,
        )
        for membership in group_memberships_to_create
    )


def bulk_set_stream_recipient_fields(
    streams: Collection[Stream] | QuerySet[Stream],
    recipients: Iterable[Recipient] | None = None,
) -> None:
    if recipients is None:
        stream_ids = [stream.id for stream in streams]
        recipients = Recipient.objects.filter(type=Recipient.STREAM, type_id__in=stream_ids)

    streams_dict = {stream.id: stream for stream in streams}

    streams_to_update = set()
    for recipient in recipients:
        assert recipient.type == Recipient.STREAM
        stream = streams_dict.get(recipient.type_id)
        if stream is not None:
            stream.recipient = recipient
            streams_to_update.add(stream)
    Stream.objects.bulk_update(streams_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():
        creator = options.get("creator", None)
        if name.lower() not in existing_streams:
            stream = Stream(
                realm=realm,
                name=name,
                description=options["description"],
                rendered_description=render_stream_description(options["description"], realm),
                invite_only=options.get("invite_only", False),
                history_public_to_subscribers=options.get("history_public_to_subscribers", True),
                is_web_public=options.get("is_web_public", False),
                creator=options.get("creator", None),
                folder_id=options.get("folder_id", None),
                topics_policy=options.get("topics_policy", StreamTopicsPolicyEnum.inherit.value),
                **get_default_values_for_stream_permission_group_settings(realm, creator),
            )
            if "can_send_message_group" in options:
                stream.can_send_message_group = options["can_send_message_group"]

            streams_to_create.append(stream)
    # 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 = [
        Recipient(type_id=stream["id"], type=Recipient.STREAM)
        for stream in Stream.objects.filter(realm=realm).values("id", "name")
        if stream["name"].lower() not in existing_streams
    ]
    Recipient.objects.bulk_create(recipients_to_create)

    bulk_set_stream_recipient_fields(streams_to_create, recipients_to_create)


def create_users(
    realm: Realm, name_list: Iterable[tuple[str, str]], bot_type: int | None = None
) -> None:
    user_set = {(email, full_name, True) for full_name, email in name_list}
    bulk_create_users(realm, user_set, bot_type)