Mirror
/
zulip
mirror of https://github.com/zulip/zulip.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

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.
This commit is contained in:
Sahil Batra 2024-08-19 19:43:06 +05:30 committed by Tim Abbott
parent a892f06cb5
commit 3dc82eed57
4 changed files with 67 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

93
zerver/lib/user_groups.py
View File

@ -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"])

4
zerver/tests/test_event_system.py
View File

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

6
zerver/tests/test_home.py
View File

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

14
zerver/tests/test_user_groups.py
View File

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