paulmataruso/zulip
1
0
Fork 0
mirror of https://github.com/zulip/zulip.git synced 2025-11-06 06:53:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

create stream: Bulk conversion of principals to UserProfiles.

Browse Source 
Previously, this logic did the database queries to look up UserProfile
objects in a loop.

Fixes #21820.

Significantly improves Stream creation time and also unsusbcribing users.
Tested stream creation with 10k stream subscribers:

- before: 127 seconds ~2 mins
- after: 17 seconds ~0.3 min

Add a test case for user unsubscribing themself.
This commit is contained in:
bedo
2024-08-03 18:05:16 +03:00
committed by Tim Abbott
parent 3d58a7ec04
commit 00c9f36434
3 changed files with 147 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
zerver/lib/users.py
View File

@@ -10,6 +10,7 @@ from typing import Any, TypedDict
from django.conf import settings
from django.core.exceptions import ValidationError
from django.db.models import Q, QuerySet
from django.db.models.functions import Upper
from django.utils.translation import gettext as _
from django_otp.middleware import is_verified
from typing_extensions import NotRequired
@@ -347,6 +348,74 @@ def access_user_by_email(
return access_user_common(target, user_profile, allow_deactivated, allow_bots, for_admin)
def bulk_access_users_by_email(
emails: list[str],
*,
acting_user: UserProfile,
allow_deactivated: bool = False,
allow_bots: bool = False,
for_admin: bool,
) -> set[UserProfile]:
# We upper-case the email addresses ourselves here, because
# `email__iexact__in=emails` is not supported by Django.
target_emails_upper = [email.strip().upper() for email in emails]
users = (
UserProfile.objects.annotate(email_upper=Upper("email"))
.select_related(
"realm",
"realm__can_access_all_users_group",
"realm__can_access_all_users_group__named_user_group",
"realm__direct_message_initiator_group",
"realm__direct_message_initiator_group__named_user_group",
"realm__direct_message_permission_group",
"realm__direct_message_permission_group__named_user_group",
"bot_owner",
)
.filter(email_upper__in=target_emails_upper, realm=acting_user.realm)
)
valid_emails_upper = {user_profile.email_upper for user_profile in users}
all_users_exist = all(email in valid_emails_upper for email in target_emails_upper)
if not all_users_exist:
raise JsonableError(_("No such user"))
return {
access_user_common(user_profile, acting_user, allow_deactivated, allow_bots, for_admin)
for user_profile in users
}
def bulk_access_users_by_id(
user_ids: list[int],
*,
acting_user: UserProfile,
allow_deactivated: bool = False,
allow_bots: bool = False,
for_admin: bool,
) -> set[UserProfile]:
users = UserProfile.objects.select_related(
"realm",
"realm__can_access_all_users_group",
"realm__can_access_all_users_group__named_user_group",
"realm__direct_message_initiator_group",
"realm__direct_message_initiator_group__named_user_group",
"realm__direct_message_permission_group",
"realm__direct_message_permission_group__named_user_group",
"bot_owner",
).filter(id__in=user_ids, realm=acting_user.realm)
valid_user_ids = {user_profile.id for user_profile in users}
all_users_exist = all(user_id in valid_user_ids for user_id in user_ids)
if not all_users_exist:
raise JsonableError(_("No such user"))
return {
access_user_common(user_profile, acting_user, allow_deactivated, allow_bots, for_admin)
for user_profile in users
}
class Account(TypedDict):
realm_name: str
realm_id: int

58
zerver/tests/test_subs.py
View File

@@ -2684,7 +2684,7 @@ class StreamAdminTest(ZulipTestCase):
for name in ["cordelia", "prospero", "iago", "hamlet", "outgoing_webhook_bot"]
]
result = self.attempt_unsubscribe_of_principal(
query_count=28,
query_count=24,
cache_count=8,
target_users=target_users,
is_realm_admin=True,
@@ -2759,7 +2759,7 @@ class StreamAdminTest(ZulipTestCase):
def test_admin_remove_multiple_users_from_stream_legacy_emails(self) -> None:
result = self.attempt_unsubscribe_of_principal(
query_count=20,
query_count=19,
target_users=[self.example_user("cordelia"), self.example_user("prospero")],
is_realm_admin=True,
is_subbed=True,
@@ -2773,7 +2773,7 @@ class StreamAdminTest(ZulipTestCase):
def test_remove_unsubbed_user_along_with_subbed(self) -> None:
result = self.attempt_unsubscribe_of_principal(
query_count=17,
query_count=16,
target_users=[self.example_user("cordelia"), self.example_user("iago")],
is_realm_admin=True,
is_subbed=True,
@@ -2909,6 +2909,43 @@ class StreamAdminTest(ZulipTestCase):
)
self.assert_json_error(result, "No such user", status_code=400)
def test_user_unsubscribe_theirself(self) -> None:
"""
User trying to unsubscribe theirself from the stream, where
principals has the id of the acting_user performing the
unsubscribe action.
"""
admin = self.example_user("iago")
self.login_user(admin)
self.assertTrue(admin.is_realm_admin)
stream_name = "hümbüǵ"
self.make_stream(stream_name)
self.subscribe(admin, stream_name)
# unsubscribing when subscribed.
result = self.client_delete(
"/json/users/me/subscriptions",
{
"subscriptions": orjson.dumps([stream_name]).decode(),
"principals": orjson.dumps([admin.id]).decode(),
},
)
json = self.assert_json_success(result)
self.assert_length(json["removed"], 1)
# unsubscribing after already being unsubscribed.
result = self.client_delete(
"/json/users/me/subscriptions",
{
"subscriptions": orjson.dumps([stream_name]).decode(),
"principals": orjson.dumps([admin.id]).decode(),
},
)
json = self.assert_json_success(result)
self.assert_length(json["not_removed"], 1)
class DefaultStreamTest(ZulipTestCase):
def get_default_stream_names(self, realm: Realm) -> set[str]:
@@ -4686,7 +4723,7 @@ class SubscriptionAPITest(ZulipTestCase):
streams_to_sub = ["multi_user_stream"]
with (
self.capture_send_event_calls(expected_num_events=5) as events,
self.assert_database_query_count(38),
self.assert_database_query_count(37),
):
self.common_subscribe_to_streams(
self.test_user,
@@ -5157,11 +5194,8 @@ class SubscriptionAPITest(ZulipTestCase):
test_user_ids = [user.id for user in test_users]
# The only known O(N) behavior here is that we call
# principal_to_user_profile for each of our users, but it
# should be cached.
with (
self.assert_database_query_count(21),
self.assert_database_query_count(16),
self.assert_memcached_count(3),
mock.patch("zerver.views.streams.send_messages_for_new_subscribers"),
):
@@ -5517,7 +5551,7 @@ class SubscriptionAPITest(ZulipTestCase):
]
# Test creating a public stream when realm does not have a notification stream.
with self.assert_database_query_count(38):
with self.assert_database_query_count(37):
self.common_subscribe_to_streams(
self.test_user,
[new_streams[0]],
@@ -5525,7 +5559,7 @@ class SubscriptionAPITest(ZulipTestCase):
)
# Test creating private stream.
with self.assert_database_query_count(40):
with self.assert_database_query_count(39):
self.common_subscribe_to_streams(
self.test_user,
[new_streams[1]],
@@ -5537,7 +5571,7 @@ class SubscriptionAPITest(ZulipTestCase):
new_stream_announcements_stream = get_stream(self.streams[0], self.test_realm)
self.test_realm.new_stream_announcements_stream_id = new_stream_announcements_stream.id
self.test_realm.save()
with self.assert_database_query_count(49):
with self.assert_database_query_count(48):
self.common_subscribe_to_streams(
self.test_user,
[new_streams[2]],
@@ -6015,7 +6049,7 @@ class GetSubscribersTest(ZulipTestCase):
polonius.id,
]
with self.assert_database_query_count(46):
with self.assert_database_query_count(43):
self.common_subscribe_to_streams(
self.user_profile,
streams,

48
zerver/views/streams.py
View File

@@ -75,20 +75,37 @@ from zerver.lib.topic import (
from zerver.lib.typed_endpoint import ApiParamConfig, PathOnly, typed_endpoint
from zerver.lib.typed_endpoint_validators import check_color, check_int_in_validator
from zerver.lib.user_groups import access_user_group_for_setting
from zerver.lib.users import access_user_by_email, access_user_by_id
from zerver.lib.users import bulk_access_users_by_email, bulk_access_users_by_id
from zerver.lib.utils import assert_is_not_none
from zerver.models import NamedUserGroup, Realm, Stream, UserProfile
from zerver.models.users import get_system_bot
def principal_to_user_profile(agent: UserProfile, principal: str | int) -> UserProfile:
if isinstance(principal, str):
return access_user_by_email(
agent, principal, allow_deactivated=False, allow_bots=True, for_admin=False
def bulk_principals_to_user_profiles(
principals: list[str] | list[int],
acting_user: UserProfile,
) -> set[UserProfile]:
# Since principals is guaranteed to be non-empty and to have the same type of elements,
# the following if/else is safe and enough.
# principals are user emails.
if isinstance(principals[0], str):
return bulk_access_users_by_email(
principals, # type: ignore[arg-type] # principals guaranteed to be list[str] only.
acting_user=acting_user,
allow_deactivated=False,
allow_bots=True,
for_admin=False,
)
# principals are user ids.
else:
return access_user_by_id(
agent, principal, allow_deactivated=False, allow_bots=True, for_admin=False
return bulk_access_users_by_id(
principals, # type: ignore[arg-type] # principals guaranteed to be list[int] only.
acting_user=acting_user,
allow_deactivated=False,
allow_bots=True,
for_admin=False,
)
@@ -473,12 +490,11 @@ def remove_subscriptions_backend(
unsubscribing_others = False
if principals:
people_to_unsub = set()
for principal in principals:
target_user = principal_to_user_profile(user_profile, principal)
people_to_unsub.add(target_user)
if not user_directly_controls_user(user_profile, target_user):
unsubscribing_others = True
people_to_unsub = bulk_principals_to_user_profiles(principals, user_profile)
unsubscribing_others = any(
not user_directly_controls_user(user_profile, target) for target in people_to_unsub
)
else:
people_to_unsub = {user_profile}
@@ -551,6 +567,7 @@ def add_subscriptions_backend(
realm = user_profile.realm
stream_dicts = []
color_map = {}
# UserProfile ids or emails.
if principals is None:
principals = []
@@ -607,9 +624,8 @@ def add_subscriptions_backend(
# Guest users case will not be handled here as it will
# be handled by the decorator above.
raise JsonableError(_("Insufficient permission"))
subscribers = {
principal_to_user_profile(user_profile, principal) for principal in principals
}
subscribers = bulk_principals_to_user_profiles(principals, user_profile)
else:
subscribers = {user_profile}