From c9a7c13ea762ff1bbf3ff92b8eccddf1aa6c301f Mon Sep 17 00:00:00 2001 From: Sahil Batra Date: Tue, 16 Apr 2024 14:40:31 +0530 Subject: [PATCH] import: Add code to support import and export of NamedUserGroups. --- zerver/lib/export.py | 11 ++++++-- zerver/lib/import_realm.py | 40 ++++++++++++++++++++++++++++++ zerver/tests/test_import_export.py | 15 +++++++++++ 3 files changed, 64 insertions(+), 2 deletions(-) diff --git a/zerver/lib/export.py b/zerver/lib/export.py index d5387efddd..533d2a3809 100644 --- a/zerver/lib/export.py +++ b/zerver/lib/export.py @@ -45,6 +45,7 @@ from zerver.models import ( Huddle, Message, MutedUser, + NamedUserGroup, OnboardingStep, Reaction, Realm, @@ -244,8 +245,6 @@ NON_EXPORTED_TABLES = { "zerver_submessage", # Drafts don't need to be exported as they are supposed to be more ephemeral. "zerver_draft", - # NamedUserGroup is not exported temporarily, will be done in the next few commits. - "zerver_namedusergroup", # For any tables listed below here, it's a bug that they are not present in the export. } @@ -779,6 +778,14 @@ def get_realm_config() -> Config: exclude=["direct_members", "direct_subgroups"], ) + Config( + table="zerver_namedusergroup", + model=NamedUserGroup, + normal_parent=realm_config, + include_rows="realm_for_sharding_id__in", + exclude=["realm", "direct_members", "direct_subgroups"], + ) + Config( table="zerver_usergroupmembership", model=UserGroupMembership, diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py index a398ef3460..a5d51690ff 100644 --- a/zerver/lib/import_realm.py +++ b/zerver/lib/import_realm.py @@ -689,6 +689,29 @@ def bulk_import_model(data: TableData, model: Any, dump_file_id: Optional[str] = logging.info("Successfully imported %s from %s[%s].", model, table, dump_file_id) +def bulk_import_named_user_groups(data: TableData) -> None: + vals = [ + ( + group["usergroup_ptr_id"], + group["realm_for_sharding_id"], + group["named_group_name"], + group["named_group_description"], + group["named_group_is_system_group"], + group["named_group_can_mention_group_id"], + ) + for group in data["zerver_namedusergroup"] + ] + + query = SQL( + """ + INSERT INTO zerver_namedusergroup (usergroup_ptr_id, realm_id, name, description, is_system_group, can_mention_group_id) + VALUES %s + """ + ) + with connection.cursor() as cursor: + execute_values(cursor.cursor, query, vals) + + # Client is a table shared by multiple realms, so in order to # correctly import multiple realms into the same server, we need to # check if a Client object already exists, and so we need to support @@ -1039,6 +1062,23 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int = 1) -> Rea ) bulk_import_model(data, UserGroup) + if "zerver_namedusergroup" in data: + re_map_foreign_keys( + data, "zerver_namedusergroup", "usergroup_ptr", related_table="usergroup" + ) + re_map_foreign_keys( + data, "zerver_namedusergroup", "realm_for_sharding", related_table="realm" + ) + for setting_name in UserGroup.GROUP_PERMISSION_SETTINGS: + named_group_setting_name = "named_group_" + setting_name + re_map_foreign_keys( + data, + "zerver_namedusergroup", + named_group_setting_name, + related_table="usergroup", + ) + bulk_import_named_user_groups(data) + # 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, NamedUserGroup]] = None diff --git a/zerver/tests/test_import_export.py b/zerver/tests/test_import_export.py index a708f5afa6..3dfbd08c2e 100644 --- a/zerver/tests/test_import_export.py +++ b/zerver/tests/test_import_export.py @@ -66,6 +66,7 @@ from zerver.models import ( Huddle, Message, MutedUser, + NamedUserGroup, OnboardingStep, Reaction, Realm, @@ -457,6 +458,15 @@ class RealmImportExportTest(ExportFile): self.assertFalse("direct_members" in exported_usergroups[2]) self.assertFalse("direct_subgroups" in exported_usergroups[2]) + exported_namedusergroups = data["zerver_namedusergroup"] + self.assert_length(exported_namedusergroups, 9) + self.assertEqual(exported_namedusergroups[2]["named_group_name"], "role:administrators") + self.assertTrue("usergroup_ptr" in exported_namedusergroups[2]) + self.assertTrue("realm_for_sharding" in exported_namedusergroups[2]) + self.assertFalse("realm" in exported_namedusergroups[2]) + self.assertFalse("direct_members" in exported_namedusergroups[2]) + self.assertFalse("direct_subgroups" in exported_namedusergroups[2]) + data = read_json("messages-000001.json") um = UserMessage.objects.all()[0] exported_um = self.find_by_id(data["zerver_usermessage"], um.id) @@ -1223,6 +1233,10 @@ class RealmImportExportTest(ExportFile): def get_user_group_names(r: Realm) -> Set[str]: return {group.name for group in UserGroup.objects.filter(realm=r)} + @getter + def get_named_user_group_names(r: Realm) -> Set[str]: + return {group.name for group in NamedUserGroup.objects.filter(realm=r)} + @getter def get_user_membership(r: Realm) -> Set[str]: usergroup = UserGroup.objects.get(realm=r, name="hamletcharacters") @@ -1669,6 +1683,7 @@ class RealmImportExportTest(ExportFile): # Simulate an external export where user groups are missing. data = read_json("realm.json") data.pop("zerver_usergroup") + data.pop("zerver_namedusergroup") data.pop("zerver_realmauditlog") # User groups data is missing. So, all the realm group based settings