From e23af200798df506ba25b8539144036223930371 Mon Sep 17 00:00:00 2001 From: Sahil Batra Date: Wed, 12 Jun 2024 19:13:02 +0530 Subject: [PATCH] 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. --- zerver/lib/events.py | 6 ++++++ zerver/models/realms.py | 13 +++++++++++++ zerver/models/users.py | 6 ------ zerver/tests/test_event_system.py | 13 +++++++++---- zerver/tests/test_home.py | 6 +++--- zerver/tests/test_subs.py | 6 +++--- 6 files changed, 34 insertions(+), 16 deletions(-) diff --git a/zerver/lib/events.py b/zerver/lib/events.py index ba7f54714a..c80a19400f 100644 --- a/zerver/lib/events.py +++ b/zerver/lib/events.py @@ -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 diff --git a/zerver/models/realms.py b/zerver/models/realms.py index 4dd2549c06..9e7a9d2764 100644 --- a/zerver/models/realms.py +++ b/zerver/models/realms.py @@ -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. diff --git a/zerver/models/users.py b/zerver/models/users.py index d8623c49ff..468b7d526e 100644 --- a/zerver/models/users.py +++ b/zerver/models/users.py @@ -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) diff --git a/zerver/tests/test_event_system.py b/zerver/tests/test_event_system.py index ec59e6ba1a..41112f2bb6 100644 --- a/zerver/tests/test_event_system.py +++ b/zerver/tests/test_event_system.py @@ -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): diff --git a/zerver/tests/test_home.py b/zerver/tests/test_home.py index 06e72668bc..15a9fc4f56 100644 --- a/zerver/tests/test_home.py +++ b/zerver/tests/test_home.py @@ -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. diff --git a/zerver/tests/test_subs.py b/zerver/tests/test_subs.py index d93cd7df0b..90898ffe14 100644 --- a/zerver/tests/test_subs.py +++ b/zerver/tests/test_subs.py @@ -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]],