From 343741f621789c20cbc69bb238a17d8026fc97eb Mon Sep 17 00:00:00 2001 From: Shubham Padia Date: Fri, 29 Nov 2024 13:26:08 +0000 Subject: [PATCH] stream: Fetch system groups in one query for group setting defaults. Right now, the number of queries has remained the same, but when we add more settings in the future, we won't be increasing the number of queries when iterating over stream permission group settings. --- zerver/lib/import_realm.py | 27 ++++++++++++++++++--------- zerver/lib/streams.py | 21 +++++++++++++-------- zerver/lib/user_groups.py | 8 ++------ zerver/views/streams.py | 8 +++++++- 4 files changed, 40 insertions(+), 24 deletions(-) diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py index e42271183e..c7df38ac3f 100644 --- a/zerver/lib/import_realm.py +++ b/zerver/lib/import_realm.py @@ -247,11 +247,15 @@ def fix_upload_links(data: TableData, message_table: TableName) -> None: ) -def fix_stream_permission_group_settings(data: TableData, realm: Realm) -> None: +def fix_stream_permission_group_settings( + data: TableData, system_groups_name_dict: dict[str, NamedUserGroup] +) -> None: table = get_db_table(Stream) for stream in data[table]: for setting_name in Stream.stream_permission_group_settings: - stream[setting_name] = get_stream_permission_default_group(setting_name, realm) + stream[setting_name] = get_stream_permission_default_group( + setting_name, system_groups_name_dict + ) def create_subscription_events(data: TableData, realm_id: int) -> None: @@ -1269,9 +1273,9 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int = 1) -> Rea # We expect Zulip server exports to contain these system groups, # this logic here is needed to handle the imports from other services. - role_system_groups_dict: dict[int, NamedUserGroup] | None = None + system_groups_name_dict: dict[str, NamedUserGroup] | None = None if "zerver_usergroup" not in data: - role_system_groups_dict = create_system_user_groups_for_realm(realm) + system_groups_name_dict = create_system_user_groups_for_realm(realm) # Email tokens will automatically be randomly generated when the # Stream objects are created by Django. @@ -1289,11 +1293,11 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int = 1) -> Rea creator_id = stream.pop("creator_id", None) stream_id_to_creator_id[stream["id"]] = creator_id - if role_system_groups_dict is not None: + if system_groups_name_dict is not None: # Because the system user groups are missing, we manually set up # the defaults for stream permission settings for all the # streams. - fix_stream_permission_group_settings(data, realm) + fix_stream_permission_group_settings(data, system_groups_name_dict) else: for setting_name in Stream.stream_permission_group_settings: re_map_foreign_keys(data, "zerver_stream", setting_name, related_table="usergroup") @@ -1540,8 +1544,8 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int = 1) -> Rea # We expect Zulip server exports to contain UserGroupMembership objects # for system groups, this logic here is needed to handle the imports from # other services. - if role_system_groups_dict is not None: - add_users_to_system_user_groups(realm, user_profiles, role_system_groups_dict) + if system_groups_name_dict is not None: + add_users_to_system_user_groups(realm, user_profiles, system_groups_name_dict) if "zerver_botstoragedata" in data: re_map_foreign_keys( @@ -2078,7 +2082,7 @@ def import_analytics_data(realm: Realm, import_dir: Path, crossrealm_user_ids: s def add_users_to_system_user_groups( realm: Realm, user_profiles: list[UserProfile], - role_system_groups_dict: dict[int, NamedUserGroup], + system_groups_name_dict: dict[str, NamedUserGroup], ) -> None: full_members_system_group = NamedUserGroup.objects.get( name=SystemGroups.FULL_MEMBERS, @@ -2086,6 +2090,11 @@ def add_users_to_system_user_groups( is_system_group=True, ) + role_system_groups_dict: dict[int, NamedUserGroup] = dict() + for role in NamedUserGroup.SYSTEM_USER_GROUP_ROLE_MAP: + group_name = NamedUserGroup.SYSTEM_USER_GROUP_ROLE_MAP[role]["name"] + role_system_groups_dict[role] = system_groups_name_dict[group_name] + usergroup_memberships = [] for user_profile in user_profiles: user_group = role_system_groups_dict[user_profile.role] diff --git a/zerver/lib/streams.py b/zerver/lib/streams.py index 3ebb4ba9b6..93c8d21c51 100644 --- a/zerver/lib/streams.py +++ b/zerver/lib/streams.py @@ -25,6 +25,7 @@ from zerver.lib.timestamp import datetime_to_timestamp from zerver.lib.types import AnonymousSettingGroupDict, APIStreamDict from zerver.lib.user_groups import ( get_group_setting_value_for_api, + get_role_based_system_groups_dict, user_has_permission_for_group_setting, ) from zerver.models import ( @@ -139,20 +140,19 @@ def send_stream_creation_event( send_event_on_commit(realm, event, user_ids) -def get_stream_permission_default_group(setting_name: str, realm: Realm) -> UserGroup: +def get_stream_permission_default_group( + setting_name: str, system_groups_name_dict: dict[str, NamedUserGroup] +) -> UserGroup: setting_default_name = Stream.stream_permission_group_settings[setting_name].default_group_name - return NamedUserGroup.objects.get( - name=setting_default_name, - realm=realm, - is_system_group=True, - ) + return system_groups_name_dict[setting_default_name] def get_default_group_setting_values(realm: Realm) -> dict[str, UserGroup]: group_setting_values = {} + system_groups_name_dict = get_role_based_system_groups_dict(realm) for setting_name in Stream.stream_permission_group_settings: group_setting_values[setting_name] = get_stream_permission_default_group( - setting_name, realm + setting_name, system_groups_name_dict ) return group_setting_values @@ -179,13 +179,18 @@ def create_stream_if_needed( group_setting_values = {} request_settings_dict = locals() + # We don't want to calculate this value if no default values are + # needed. + system_groups_name_dict = None for setting_name in Stream.stream_permission_group_settings: if setting_name not in request_settings_dict: # nocoverage continue if request_settings_dict[setting_name] is None: + if system_groups_name_dict is None: + system_groups_name_dict = get_role_based_system_groups_dict(realm) group_setting_values[setting_name] = get_stream_permission_default_group( - setting_name, realm + setting_name, system_groups_name_dict ) else: group_setting_values[setting_name] = request_settings_dict[setting_name] diff --git a/zerver/lib/user_groups.py b/zerver/lib/user_groups.py index 2946b85526..f73c695c6f 100644 --- a/zerver/lib/user_groups.py +++ b/zerver/lib/user_groups.py @@ -868,12 +868,11 @@ def bulk_create_system_user_groups(groups: list[dict[str, str]], realm: Realm) - @transaction.atomic(savepoint=False) -def create_system_user_groups_for_realm(realm: Realm) -> dict[int, NamedUserGroup]: +def create_system_user_groups_for_realm(realm: Realm) -> dict[str, NamedUserGroup]: """Any changes to this function likely require a migration to adjust existing realms. See e.g. migration 0382_create_role_based_system_groups.py, which is a copy of this function from when we introduced system groups. """ - role_system_groups_dict: dict[int, NamedUserGroup] = {} system_groups_info_list: list[dict[str, str]] = [] @@ -906,9 +905,6 @@ def create_system_user_groups_for_realm(realm: Realm) -> dict[int, NamedUserGrou bulk_create_system_user_groups(system_groups_info_list, realm) system_groups_name_dict: dict[str, NamedUserGroup] = get_role_based_system_groups_dict(realm) - for role in NamedUserGroup.SYSTEM_USER_GROUP_ROLE_MAP: - group_name = NamedUserGroup.SYSTEM_USER_GROUP_ROLE_MAP[role]["name"] - role_system_groups_dict[role] = system_groups_name_dict[group_name] # Order of this list here is important to create correct GroupGroupMembership objects # Note that because we do not create user memberships here, no audit log entries for @@ -982,7 +978,7 @@ def create_system_user_groups_for_realm(realm: Realm) -> dict[int, NamedUserGrou GroupGroupMembership.objects.bulk_create(subgroup_objects) RealmAuditLog.objects.bulk_create(realmauditlog_objects) - return role_system_groups_dict + return system_groups_name_dict def get_system_user_group_for_user(user_profile: UserProfile) -> NamedUserGroup: diff --git a/zerver/views/streams.py b/zerver/views/streams.py index 60da8c9cff..19c69fa79b 100644 --- a/zerver/views/streams.py +++ b/zerver/views/streams.py @@ -81,6 +81,7 @@ from zerver.lib.user_groups import ( GroupSettingChangeRequest, access_user_group_for_setting, get_group_setting_value_for_api, + get_role_based_system_groups_dict, parse_group_setting_value, validate_group_setting_value_change, ) @@ -594,6 +595,9 @@ def add_subscriptions_backend( setting_groups_dict = {} group_settings_map = {} request_settings_dict = locals() + # We don't want to calculate this value if no default values are + # needed. + system_groups_name_dict = None for setting_name, permission_configuration in Stream.stream_permission_group_settings.items(): assert setting_name in request_settings_dict if request_settings_dict[setting_name] is not None: @@ -607,8 +611,10 @@ def add_subscriptions_backend( ) setting_groups_dict[group_settings_map[setting_name].id] = setting_value else: + if system_groups_name_dict is None: + system_groups_name_dict = get_role_based_system_groups_dict(realm) group_settings_map[setting_name] = get_stream_permission_default_group( - setting_name, realm + setting_name, system_groups_name_dict ) setting_groups_dict[group_settings_map[setting_name].id] = group_settings_map[ setting_name