import: Add code to support import and export of NamedUserGroups.

This commit is contained in:
Sahil Batra 2024-04-16 14:40:31 +05:30 committed by Tim Abbott
parent b60d5b6fc9
commit c9a7c13ea7
3 changed files with 64 additions and 2 deletions

View File

@ -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,

View File

@ -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

View File

@ -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