mirror of
https://github.com/zulip/zulip.git
synced 2025-11-04 14:03:30 +00:00
8acfa17fe6
This is adds foreign keys to the corresponding Recipient object in the UserProfile on Stream tables, a denormalization intended to improve performance as this is a common query. In the migration for setting the field correctly for existing users, we do a direct SQL query (because Django 1.11 doesn't provide any good method for doing it properly in bulk using the ORM.). A consequence of this change to the model is that a bit of code needs to be added to the functions responsible for creating new users (to set the field after the Recipient object gets created). Fortunately, there's only a few code paths for doing that. Also an adjustment is needed in the import system - this introduces a circular relation between Recipient and UserProfile. The field cannot be set until the Recipient objects have been created, but UserProfiles need to be created before their corresponding Recipients. We deal with this by first importing UserProfiles same way as before, but we leave the personal_recipient field uninitialized. After creating the Recipient objects, we call a function to set the field for all the imported users in bulk. A similar change is made for managing Stream objects.
141 lines
7.0 KiB
Python
141 lines
7.0 KiB
Python
from django.db.models import Model
|
|
|
|
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
|
|
|
|
from zerver.lib.initial_password import initial_password
|
|
from zerver.models import Realm, Stream, UserProfile, \
|
|
Subscription, Recipient, RealmAuditLog
|
|
from zerver.lib.create_user import create_user_profile
|
|
|
|
def bulk_create_users(realm: Realm,
|
|
users_raw: Set[Tuple[str, 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])
|
|
|
|
# If we have a different email_address_visibility mode, the code
|
|
# below doesn't have the logic to set user_profile.email properly.
|
|
assert realm.email_address_visibility == Realm.EMAIL_ADDRESS_VISIBILITY_EVERYONE
|
|
|
|
# Now create user_profiles
|
|
profiles_to_create = [] # type: List[UserProfile]
|
|
for (email, full_name, short_name, active) in users:
|
|
profile = create_user_profile(realm, email,
|
|
initial_password(email), active, bot_type,
|
|
full_name, short_name, bot_owner, False, tos_version,
|
|
timezone, tutorial_status=UserProfile.TUTORIAL_FINISHED,
|
|
enter_sends=True)
|
|
profiles_to_create.append(profile)
|
|
UserProfile.objects.bulk_create(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])
|
|
|
|
profiles_by_email = {} # type: Dict[str, UserProfile]
|
|
profiles_by_id = {} # type: Dict[int, UserProfile]
|
|
for profile in UserProfile.objects.select_related().filter(realm=realm):
|
|
profiles_by_email[profile.email] = profile
|
|
profiles_by_id[profile.id] = profile
|
|
|
|
recipients_to_create = [] # type: List[Recipient]
|
|
for (email, full_name, short_name, active) in users:
|
|
recipients_to_create.append(Recipient(type_id=profiles_by_email[email].id,
|
|
type=Recipient.PERSONAL))
|
|
Recipient.objects.bulk_create(recipients_to_create)
|
|
|
|
bulk_set_users_or_streams_recipient_fields(UserProfile, profiles_to_create, recipients_to_create)
|
|
|
|
recipients_by_email = {} # type: Dict[str, Recipient]
|
|
for recipient in recipients_to_create:
|
|
recipients_by_email[profiles_by_id[recipient.type_id].email] = recipient
|
|
|
|
subscriptions_to_create = [] # type: List[Subscription]
|
|
for (email, full_name, short_name, active) in users:
|
|
subscriptions_to_create.append(
|
|
Subscription(user_profile_id=profiles_by_email[email].id,
|
|
recipient=recipients_by_email[email]))
|
|
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 = dict((obj.id, obj) for obj in objects)
|
|
|
|
for recipient in recipients:
|
|
assert recipient.type == recipient_type
|
|
result = objects_dict.get(recipient.type_id)
|
|
if result is not None:
|
|
result.recipient = recipient
|
|
# TODO: Django 2.2 has a bulk_update method, so once we manage to migrate to that version,
|
|
# we take adventage of this, instead of calling save individually.
|
|
result.save(update_fields=['recipient'])
|
|
|
|
# This is only sed in populate_db, so doesn't realy need tests
|
|
def bulk_create_streams(realm: Realm,
|
|
stream_dict: Dict[str, Dict[str, Any]]) -> None: # nocoverage
|
|
existing_streams = frozenset([name.lower() for name in
|
|
Stream.objects.filter(realm=realm)
|
|
.values_list('name', flat=True)])
|
|
streams_to_create = [] # type: 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:
|
|
from zerver.lib.actions import render_stream_description
|
|
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),
|
|
is_announcement_only=options.get("is_announcement_only", False),
|
|
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 = [] # type: 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)
|