realm: Prefetch group settings only for "/register" response.

This commit updates code to prefetch realm group settings like
"can_create_public_channel_group" only when computing settings
for "/register" response by refetching the realm object with
select_related instead of fetching those settings in UserProfile
query.

This change is done because we do not need to prefetch these
settings for every UserProfile object and for most of the cases
where these settings are actually accessed, we can afford extra
query like when checking permission to create streams. But we
cannot afford one query extra for each setting when computing
these settings for "/register" response, so we re-fetch the
realm object with select_related leading to only one extra
query.

The query count changes in tests are -
- Query count increases by 1 when calling fetch_initial_state_data
for computing can_create_public_streams because Realm object from
UserProfile does not have prefetched setting fields.

- Query count increases by one in test_subs where streams are
created which is as expected due to the setting not being prefetched.

- Query count increases by 2 in tests in test_home.py where one
query is to refetch the realm object and one for computing
can_create_public_streams as mentioned above.
This commit is contained in:
Sahil Batra
2024-06-12 19:13:02 +05:30
committed by Tim Abbott
parent dea62a3fd9
commit e23af20079
6 changed files with 34 additions and 16 deletions

View File

@@ -94,6 +94,7 @@ from zerver.models.realms import (
EditTopicPolicyEnum,
get_corresponding_policy_value_for_group_setting,
get_realm_domains,
get_realm_with_settings,
)
from zerver.models.streams import get_default_stream_groups
from zerver.tornado.django_api import get_user_events, request_event_queue
@@ -1642,6 +1643,11 @@ def do_events_register(
else:
event_types_set = None
# Fetch the realm object again to prefetch all the
# group settings which support anonymous groups
# to avoid unnecessary DB queries.
realm = get_realm_with_settings(realm_id=realm.id)
if user_profile is None:
# TODO: Unify the two fetch_initial_state_data code paths.
assert client_gravatar is False

View File

@@ -1074,6 +1074,19 @@ def get_realm_by_id(realm_id: int) -> Realm:
return Realm.objects.get(id=realm_id)
def get_realm_with_settings(realm_id: int) -> Realm:
# Prefetch all the settings that can be set to anonymous groups.
# This also prefetches can_access_all_users_group setting,
# even when it cannot be set to anonymous groups because
# the setting is used when fetching users in the realm.
return Realm.objects.select_related(
"can_access_all_users_group",
"can_access_all_users_group__named_user_group",
"can_create_public_channel_group",
"can_create_public_channel_group__named_user_group",
).get(id=realm_id)
def require_unique_names(realm: Optional[Realm]) -> bool:
if realm is None:
# realm is None when a new realm is being created.

View File

@@ -882,8 +882,6 @@ def get_user_profile_by_id(user_profile_id: int) -> UserProfile:
"realm",
"realm__can_access_all_users_group",
"realm__can_access_all_users_group__named_user_group",
"realm__can_create_public_channel_group",
"realm__can_create_public_channel_group__named_user_group",
"bot_owner",
).get(id=user_profile_id)
@@ -905,8 +903,6 @@ def maybe_get_user_profile_by_api_key(api_key: str) -> Optional[UserProfile]:
"realm",
"realm__can_access_all_users_group",
"realm__can_access_all_users_group__named_user_group",
"realm__can_create_public_channel_group",
"realm__can_create_public_channel_group__named_user_group",
"bot_owner",
).get(api_key=api_key)
except UserProfile.DoesNotExist:
@@ -936,8 +932,6 @@ def get_user_by_delivery_email(email: str, realm: "Realm") -> UserProfile:
"realm",
"realm__can_access_all_users_group",
"realm__can_access_all_users_group__named_user_group",
"realm__can_create_public_channel_group",
"realm__can_create_public_channel_group__named_user_group",
"bot_owner",
).get(delivery_email__iexact=email.strip(), realm=realm)

View File

@@ -29,7 +29,7 @@ from zerver.lib.test_helpers import (
from zerver.lib.users import get_api_key, get_users_for_api
from zerver.models import CustomProfileField, UserMessage, UserPresence, UserProfile
from zerver.models.clients import get_client
from zerver.models.realms import get_realm
from zerver.models.realms import get_realm, get_realm_with_settings
from zerver.models.streams import get_stream
from zerver.models.users import get_system_bot
from zerver.tornado.event_queue import (
@@ -1165,9 +1165,14 @@ class FetchQueriesTest(ZulipTestCase):
self.login_user(user)
with self.assert_database_query_count(43):
# Fetch realm like it is done when calling fetch_initial_state_data
# in production to match the query counts with the actual query
# count in production.
realm = get_realm_with_settings(realm_id=user.realm_id)
with self.assert_database_query_count(44):
with mock.patch("zerver.lib.events.always_want") as want_mock:
fetch_initial_state_data(user, realm=user.realm)
fetch_initial_state_data(user, realm=realm)
expected_counts = dict(
alert_words=1,
@@ -1220,7 +1225,7 @@ class FetchQueriesTest(ZulipTestCase):
else:
event_types = [event_type]
fetch_initial_state_data(user, realm=user.realm, event_types=event_types)
fetch_initial_state_data(user, realm=realm, event_types=event_types)
class TestEventsRegisterAllPublicStreamsDefaults(ZulipTestCase):

View File

@@ -259,7 +259,7 @@ class HomeTest(ZulipTestCase):
self.client_post("/json/bots", bot_info)
# Verify succeeds once logged-in
with self.assert_database_query_count(53):
with self.assert_database_query_count(55):
with patch("zerver.lib.cache.cache_set") as cache_mock:
result = self._get_home_page(stream="Denmark")
self.check_rendered_logged_in_app(result)
@@ -564,7 +564,7 @@ class HomeTest(ZulipTestCase):
def test_num_queries_for_realm_admin(self) -> None:
# Verify number of queries for Realm admin isn't much higher than for normal users.
self.login("iago")
with self.assert_database_query_count(53):
with self.assert_database_query_count(55):
with patch("zerver.lib.cache.cache_set") as cache_mock:
result = self._get_home_page()
self.check_rendered_logged_in_app(result)
@@ -595,7 +595,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(48):
with self.assert_database_query_count(50):
result = self._get_home_page()
# Do a sanity check that our new streams were in the payload.

View File

@@ -4767,7 +4767,7 @@ class SubscriptionAPITest(ZulipTestCase):
realm = get_realm("zulip")
streams_to_sub = ["multi_user_stream"]
with self.capture_send_event_calls(expected_num_events=5) as events:
with self.assert_database_query_count(37):
with self.assert_database_query_count(38):
self.common_subscribe_to_streams(
self.test_user,
streams_to_sub,
@@ -5591,7 +5591,7 @@ class SubscriptionAPITest(ZulipTestCase):
]
# Test creating a public stream when realm does not have a notification stream.
with self.assert_database_query_count(37):
with self.assert_database_query_count(38):
self.common_subscribe_to_streams(
self.test_user,
[new_streams[0]],
@@ -5611,7 +5611,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(48):
with self.assert_database_query_count(49):
self.common_subscribe_to_streams(
self.test_user,
[new_streams[2]],