From 3dc82eed5757b095c747d8b180da63d52d39a85d Mon Sep 17 00:00:00 2001 From: Sahil Batra Date: Mon, 19 Aug 2024 19:43:06 +0530 Subject: [PATCH] user_groups: Refactor code to compute user group objects. This commit refactors code in user_groups_in_realm_serialized such that we do not prefetch "can_mention_group__direct_members" and "can_mention_group__direct_subgroups" using prefetch_related and instead fetch members and subgroups for all groups in separate queries and then use that data to find the members and subgroups of the group used for that setting. This change helps us in avoiding two prefetch queries for each setting when we add more group settings. --- zerver/lib/user_groups.py | 93 ++++++++++++++++++------------- zerver/tests/test_event_system.py | 4 +- zerver/tests/test_home.py | 6 +- zerver/tests/test_user_groups.py | 14 ++--- 4 files changed, 67 insertions(+), 50 deletions(-) diff --git a/zerver/lib/user_groups.py b/zerver/lib/user_groups.py index 65512e63fb..9440427681 100644 --- a/zerver/lib/user_groups.py +++ b/zerver/lib/user_groups.py @@ -1,3 +1,4 @@ +from collections import defaultdict from collections.abc import Collection, Iterable, Iterator, Mapping from contextlib import contextmanager from dataclasses import asdict, dataclass @@ -5,7 +6,7 @@ from typing import TypedDict from django.conf import settings from django.db import connection, transaction -from django.db.models import F, Prefetch, QuerySet +from django.db.models import F, QuerySet from django.utils.timezone import now as timezone_now from django.utils.translation import gettext as _ from django_cte import With @@ -345,60 +346,76 @@ def get_group_setting_value_for_api( ) +def get_setting_value_for_user_group_object( + setting_value_group: UserGroup, + direct_members_dict: dict[int, list[int]], + direct_subgroups_dict: dict[int, list[int]], +) -> int | AnonymousSettingGroupDict: + if hasattr(setting_value_group, "named_user_group"): + return setting_value_group.id + + direct_members = [] + if setting_value_group.id in direct_members_dict: + direct_members = direct_members_dict[setting_value_group.id] + + direct_subgroups = [] + if setting_value_group.id in direct_subgroups_dict: + direct_subgroups = direct_subgroups_dict[setting_value_group.id] + + return AnonymousSettingGroupDict( + direct_members=direct_members, + direct_subgroups=direct_subgroups, + ) + + def user_groups_in_realm_serialized(realm: Realm) -> list[UserGroupDict]: """This function is used in do_events_register code path so this code should be performant. We need to do 2 database queries because Django's ORM doesn't properly support the left join between UserGroup and UserGroupMembership that we need. """ - realm_groups = ( - NamedUserGroup.objects.select_related( - "can_mention_group", "can_mention_group__named_user_group" - ) - # Using prefetch_related results in one query for each field. This is fine - # for now but would be problematic when more settings would be added. - # - # TODO: We should refactor it such that we only make two queries - one - # to fetch all the realm groups and the second to fetch all the groups - # that they point to and then set the setting fields for realm groups - # accordingly in Python. - .prefetch_related( - Prefetch("can_mention_group__direct_members", queryset=UserProfile.objects.only("id")), - Prefetch( - "can_mention_group__direct_subgroups", queryset=NamedUserGroup.objects.only("id") - ), - ) - .filter(realm=realm) + realm_groups = NamedUserGroup.objects.select_related( + "can_mention_group", "can_mention_group__named_user_group" + ).filter(realm=realm) + + membership = UserGroupMembership.objects.filter(user_group__realm=realm).values_list( + "user_group_id", "user_profile_id" ) + group_membership = GroupGroupMembership.objects.filter(subgroup__realm=realm).values_list( + "subgroup_id", "supergroup_id" + ) + + group_members = defaultdict(list) + for user_group_id, user_profile_id in membership: + group_members[user_group_id].append(user_profile_id) + + group_subgroups = defaultdict(list) + for subgroup_id, supergroup_id in group_membership: + group_subgroups[supergroup_id].append(subgroup_id) + group_dicts: dict[int, UserGroupDict] = {} for user_group in realm_groups: + direct_member_ids = [] + if user_group.id in group_members: + direct_member_ids = group_members[user_group.id] + + direct_subgroup_ids = [] + if user_group.id in group_subgroups: + direct_subgroup_ids = group_subgroups[user_group.id] + group_dicts[user_group.id] = dict( id=user_group.id, name=user_group.name, description=user_group.description, - members=[], - direct_subgroup_ids=[], + members=direct_member_ids, + direct_subgroup_ids=direct_subgroup_ids, is_system_group=user_group.is_system_group, - can_mention_group=get_group_setting_value_for_api(user_group.can_mention_group), + can_mention_group=get_setting_value_for_user_group_object( + user_group.can_mention_group, group_members, group_subgroups + ), ) - membership = ( - UserGroupMembership.objects.filter(user_group__realm=realm) - .exclude(user_group__named_user_group=None) - .values_list("user_group_id", "user_profile_id") - ) - for user_group_id, user_profile_id in membership: - group_dicts[user_group_id]["members"].append(user_profile_id) - - group_membership = ( - GroupGroupMembership.objects.filter(subgroup__realm=realm) - .exclude(supergroup__named_user_group=None) - .values_list("subgroup_id", "supergroup_id") - ) - for subgroup_id, supergroup_id in group_membership: - group_dicts[supergroup_id]["direct_subgroup_ids"].append(subgroup_id) - for group_dict in group_dicts.values(): group_dict["members"] = sorted(group_dict["members"]) group_dict["direct_subgroup_ids"] = sorted(group_dict["direct_subgroup_ids"]) diff --git a/zerver/tests/test_event_system.py b/zerver/tests/test_event_system.py index 582f404b85..2110151980 100644 --- a/zerver/tests/test_event_system.py +++ b/zerver/tests/test_event_system.py @@ -1176,7 +1176,7 @@ class FetchQueriesTest(ZulipTestCase): realm = get_realm_with_settings(realm_id=user.realm_id) with ( - self.assert_database_query_count(41), + self.assert_database_query_count(39), mock.patch("zerver.lib.events.always_want") as want_mock, ): fetch_initial_state_data(user, realm=realm) @@ -1202,7 +1202,7 @@ class FetchQueriesTest(ZulipTestCase): realm_linkifiers=0, realm_playgrounds=1, realm_user=4, - realm_user_groups=5, + realm_user_groups=3, realm_user_settings_defaults=1, recent_private_conversations=1, scheduled_messages=1, diff --git a/zerver/tests/test_home.py b/zerver/tests/test_home.py index 8d8b458779..a6d3fbf7be 100644 --- a/zerver/tests/test_home.py +++ b/zerver/tests/test_home.py @@ -268,7 +268,7 @@ class HomeTest(ZulipTestCase): # Verify succeeds once logged-in with ( - self.assert_database_query_count(52), + self.assert_database_query_count(50), patch("zerver.lib.cache.cache_set") as cache_mock, ): result = self._get_home_page(stream="Denmark") @@ -572,7 +572,7 @@ class HomeTest(ZulipTestCase): # Verify number of queries for Realm admin isn't much higher than for normal users. self.login("iago") with ( - self.assert_database_query_count(52), + self.assert_database_query_count(50), patch("zerver.lib.cache.cache_set") as cache_mock, ): result = self._get_home_page() @@ -604,7 +604,7 @@ class HomeTest(ZulipTestCase): self._get_home_page() # Then for the second page load, measure the number of queries. - with self.assert_database_query_count(47): + with self.assert_database_query_count(45): result = self._get_home_page() # Do a sanity check that our new streams were in the payload. diff --git a/zerver/tests/test_user_groups.py b/zerver/tests/test_user_groups.py index 1168541bce..80103f5302 100644 --- a/zerver/tests/test_user_groups.py +++ b/zerver/tests/test_user_groups.py @@ -106,8 +106,9 @@ class UserGroupTestCase(ZulipTestCase): self.assertEqual(user_groups[9]["can_mention_group"], everyone_group.id) othello = self.example_user("othello") + hamletcharacters_group = NamedUserGroup.objects.get(name="hamletcharacters", realm=realm) setting_group = self.create_or_update_anonymous_group_for_setting( - [othello], [admins_system_group] + [othello], [admins_system_group, hamletcharacters_group] ) new_user_group = check_add_user_group( realm, @@ -121,12 +122,11 @@ class UserGroupTestCase(ZulipTestCase): self.assertEqual(user_groups[10]["name"], "newgroup2") self.assertEqual(user_groups[10]["description"], "") self.assertEqual(user_groups[10]["members"], [othello.id]) - self.assertEqual( - user_groups[10]["can_mention_group"], - AnonymousSettingGroupDict( - direct_members=[othello.id], - direct_subgroups=[admins_system_group.id], - ), + assert isinstance(user_groups[10]["can_mention_group"], AnonymousSettingGroupDict) + self.assertEqual(user_groups[10]["can_mention_group"].direct_members, [othello.id]) + self.assertCountEqual( + user_groups[10]["can_mention_group"].direct_subgroups, + [admins_system_group.id, hamletcharacters_group.id], ) def test_get_direct_user_groups(self) -> None: