create stream: Bulk conversion of principals to UserProfiles.

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.

zerver/lib/users.py

@@ -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

zerver/tests/test_subs.py

@@ -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,

zerver/views/streams.py

@@ -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:
+def bulk_principals_to_user_profiles(
-    if isinstance(principal, str):
+    principals: list[str] | list[int],
-        return access_user_by_email(
+    acting_user: UserProfile,
-            agent, principal, allow_deactivated=False, allow_bots=True, for_admin=False
+) -> 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(
+        return bulk_access_users_by_id(
-            agent, principal, allow_deactivated=False, allow_bots=True, for_admin=False
+            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()
+        people_to_unsub = bulk_principals_to_user_profiles(principals, user_profile)
-        for principal in principals:
+        unsubscribing_others = any(
-            target_user = principal_to_user_profile(user_profile, principal)
+            not user_directly_controls_user(user_profile, target) for target in people_to_unsub
-            people_to_unsub.add(target_user)
+        )
-            if not user_directly_controls_user(user_profile, target_user):
+
-                unsubscribing_others = True
     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 = {
+        subscribers = bulk_principals_to_user_profiles(principals, user_profile)
-            principal_to_user_profile(user_profile, principal) for principal in principals
+
-        }
     else:
         subscribers = {user_profile}