diff --git a/zerver/lib/export.py b/zerver/lib/export.py index 91357557e3..32c5a9dfa4 100644 --- a/zerver/lib/export.py +++ b/zerver/lib/export.py @@ -540,7 +540,7 @@ def get_realm_config() -> Config: custom_fetch=fetch_user_profile, ) - Config( + user_groups_config = Config( table='zerver_usergroup', model=UserGroup, normal_parent=realm_config, @@ -550,7 +550,7 @@ def get_realm_config() -> Config: Config( table='zerver_usergroupmembership', model=UserGroupMembership, - normal_parent=user_profile_config, + normal_parent=user_groups_config, parent_key='user_group__in', ) diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py index 73d979abaf..7d5a13be17 100644 --- a/zerver/lib/import_realm.py +++ b/zerver/lib/import_realm.py @@ -27,7 +27,7 @@ from zerver.models import UserProfile, Realm, Client, Huddle, Stream, \ UserPresence, UserActivity, UserActivityInterval, Reaction, \ CustomProfileField, CustomProfileFieldValue, RealmAuditLog, \ Attachment, get_system_bot, email_to_username, get_huddle_hash, \ - UserHotspot, MutedTopic, Service + UserHotspot, MutedTopic, Service, UserGroup, UserGroupMembership # Code from here is the realm import code path @@ -66,6 +66,8 @@ id_maps = { 'userhotspot': {}, 'mutedtopic': {}, 'service': {}, + 'usergroup': {}, + 'usergroupmembership': {}, } # type: Dict[str, Dict[int, int]] id_map_to_list = { @@ -702,6 +704,21 @@ def do_import_realm(import_dir: Path, subdomain: str) -> Realm: update_model_ids(Service, data, 'zerver_service', 'service') bulk_import_model(data, Service, 'zerver_service') + if 'zerver_usergroup' in data: + re_map_foreign_keys(data, 'zerver_usergroup', 'realm', related_table='realm') + re_map_foreign_keys_many_to_many(data, 'zerver_usergroup', + 'members', related_table='user_profile') + update_model_ids(UserGroup, data, 'zerver_usergroup', 'usergroup') + bulk_import_model(data, UserGroup, 'zerver_usergroup') + + re_map_foreign_keys(data, 'zerver_usergroupmembership', + 'user_group', related_table='usergroup') + re_map_foreign_keys(data, 'zerver_usergroupmembership', + 'user_profile', related_table='user_profile') + update_model_ids(UserGroupMembership, data, 'zerver_usergroupmembership', + 'usergroupmembership') + bulk_import_model(data, UserGroupMembership, 'zerver_usergroupmembership') + fix_datetime_fields(data, 'zerver_userpresence') re_map_foreign_keys(data, 'zerver_userpresence', 'user_profile', related_table="user_profile") re_map_foreign_keys(data, 'zerver_userpresence', 'client', related_table='client') diff --git a/zerver/tests/test_import_export.py b/zerver/tests/test_import_export.py index 6ce031e30f..bdcad3d444 100644 --- a/zerver/tests/test_import_export.py +++ b/zerver/tests/test_import_export.py @@ -63,6 +63,8 @@ from zerver.models import ( Huddle, UserHotspot, MutedTopic, + UserGroup, + UserGroupMembership, get_active_streams, get_stream, get_stream_recipient, @@ -653,6 +655,19 @@ class ImportExportTest(ZulipTestCase): assert_realm_values(get_muted_topics) + # test usergroups + assert_realm_values( + lambda r: {group.name for group in UserGroup.objects.filter(realm=r)} + ) + + def get_user_membership(r: str) -> Set[str]: + usergroup = UserGroup.objects.get(realm=r, name='hamletcharacters') + usergroup_membership = UserGroupMembership.objects.filter(user_group=usergroup) + users = {membership.user_profile.email for membership in usergroup_membership} + return users + + assert_realm_values(get_user_membership) + # test messages def get_stream_messages(r: Realm) -> Message: recipient = get_recipient_stream(r)