From 86f73fcb3d376cfffcbe31e345848e60a54aafa2 Mon Sep 17 00:00:00 2001 From: Sahil Batra Date: Tue, 16 Apr 2024 19:35:43 +0530 Subject: [PATCH] user_groups: Create NamedUserGroup objects when creating new groups. --- zerver/actions/user_groups.py | 22 ++- zerver/lib/import_realm.py | 7 +- zerver/lib/user_groups.py | 165 ++++++++++++------- zerver/models/__init__.py | 1 + zerver/tests/test_audit_log.py | 16 +- zerver/tests/test_user_groups.py | 37 +++-- zerver/transaction_tests/test_user_groups.py | 4 +- 7 files changed, 166 insertions(+), 86 deletions(-) diff --git a/zerver/actions/user_groups.py b/zerver/actions/user_groups.py index 97bc792847..a8dfee5d7a 100644 --- a/zerver/actions/user_groups.py +++ b/zerver/actions/user_groups.py @@ -13,6 +13,7 @@ from zerver.lib.user_groups import ( ) from zerver.models import ( GroupGroupMembership, + NamedUserGroup, Realm, RealmAuditLog, UserGroup, @@ -40,13 +41,22 @@ def create_user_group_in_database( description: str = "", group_settings_map: Mapping[str, UserGroup] = {}, is_system_group: bool = False, -) -> UserGroup: - user_group = UserGroup( - name=name, realm=realm, description=description, is_system_group=is_system_group +) -> NamedUserGroup: + user_group = NamedUserGroup( + name=name, + realm=realm, + description=description, + is_system_group=is_system_group, + named_group_name=name, + named_group_description=description, + named_group_is_system_group=is_system_group, + realm_for_sharding=realm, ) for setting_name, setting_value in group_settings_map.items(): setattr(user_group, setting_name, setting_value) + named_group_setting_name = "named_group_" + setting_name + setattr(user_group, named_group_setting_name, setting_value) system_groups_name_dict = get_role_based_system_groups_dict(realm) user_group = set_defaults_for_group_settings( @@ -156,7 +166,9 @@ def promote_new_full_members() -> None: def do_send_create_user_group_event( - user_group: UserGroup, members: List[UserProfile], direct_subgroups: Sequence[UserGroup] = [] + user_group: NamedUserGroup, + members: List[UserProfile], + direct_subgroups: Sequence[UserGroup] = [], ) -> None: event = dict( type="user_group", @@ -182,7 +194,7 @@ def check_add_user_group( group_settings_map: Mapping[str, UserGroup] = {}, *, acting_user: Optional[UserProfile], -) -> UserGroup: +) -> NamedUserGroup: try: user_group = create_user_group_in_database( name, diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py index 699282007f..a398ef3460 100644 --- a/zerver/lib/import_realm.py +++ b/zerver/lib/import_realm.py @@ -52,6 +52,7 @@ from zerver.models import ( Huddle, Message, MutedUser, + NamedUserGroup, OnboardingStep, Reaction, Realm, @@ -1040,7 +1041,7 @@ 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: Optional[Dict[int, UserGroup]] = None + role_system_groups_dict: Optional[Dict[int, NamedUserGroup]] = None if "zerver_usergroup" not in data: role_system_groups_dict = create_system_user_groups_for_realm(realm) @@ -1736,7 +1737,9 @@ 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, UserGroup] + realm: Realm, + user_profiles: List[UserProfile], + role_system_groups_dict: Dict[int, NamedUserGroup], ) -> None: full_members_system_group = UserGroup.objects.get( name=SystemGroups.FULL_MEMBERS, diff --git a/zerver/lib/user_groups.py b/zerver/lib/user_groups.py index 13377101d2..42d601d46a 100644 --- a/zerver/lib/user_groups.py +++ b/zerver/lib/user_groups.py @@ -2,17 +2,19 @@ from contextlib import contextmanager from dataclasses import dataclass from typing import Collection, Dict, Iterable, Iterator, List, Mapping, TypedDict -from django.db import transaction +from django.db import connection, transaction 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 from django_stubs_ext import ValuesQuerySet +from psycopg2.sql import SQL, Literal from zerver.lib.exceptions import JsonableError from zerver.lib.types import GroupPermissionSetting, ServerSupportedPermissionSettings from zerver.models import ( GroupGroupMembership, + NamedUserGroup, Realm, RealmAuditLog, Stream, @@ -394,8 +396,8 @@ def get_recursive_subgroups_for_groups( return recursive_subgroups -def get_role_based_system_groups_dict(realm: Realm) -> Dict[str, UserGroup]: - system_groups = UserGroup.objects.filter(realm=realm, is_system_group=True) +def get_role_based_system_groups_dict(realm: Realm) -> Dict[str, NamedUserGroup]: + system_groups = NamedUserGroup.objects.filter(realm=realm, is_system_group=True) system_groups_name_dict = {} for group in system_groups: system_groups_name_dict[group.name] = group @@ -404,10 +406,10 @@ def get_role_based_system_groups_dict(realm: Realm) -> Dict[str, UserGroup]: def set_defaults_for_group_settings( - user_group: UserGroup, + user_group: NamedUserGroup, group_settings_map: Mapping[str, UserGroup], - system_groups_name_dict: Dict[str, UserGroup], -) -> UserGroup: + system_groups_name_dict: Dict[str, NamedUserGroup], +) -> NamedUserGroup: for setting_name, permission_config in UserGroup.GROUP_PERMISSION_SETTINGS.items(): if setting_name in group_settings_map: # We skip the settings for which a value is passed @@ -419,73 +421,121 @@ def set_defaults_for_group_settings( else: default_group_name = permission_config.default_group_name - default_group = system_groups_name_dict[default_group_name] + default_group = system_groups_name_dict[default_group_name].usergroup_ptr setattr(user_group, setting_name, default_group) + setting_name_for_named_object = "named_group_" + setting_name + setattr(user_group, setting_name_for_named_object, default_group) + return user_group -@transaction.atomic(savepoint=False) -def create_system_user_groups_for_realm(realm: Realm) -> Dict[int, UserGroup]: - """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, UserGroup] = {} - +def bulk_create_system_user_groups(groups: List[Dict[str, str]], realm: Realm) -> None: # This value will be used to set the temporary initial value for different # settings since we can only set them to the correct values after the groups # are created. initial_group_setting_value = -1 - - for role in UserGroup.SYSTEM_USER_GROUP_ROLE_MAP: - user_group_params = UserGroup.SYSTEM_USER_GROUP_ROLE_MAP[role] - user_group = UserGroup( - name=user_group_params["name"], - description=user_group_params["description"], - realm=realm, - is_system_group=True, - can_mention_group_id=initial_group_setting_value, + rows = [ + SQL("({},{},{},{},{})").format( + Literal(realm.id), + Literal(group["name"]), + Literal(group["description"]), + Literal(True), + Literal(initial_group_setting_value), ) - role_system_groups_dict[role] = user_group + for group in groups + ] + + query = SQL( + """ + INSERT INTO zerver_usergroup (realm_id, name, description, is_system_group, can_mention_group_id) + VALUES {rows} + RETURNING id + """ + ).format(rows=SQL(", ").join(rows)) + with connection.cursor() as cursor: + cursor.execute(query) + user_group_ids = [id for (id,) in cursor.fetchall()] + + rows = [ + SQL("({},{},{},{},{},{})").format( + Literal(user_group_ids[idx]), + Literal(realm.id), + Literal(group["name"]), + Literal(group["description"]), + Literal(True), + Literal(initial_group_setting_value), + ) + for idx, group in enumerate(groups) + ] + query = SQL( + """ + INSERT INTO zerver_namedusergroup (usergroup_ptr_id, realm_id, name, description, is_system_group, can_mention_group_id) + VALUES {rows} + """ + ).format(rows=SQL(", ").join(rows)) + with connection.cursor() as cursor: + cursor.execute(query) + + +@transaction.atomic(savepoint=False) +def create_system_user_groups_for_realm(realm: Realm) -> Dict[int, 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]] = [] + + nobody_group_info = { + "name": SystemGroups.NOBODY, + "description": "Nobody", + } + + full_members_group_info = { + "name": SystemGroups.FULL_MEMBERS, + "description": "Members of this organization, not including new accounts and guests", + } + + everyone_on_internet_group_info = { + "name": SystemGroups.EVERYONE_ON_INTERNET, + "description": "Everyone on the Internet", + } + + system_groups_info_list = [ + nobody_group_info, + UserGroup.SYSTEM_USER_GROUP_ROLE_MAP[UserProfile.ROLE_REALM_OWNER], + UserGroup.SYSTEM_USER_GROUP_ROLE_MAP[UserProfile.ROLE_REALM_ADMINISTRATOR], + UserGroup.SYSTEM_USER_GROUP_ROLE_MAP[UserProfile.ROLE_MODERATOR], + full_members_group_info, + UserGroup.SYSTEM_USER_GROUP_ROLE_MAP[UserProfile.ROLE_MEMBER], + UserGroup.SYSTEM_USER_GROUP_ROLE_MAP[UserProfile.ROLE_GUEST], + everyone_on_internet_group_info, + ] + + 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 UserGroup.SYSTEM_USER_GROUP_ROLE_MAP: + group_name = UserGroup.SYSTEM_USER_GROUP_ROLE_MAP[role]["name"] + role_system_groups_dict[role] = system_groups_name_dict[group_name] - full_members_system_group = UserGroup( - name=SystemGroups.FULL_MEMBERS, - description="Members of this organization, not including new accounts and guests", - realm=realm, - is_system_group=True, - can_mention_group_id=initial_group_setting_value, - ) - everyone_on_internet_system_group = UserGroup( - name=SystemGroups.EVERYONE_ON_INTERNET, - description="Everyone on the Internet", - realm=realm, - is_system_group=True, - can_mention_group_id=initial_group_setting_value, - ) - nobody_system_group = UserGroup( - name=SystemGroups.NOBODY, - description="Nobody", - realm=realm, - is_system_group=True, - can_mention_group_id=initial_group_setting_value, - ) # 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 # user memberships are populated either. system_user_groups_list = [ - nobody_system_group, - role_system_groups_dict[UserProfile.ROLE_REALM_OWNER], - role_system_groups_dict[UserProfile.ROLE_REALM_ADMINISTRATOR], - role_system_groups_dict[UserProfile.ROLE_MODERATOR], - full_members_system_group, - role_system_groups_dict[UserProfile.ROLE_MEMBER], - role_system_groups_dict[UserProfile.ROLE_GUEST], - everyone_on_internet_system_group, + system_groups_name_dict[SystemGroups.NOBODY], + system_groups_name_dict[SystemGroups.OWNERS], + system_groups_name_dict[SystemGroups.ADMINISTRATORS], + system_groups_name_dict[SystemGroups.MODERATORS], + system_groups_name_dict[SystemGroups.FULL_MEMBERS], + system_groups_name_dict[SystemGroups.MEMBERS], + system_groups_name_dict[SystemGroups.EVERYONE], + system_groups_name_dict[SystemGroups.EVERYONE_ON_INTERNET], ] creation_time = timezone_now() - UserGroup.objects.bulk_create(system_user_groups_list) realmauditlog_objects = [ RealmAuditLog( realm=realm, @@ -498,11 +548,12 @@ def create_system_user_groups_for_realm(realm: Realm) -> Dict[int, UserGroup]: ] groups_with_updated_settings = [] - system_groups_name_dict = get_role_based_system_groups_dict(realm) for group in system_user_groups_list: user_group = set_defaults_for_group_settings(group, {}, system_groups_name_dict) groups_with_updated_settings.append(user_group) - UserGroup.objects.bulk_update(groups_with_updated_settings, ["can_mention_group"]) + NamedUserGroup.objects.bulk_update( + groups_with_updated_settings, ["can_mention_group", "named_group_can_mention_group"] + ) subgroup_objects: List[GroupGroupMembership] = [] # "Nobody" system group is not a subgroup of any user group, since it is already empty. diff --git a/zerver/models/__init__.py b/zerver/models/__init__.py index ce60472b55..b6cf65b5a6 100644 --- a/zerver/models/__init__.py +++ b/zerver/models/__init__.py @@ -8,6 +8,7 @@ from zerver.models.custom_profile_fields import CustomProfileField as CustomProf from zerver.models.custom_profile_fields import CustomProfileFieldValue as CustomProfileFieldValue from zerver.models.drafts import Draft as Draft from zerver.models.groups import GroupGroupMembership as GroupGroupMembership +from zerver.models.groups import NamedUserGroup as NamedUserGroup from zerver.models.groups import UserGroup as UserGroup from zerver.models.groups import UserGroupMembership as UserGroupMembership from zerver.models.linkifiers import RealmFilter as RealmFilter diff --git a/zerver/tests/test_audit_log.py b/zerver/tests/test_audit_log.py index bda3f6c75a..403973a7b1 100644 --- a/zerver/tests/test_audit_log.py +++ b/zerver/tests/test_audit_log.py @@ -1167,7 +1167,7 @@ class TestRealmAuditLog(ZulipTestCase): ) self.assert_length(audit_log_entries, 1) self.assertIsNone(audit_log_entries[0].modified_user) - self.assertEqual(audit_log_entries[0].modified_user_group, user_group) + self.assertEqual(audit_log_entries[0].modified_user_group, user_group.usergroup_ptr) audit_log_entries = RealmAuditLog.objects.filter( acting_user=hamlet, @@ -1212,19 +1212,21 @@ class TestRealmAuditLog(ZulipTestCase): hamlet = self.example_user("hamlet") user_group = check_add_user_group(hamlet.realm, "main", [], acting_user=None) subgroups = [ - check_add_user_group(hamlet.realm, f"subgroup{num}", [], acting_user=hamlet) + check_add_user_group( + hamlet.realm, f"subgroup{num}", [], acting_user=hamlet + ).usergroup_ptr for num in range(3) ] now = timezone_now() - add_subgroups_to_user_group(user_group, subgroups, acting_user=hamlet) + add_subgroups_to_user_group(user_group.usergroup_ptr, subgroups, acting_user=hamlet) # Only one audit log entry for the subgroup membership is expected. audit_log_entry = RealmAuditLog.objects.get( realm=hamlet.realm, event_time__gte=now, event_type=RealmAuditLog.USER_GROUP_DIRECT_SUBGROUP_MEMBERSHIP_ADDED, ) - self.assertEqual(audit_log_entry.modified_user_group, user_group) + self.assertEqual(audit_log_entry.modified_user_group, user_group.usergroup_ptr) self.assertEqual(audit_log_entry.acting_user, hamlet) self.assertDictEqual( audit_log_entry.extra_data, @@ -1244,13 +1246,15 @@ class TestRealmAuditLog(ZulipTestCase): {"supergroup_ids": [user_group.id]}, ) - remove_subgroups_from_user_group(user_group, subgroups[:2], acting_user=hamlet) + remove_subgroups_from_user_group( + user_group.usergroup_ptr, subgroups[:2], acting_user=hamlet + ) audit_log_entry = RealmAuditLog.objects.get( realm=hamlet.realm, event_time__gte=now, event_type=RealmAuditLog.USER_GROUP_DIRECT_SUBGROUP_MEMBERSHIP_REMOVED, ) - self.assertEqual(audit_log_entry.modified_user_group, user_group) + self.assertEqual(audit_log_entry.modified_user_group, user_group.usergroup_ptr) self.assertEqual(audit_log_entry.acting_user, hamlet) self.assertDictEqual( audit_log_entry.extra_data, diff --git a/zerver/tests/test_user_groups.py b/zerver/tests/test_user_groups.py index 310c9fe985..0c7506aa06 100644 --- a/zerver/tests/test_user_groups.py +++ b/zerver/tests/test_user_groups.py @@ -113,20 +113,29 @@ class UserGroupTestCase(ZulipTestCase): everyone_group = check_add_user_group(realm, "Everyone", [shiva], acting_user=None) GroupGroupMembership.objects.create(supergroup=everyone_group, subgroup=staff_group) - self.assertCountEqual(list(get_recursive_subgroups(leadership_group)), [leadership_group]) self.assertCountEqual( - list(get_recursive_subgroups(staff_group)), [leadership_group, staff_group] + list(get_recursive_subgroups(leadership_group)), [leadership_group.usergroup_ptr] + ) + self.assertCountEqual( + list(get_recursive_subgroups(staff_group)), + [leadership_group.usergroup_ptr, staff_group.usergroup_ptr], ) self.assertCountEqual( list(get_recursive_subgroups(everyone_group)), - [leadership_group, staff_group, everyone_group], + [ + leadership_group.usergroup_ptr, + staff_group.usergroup_ptr, + everyone_group.usergroup_ptr, + ], ) self.assertCountEqual(list(get_recursive_strict_subgroups(leadership_group)), []) - self.assertCountEqual(list(get_recursive_strict_subgroups(staff_group)), [leadership_group]) + self.assertCountEqual( + list(get_recursive_strict_subgroups(staff_group)), [leadership_group.usergroup_ptr] + ) self.assertCountEqual( list(get_recursive_strict_subgroups(everyone_group)), - [leadership_group, staff_group], + [leadership_group.usergroup_ptr, staff_group.usergroup_ptr], ) self.assertCountEqual(list(get_recursive_group_members(leadership_group)), [desdemona]) @@ -135,14 +144,14 @@ class UserGroupTestCase(ZulipTestCase): list(get_recursive_group_members(everyone_group)), [desdemona, iago, shiva] ) - self.assertIn(leadership_group, get_recursive_membership_groups(desdemona)) - self.assertIn(staff_group, get_recursive_membership_groups(desdemona)) - self.assertIn(everyone_group, get_recursive_membership_groups(desdemona)) + self.assertIn(leadership_group.usergroup_ptr, get_recursive_membership_groups(desdemona)) + self.assertIn(staff_group.usergroup_ptr, get_recursive_membership_groups(desdemona)) + self.assertIn(everyone_group.usergroup_ptr, get_recursive_membership_groups(desdemona)) - self.assertIn(staff_group, get_recursive_membership_groups(iago)) - self.assertIn(everyone_group, get_recursive_membership_groups(iago)) + self.assertIn(staff_group.usergroup_ptr, get_recursive_membership_groups(iago)) + self.assertIn(everyone_group.usergroup_ptr, get_recursive_membership_groups(iago)) - self.assertIn(everyone_group, get_recursive_membership_groups(shiva)) + self.assertIn(everyone_group.usergroup_ptr, get_recursive_membership_groups(shiva)) def test_subgroups_of_role_based_system_groups(self) -> None: realm = get_realm("zulip") @@ -368,7 +377,7 @@ class UserGroupAPITestCase(UserGroupTestCase): result = self.client_post("/json/user_groups/create", info=params) self.assert_json_success(result) test_group = UserGroup.objects.get(name="test", realm=hamlet.realm) - self.assertEqual(test_group.can_mention_group, leadership_group) + self.assertEqual(test_group.can_mention_group, leadership_group.usergroup_ptr) nobody_group = UserGroup.objects.get( name="role:nobody", realm=hamlet.realm, is_system_group=True @@ -538,7 +547,7 @@ class UserGroupAPITestCase(UserGroupTestCase): result = self.client_patch(f"/json/user_groups/{support_group.id}", info=params) self.assert_json_success(result) support_group = UserGroup.objects.get(name="support", realm=hamlet.realm) - self.assertEqual(support_group.can_mention_group, marketing_group) + self.assertEqual(support_group.can_mention_group, marketing_group.usergroup_ptr) nobody_group = UserGroup.objects.get( name="role:nobody", realm=hamlet.realm, is_system_group=True @@ -640,7 +649,7 @@ class UserGroupAPITestCase(UserGroupTestCase): for i in range(50) ] - with self.assert_database_query_count(4): + with self.assert_database_query_count(5): user_group = create_user_group_in_database( name="support", members=[hamlet, cordelia, *original_users], diff --git a/zerver/transaction_tests/test_user_groups.py b/zerver/transaction_tests/test_user_groups.py index 7655530047..caa592bd8f 100644 --- a/zerver/transaction_tests/test_user_groups.py +++ b/zerver/transaction_tests/test_user_groups.py @@ -12,7 +12,7 @@ from zerver.lib.exceptions import JsonableError from zerver.lib.test_classes import ZulipTransactionTestCase from zerver.lib.test_helpers import HostRequestMock from zerver.lib.user_groups import access_user_group_by_id -from zerver.models import Realm, UserGroup, UserProfile +from zerver.models import NamedUserGroup, Realm, UserGroup, UserProfile from zerver.models.realms import get_realm from zerver.views.user_groups import update_subgroups_of_user_group @@ -77,7 +77,7 @@ class UserGroupRaceConditionTestCase(ZulipTransactionTestCase): super().tearDown() - def create_user_group_chain(self, realm: Realm) -> List[UserGroup]: + def create_user_group_chain(self, realm: Realm) -> List[NamedUserGroup]: """Build a user groups forming a chain through group-group memberships returning a list where each group is the supergroup of its subsequent group. """