streams: Optimize computing users with metadata access.

This commit updates code to optimize computing users who have
metadata access via permission groups so that we do not have
to do DB query for each stream to get recursive members for
the groups having permissions.

zerver/actions/streams.py

@@ -43,6 +43,7 @@ from zerver.lib.streams import (
     get_stream_permission_policy_name,
     get_stream_post_policy_value_based_on_group_setting,
     get_user_ids_with_metadata_access_via_permission_groups,
+    get_users_dict_with_metadata_access_to_streams_via_permission_groups,
     render_stream_description,
     send_stream_creation_event,
     send_stream_deletion_event,
@@ -492,6 +493,7 @@ def send_stream_creation_events_for_previously_inaccessible_streams(
     stream_dict: dict[int, Stream],
     altered_user_dict: dict[int, set[int]],
     altered_guests: set[int],
+    users_with_metadata_access_via_permission_groups: dict[int, set[int]] | None = None,
 ) -> None:
     stream_ids = set(altered_user_dict.keys())
     recent_traffic = get_streams_traffic(stream_ids, realm)
@@ -504,6 +506,7 @@ def send_stream_creation_events_for_previously_inaccessible_streams(
 
         notify_user_ids = []
         if not stream.is_public():
+            assert users_with_metadata_access_via_permission_groups is not None
             # Users newly added to invite-only streams
             # need a `create` notification.  The former, because
             # they need the stream to exist before
@@ -511,13 +514,10 @@ def send_stream_creation_events_for_previously_inaccessible_streams(
             # they can manage the new stream.
             # Realm admins already have all created private streams.
             realm_admin_ids = {user.id for user in realm.get_admin_users_and_bots()}
-            user_ids_with_metadata_access_via_permission_groups = (
-                get_user_ids_with_metadata_access_via_permission_groups(stream)
-            )
             notify_user_ids = list(
                 stream_users_ids
                 - realm_admin_ids
-                - user_ids_with_metadata_access_via_permission_groups
+                - users_with_metadata_access_via_permission_groups[stream.id]
             )
         elif not stream.is_web_public:
             # Guese users need a `create` notification for
@@ -805,9 +805,19 @@ def bulk_add_subscriptions(
 
     new_streams = [stream_dict[stream_id] for stream_id in altered_user_dict]
 
+    private_streams = [stream for stream in new_streams if not stream.is_public()]
+    users_with_metadata_access_via_permission_groups = None
+    if private_streams:
+        users_with_metadata_access_via_permission_groups = (
+            get_users_dict_with_metadata_access_to_streams_via_permission_groups(
+                private_streams, realm.id
+            )
+        )
+
     subscriber_peer_info = bulk_get_subscriber_peer_info(
         realm=realm,
         streams=new_streams,
+        users_with_metadata_access_via_permission_groups=users_with_metadata_access_via_permission_groups,
     )
 
     # We now send several types of events to notify browsers.  The
@@ -820,6 +830,7 @@ def bulk_add_subscriptions(
             stream_dict=stream_dict,
             altered_user_dict=altered_user_dict,
             altered_guests=altered_guests,
+            users_with_metadata_access_via_permission_groups=users_with_metadata_access_via_permission_groups,
         )
 
         send_subscription_add_events(

zerver/lib/streams.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from collections.abc import Collection, Iterable
 from dataclasses import dataclass
 from datetime import datetime, timedelta
@@ -25,9 +26,9 @@ from zerver.lib.timestamp import datetime_to_timestamp
 from zerver.lib.types import AnonymousSettingGroupDict, APIStreamDict
 from zerver.lib.user_groups import (
     get_recursive_group_members,
-    get_recursive_group_members_union_for_groups,
     get_recursive_membership_groups,
     get_role_based_system_groups_dict,
+    get_root_id_annotated_recursive_subgroups_for_groups,
     user_has_permission_for_group_setting,
 )
 from zerver.models import (
@@ -185,20 +186,59 @@ def get_default_values_for_stream_permission_group_settings(
     return group_setting_values
 
 
-def get_user_ids_with_metadata_access_via_permission_groups(stream: Stream) -> set[int]:
+def get_users_dict_with_metadata_access_to_streams_via_permission_groups(
-    return set(
+    streams: list[Stream],
-        get_recursive_group_members_union_for_groups(
+    realm_id: int,
-            [stream.can_add_subscribers_group_id, stream.can_administer_channel_group_id]
+) -> dict[int, set[int]]:
+    can_administer_group_ids = {stream.can_administer_channel_group_id for stream in streams}
+    can_add_subscriber_group_ids = {stream.can_add_subscribers_group_id for stream in streams}
+
+    all_permission_group_ids = list(can_administer_group_ids | can_add_subscriber_group_ids)
+
+    recursive_subgroups = get_root_id_annotated_recursive_subgroups_for_groups(
+        all_permission_group_ids, realm_id
+    )
+    subgroup_root_id_dict = {}
+    all_subgroup_ids = set()
+    for group in recursive_subgroups:
+        subgroup_root_id_dict[group.id] = group.root_id  # type: ignore[attr-defined]  # root_id is an annotated field.
+        all_subgroup_ids.add(group.id)
+
+    group_members = (
+        UserGroupMembership.objects.filter(
+            user_group_id__in=list(all_subgroup_ids), user_profile__is_active=True
         )
         .exclude(
             # allow_everyone_group=False is false for both
             # can_add_subscribers_group and
             # can_administer_channel_group, so guest users cannot
             # exercise these permission to get metadata access.
-            role=UserProfile.ROLE_GUEST
+            user_profile__role=UserProfile.ROLE_GUEST
         )
-        .values_list("id", flat=True)
+        .values_list("user_group_id", "user_profile_id")
     )
+    group_members_dict = defaultdict(set)
+    for user_group_id, user_profile_id in group_members:
+        root_id = subgroup_root_id_dict[user_group_id]
+        group_members_dict[root_id].add(user_profile_id)
+
+    users_with_metadata_access_dict = defaultdict(set)
+    for stream in streams:
+        users_with_metadata_access_dict[stream.id] = (
+            group_members_dict[stream.can_administer_channel_group_id]
+            | group_members_dict[stream.can_add_subscribers_group_id]
+        )
+
+    return users_with_metadata_access_dict
+
+
+def get_user_ids_with_metadata_access_via_permission_groups(stream: Stream) -> set[int]:
+    users_with_metadata_access_dict = (
+        get_users_dict_with_metadata_access_to_streams_via_permission_groups(
+            [stream], stream.realm_id
+        )
+    )
+    return users_with_metadata_access_dict[stream.id]
 
 
 @transaction.atomic(savepoint=False)

zerver/lib/subscription_info.py

@@ -23,7 +23,7 @@ from zerver.lib.streams import (
     get_group_setting_value_dict_for_streams,
     get_setting_values_for_group_settings,
     get_stream_post_policy_value_based_on_group_setting,
-    get_user_ids_with_metadata_access_via_permission_groups,
+    get_users_dict_with_metadata_access_to_streams_via_permission_groups,
     get_web_public_streams_queryset,
     has_metadata_access_to_channel_via_groups,
     subscribed_to_stream,
@@ -521,6 +521,7 @@ def get_subscribers_query(
 def bulk_get_subscriber_peer_info(
     realm: Realm,
     streams: Collection[Stream] | QuerySet[Stream],
+    users_with_metadata_access_via_permission_groups: dict[int, set[int]] | None = None,
 ) -> SubscriberPeerInfo:
     """
     Glossary:
@@ -553,18 +554,22 @@ def bulk_get_subscriber_peer_info(
     if private_streams:
         realm_admin_ids = {user.id for user in realm.get_admin_users_and_bots()}
 
+        if users_with_metadata_access_via_permission_groups is None:
+            users_with_metadata_access_via_permission_groups = (
+                get_users_dict_with_metadata_access_to_streams_via_permission_groups(
+                    list(private_streams), realm.id
+                )
+            )
+
         for stream in private_streams:
             # Realm admins can see all private stream
             # subscribers.
             subscribed_user_ids = stream_user_ids.get(stream.id, set())
             subscribed_ids[stream.id] = subscribed_user_ids
-            user_ids_with_metadata_access_via_permission_groups = (
-                get_user_ids_with_metadata_access_via_permission_groups(stream)
-            )
             private_peer_dict[stream.id] = (
                 subscribed_user_ids
                 | realm_admin_ids
-                | user_ids_with_metadata_access_via_permission_groups
+                | users_with_metadata_access_via_permission_groups[stream.id]
             )
 
     for stream_id in public_stream_ids:

zerver/lib/user_groups.py

@@ -792,7 +792,7 @@ def get_root_id_annotated_recursive_subgroups_for_groups(
     # each group root_id and annotates it with that group.
 
     cte = With.recursive(
-        lambda cte: NamedUserGroup.objects.filter(id__in=user_group_ids, realm=realm_id)
+        lambda cte: UserGroup.objects.filter(id__in=user_group_ids, realm=realm_id)
         .values(group_id=F("id"), root_id=F("id"))
         .union(
             cte.join(NamedUserGroup, direct_supergroups=cte.col.group_id).values(
@@ -801,9 +801,7 @@ def get_root_id_annotated_recursive_subgroups_for_groups(
         )
     )
     recursive_subgroups = (
-        cte.join(NamedUserGroup, id=cte.col.group_id)
+        cte.join(UserGroup, id=cte.col.group_id).with_cte(cte).annotate(root_id=cte.col.root_id)
-        .with_cte(cte)
-        .annotate(root_id=cte.col.root_id)
     )
 
     return recursive_subgroups

zerver/tests/test_subs.py

@@ -3391,7 +3391,7 @@ class StreamAdminTest(ZulipTestCase):
         are on.
         """
         result = self.attempt_unsubscribe_of_principal(
-            query_count=16,
+            query_count=17,
             target_users=[self.example_user("cordelia")],
             is_realm_admin=True,
             is_subbed=True,
@@ -3408,7 +3408,7 @@ class StreamAdminTest(ZulipTestCase):
         streams you aren't on.
         """
         result = self.attempt_unsubscribe_of_principal(
-            query_count=16,
+            query_count=17,
             target_users=[self.example_user("cordelia")],
             is_realm_admin=True,
             is_subbed=False,
@@ -6133,15 +6133,17 @@ class SubscriptionAPITest(ZulipTestCase):
             private, "can_administer_channel_group", user6_group, acting_user=user6
         )
 
-        user7_group = self.create_or_update_anonymous_group_for_setting([user7], [])
+        user7_and_guests_group = self.create_or_update_anonymous_group_for_setting(
+            [user7, guest], []
+        )
         do_change_stream_group_based_setting(
-            private, "can_add_subscribers_group", user7_group, acting_user=user7
+            private, "can_add_subscribers_group", user7_and_guests_group, acting_user=user7
         )
 
         # Sends 3 peer-remove events, 2 unsubscribe events
         # and 2 stream delete events for private streams.
         with (
-            self.assert_database_query_count(19),
+            self.assert_database_query_count(20),
             self.assert_memcached_count(3),
             self.capture_send_event_calls(expected_num_events=7) as events,
         ):
@@ -6230,7 +6232,7 @@ class SubscriptionAPITest(ZulipTestCase):
         # Verify that peer_event events are never sent in Zephyr
         # realm. This does generate stream creation events from
         # send_stream_creation_events_for_previously_inaccessible_streams.
-        with self.assert_database_query_count(num_streams * 2 + 15):
+        with self.assert_database_query_count(num_streams + 17):
             with self.capture_send_event_calls(expected_num_events=num_streams + 1) as events:
                 self.subscribe_via_post(
                     mit_user,
@@ -6697,7 +6699,7 @@ class SubscriptionAPITest(ZulipTestCase):
             )
 
         # Test creating private stream.
-        with self.assert_database_query_count(49):
+        with self.assert_database_query_count(50):
             self.subscribe_via_post(
                 self.test_user,
                 [new_streams[1]],