From 2ea53a347a548796eead83b579bf3c8ee87c7bc4 Mon Sep 17 00:00:00 2001 From: Vishnu Ks Date: Fri, 19 Jul 2019 22:45:23 +0530 Subject: [PATCH] import: Support importing realm icon and logo. Fixes #11216 --- zerver/lib/export.py | 82 ++++++++++++++++------ zerver/lib/import_realm.py | 38 ++++++++--- zerver/tests/test_import_export.py | 105 ++++++++++++++++++++++++++++- 3 files changed, 191 insertions(+), 34 deletions(-) diff --git a/zerver/lib/export.py b/zerver/lib/export.py index 162950e77e..542412d358 100644 --- a/zerver/lib/export.py +++ b/zerver/lib/export.py @@ -1121,11 +1121,12 @@ def write_message_partial_for_query(realm: Realm, message_query: Any, dump_file_ def export_uploads_and_avatars(realm: Realm, output_dir: Path) -> None: uploads_output_dir = os.path.join(output_dir, 'uploads') avatars_output_dir = os.path.join(output_dir, 'avatars') + realm_icons_output_dir = os.path.join(output_dir, 'realm_icons') emoji_output_dir = os.path.join(output_dir, 'emoji') - for output_dir in (uploads_output_dir, avatars_output_dir, emoji_output_dir): - if not os.path.exists(output_dir): - os.makedirs(output_dir) + for dir_path in (uploads_output_dir, avatars_output_dir, realm_icons_output_dir, emoji_output_dir): + if not os.path.exists(dir_path): + os.makedirs(dir_path) if settings.LOCAL_UPLOADS_DIR: # Small installations and developers will usually just store files locally. @@ -1138,6 +1139,9 @@ def export_uploads_and_avatars(realm: Realm, output_dir: Path) -> None: export_emoji_from_local(realm, local_dir=os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars"), output_dir=emoji_output_dir) + export_realm_icons(realm, + local_dir=os.path.join(settings.LOCAL_UPLOADS_DIR), + output_dir=realm_icons_output_dir) else: # Some bigger installations will have their data stored on S3. export_files_from_s3(realm, @@ -1151,6 +1155,10 @@ def export_uploads_and_avatars(realm: Realm, output_dir: Path) -> None: settings.S3_AVATAR_BUCKET, output_dir=emoji_output_dir, processing_emoji=True) + export_files_from_s3(realm, + settings.S3_AVATAR_BUCKET, + output_dir=realm_icons_output_dir, + processing_realm_icon_and_logo=True) def _check_key_metadata(email_gateway_bot: Optional[UserProfile], key: Key, processing_avatars: bool, @@ -1183,26 +1191,34 @@ def _get_exported_s3_record( if processing_emoji: record['file_name'] = os.path.basename(key.name) - # A few early avatars don't have 'realm_id' on the object; fix their metadata - user_profile = get_user_profile_by_id(record['user_profile_id']) - if 'realm_id' not in record: - record['realm_id'] = user_profile.realm_id - record['user_profile_email'] = user_profile.email + if "user_profile_id" in record: + user_profile = get_user_profile_by_id(record['user_profile_id']) + record['user_profile_email'] = user_profile.email - # Fix the record ids - record['user_profile_id'] = int(record['user_profile_id']) - record['realm_id'] = int(record['realm_id']) + # Fix the record ids + record['user_profile_id'] = int(record['user_profile_id']) + + # A few early avatars don't have 'realm_id' on the object; fix their metadata + if 'realm_id' not in record: + record['realm_id'] = user_profile.realm_id + else: + # There are some rare cases in which 'user_profile_id' may not be present + # in S3 metadata. Eg: Exporting an organization which was created + # initially from a local export won't have the "user_profile_id" metadata + # set for realm_icons and realm_logos. + pass + + if 'realm_id' in record: + record['realm_id'] = int(record['realm_id']) + else: + raise Exception("Missing realm_id") return record -def _save_s3_object_to_file( - key: Key, - output_dir: str, - processing_avatars: bool, - processing_emoji: bool) -> None: - +def _save_s3_object_to_file(key: Key, output_dir: str, processing_avatars: bool, + processing_emoji: bool, processing_realm_icon_and_logo: bool) -> None: # Helper function for export_files_from_s3 - if processing_avatars or processing_emoji: + if processing_avatars or processing_emoji or processing_realm_icon_and_logo: filename = os.path.join(output_dir, key.name) else: fields = key.name.split('/') @@ -1216,8 +1232,8 @@ def _save_s3_object_to_file( key.get_contents_to_filename(filename) def export_files_from_s3(realm: Realm, bucket_name: str, output_dir: Path, - processing_avatars: bool=False, - processing_emoji: bool=False) -> None: + processing_avatars: bool=False, processing_emoji: bool=False, + processing_realm_icon_and_logo: bool=False) -> None: conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY) bucket = conn.get_bucket(bucket_name, validate=True) records = [] @@ -1233,7 +1249,10 @@ def export_files_from_s3(realm: Realm, bucket_name: str, output_dir: Path, avatar_hash_values.add(avatar_path) avatar_hash_values.add(avatar_path + ".original") user_ids.add(user_profile.id) - if processing_emoji: + + if processing_realm_icon_and_logo: + bucket_list = bucket.list(prefix="%s/realm/" % (realm.id,)) + elif processing_emoji: bucket_list = bucket.list(prefix="%s/emoji/images/" % (realm.id,)) else: bucket_list = bucket.list(prefix="%s/" % (realm.id,)) @@ -1254,7 +1273,8 @@ def export_files_from_s3(realm: Realm, bucket_name: str, output_dir: Path, record = _get_exported_s3_record(bucket_name, key, processing_avatars, processing_emoji) record['path'] = key.name - _save_s3_object_to_file(key, output_dir, processing_avatars, processing_emoji) + _save_s3_object_to_file(key, output_dir, processing_avatars, processing_emoji, + processing_realm_icon_and_logo) records.append(record) count += 1 @@ -1336,6 +1356,24 @@ def export_avatars_from_local(realm: Realm, local_dir: Path, output_dir: Path) - with open(os.path.join(output_dir, "records.json"), "w") as records_file: ujson.dump(records, records_file, indent=4) +def export_realm_icons(realm: Realm, local_dir: Path, output_dir: Path) -> None: + records = [] + dir_relative_path = zerver.lib.upload.upload_backend.realm_avatar_and_logo_path(realm) + icons_wildcard = os.path.join(local_dir, dir_relative_path, '*') + for icon_absolute_path in glob.glob(icons_wildcard): + icon_file_name = os.path.basename(icon_absolute_path) + icon_relative_path = os.path.join(str(realm.id), icon_file_name) + output_path = os.path.join(output_dir, icon_relative_path) + os.makedirs(str(os.path.dirname(output_path)), exist_ok=True) + shutil.copy2(str(icon_absolute_path), str(output_path)) + record = dict(realm_id=realm.id, + path=icon_relative_path, + s3_path=icon_relative_path) + records.append(record) + + with open(os.path.join(output_dir, "records.json"), "w") as records_file: + ujson.dump(records, records_file, indent=4) + def export_emoji_from_local(realm: Realm, local_dir: Path, output_dir: Path) -> None: count = 0 diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py index a358d36323..10974884da 100644 --- a/zerver/lib/import_realm.py +++ b/zerver/lib/import_realm.py @@ -584,14 +584,16 @@ def bulk_import_client(data: TableData, model: Any, table: TableName) -> None: client = Client.objects.create(name=item['name']) update_id_map(table='client', old_id=item['id'], new_id=client.id) -def import_uploads(import_dir: Path, processes: int, processing_avatars: bool=False, - processing_emojis: bool=False) -> None: +def import_uploads(realm: Realm, import_dir: Path, processes: int, processing_avatars: bool=False, + processing_emojis: bool=False, processing_realm_icons: bool=False) -> None: if processing_avatars and processing_emojis: raise AssertionError("Cannot import avatars and emojis at the same time!") if processing_avatars: logging.info("Importing avatars") elif processing_emojis: logging.info("Importing emojis") + elif processing_realm_icons: + logging.info("Importing realm icons and logos") else: logging.info("Importing uploaded files") @@ -602,14 +604,14 @@ def import_uploads(import_dir: Path, processes: int, processing_avatars: bool=Fa re_map_foreign_keys_internal(records, 'records', 'realm_id', related_table="realm", id_field=True) - if not processing_emojis: + if not processing_emojis and not processing_realm_icons: re_map_foreign_keys_internal(records, 'records', 'user_profile_id', related_table="user_profile", id_field=True) s3_uploads = settings.LOCAL_UPLOADS_DIR is None if s3_uploads: - if processing_avatars or processing_emojis: + if processing_avatars or processing_emojis or processing_realm_icons: bucket_name = settings.S3_AVATAR_BUCKET else: bucket_name = settings.S3_AUTH_UPLOADS_BUCKET @@ -641,6 +643,10 @@ def import_uploads(import_dir: Path, processes: int, processing_avatars: bool=Fa realm_id=record['realm_id'], emoji_file_name=record['file_name']) record['last_modified'] = timestamp + elif processing_realm_icons: + icon_name = os.path.basename(record["path"]) + relative_path = os.path.join(str(record['realm_id']), "realm", icon_name) + record['last_modified'] = timestamp else: # Should be kept in sync with its equivalent in zerver/lib/uploads in the # function 'upload_message_file' @@ -654,9 +660,15 @@ def import_uploads(import_dir: Path, processes: int, processing_avatars: bool=Fa if s3_uploads: key = Key(bucket) key.key = relative_path - # Exported custom emoji from tools like Slack don't have - # the data for what user uploaded them in `user_profile_id`. - if not processing_emojis: + if processing_emojis: + # Exported custom emoji from tools like Slack don't have + # the data for what user uploaded them in `user_profile_id`. + pass + elif processing_realm_icons and "user_profile_id" not in record: + # Exported realm icons and logos from local export don't have + # the value of user_profile_id in the associated record. + pass + else: user_profile_id = int(record['user_profile_id']) # Support email gateway bot and other cross-realm messages if user_profile_id in ID_MAP["user_profile"]: @@ -683,7 +695,7 @@ def import_uploads(import_dir: Path, processes: int, processing_avatars: bool=Fa key.set_contents_from_filename(os.path.join(import_dir, record['path']), headers=headers) else: - if processing_avatars or processing_emojis: + if processing_avatars or processing_emojis or processing_realm_icons: file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", relative_path) else: file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "files", relative_path) @@ -981,14 +993,18 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int=1) -> Realm bulk_import_model(data, CustomProfileFieldValue) # Import uploaded files and avatars - import_uploads(os.path.join(import_dir, "avatars"), processes, processing_avatars=True) - import_uploads(os.path.join(import_dir, "uploads"), processes) + import_uploads(realm, os.path.join(import_dir, "avatars"), processes, processing_avatars=True) + import_uploads(realm, os.path.join(import_dir, "uploads"), processes) # We need to have this check as the emoji files are only present in the data # importer from slack # For Zulip export, this doesn't exist if os.path.exists(os.path.join(import_dir, "emoji")): - import_uploads(os.path.join(import_dir, "emoji"), processes, processing_emojis=True) + import_uploads(realm, os.path.join(import_dir, "emoji"), processes, processing_emojis=True) + + if os.path.exists(os.path.join(import_dir, "realm_icons")): + import_uploads(realm, os.path.join(import_dir, "realm_icons"), processes, + processing_realm_icons=True) sender_map = { user['id']: user diff --git a/zerver/tests/test_import_export.py b/zerver/tests/test_import_export.py index 88f72e3bc8..96baf5f605 100644 --- a/zerver/tests/test_import_export.py +++ b/zerver/tests/test_import_export.py @@ -28,6 +28,8 @@ from zerver.lib.upload import ( upload_emoji_image, upload_avatar_image, ) +from zerver.lib import upload + from zerver.lib.utils import ( query_chunker, ) @@ -51,7 +53,9 @@ from zerver.lib.bot_config import ( from zerver.lib.actions import ( do_create_user, do_add_reaction, - create_stream_if_needed + create_stream_if_needed, + do_change_icon_source, + do_change_logo_source ) from zerver.lib.test_runner import slow @@ -276,6 +280,8 @@ class ImportExportTest(ZulipTestCase): result['emoji_dir_records'] = read_file(os.path.join('emoji', 'records.json')) result['avatar_dir'] = os.path.join(output_dir, 'avatars') result['avatar_dir_records'] = read_file(os.path.join('avatars', 'records.json')) + result['realm_icons_dir'] = os.path.join(output_dir, 'realm_icons') + result['realm_icons_dir_records'] = read_file(os.path.join('realm_icons', 'records.json')) return result def _setup_export_files(self) -> Tuple[str, str, str, bytes]: @@ -304,6 +310,19 @@ class ImportExportTest(ZulipTestCase): upload_avatar_image(img_file, user_profile, user_profile) with open(get_test_image_file('img.png').name, 'rb') as f: test_image = f.read() + + with get_test_image_file('img.png') as img_file: + upload.upload_backend.upload_realm_icon_image(img_file, user_profile) + do_change_icon_source(realm, Realm.ICON_UPLOADED, False) + + with get_test_image_file('img.png') as img_file: + upload.upload_backend.upload_realm_logo_image(img_file, user_profile, night=False) + do_change_logo_source(realm, Realm.LOGO_UPLOADED, False) + with get_test_image_file('img.png') as img_file: + upload.upload_backend.upload_realm_logo_image(img_file, user_profile, night=True) + do_change_logo_source(realm, Realm.LOGO_UPLOADED, True) + + test_image = get_test_image_file('img.png').read() message.sender.avatar_source = 'U' message.sender.save() @@ -340,6 +359,21 @@ class ImportExportTest(ZulipTestCase): self.assertEqual(records[0]['path'], '2/emoji/images/1.png') self.assertEqual(records[0]['s3_path'], '2/emoji/images/1.png') + # Test realm logo and icon + records = full_data['realm_icons_dir_records'] + image_files = set() + for record in records: + image_path = os.path.join(full_data['realm_icons_dir'], record["path"]) + if image_path[-9:] == ".original": + image_data = open(image_path, 'rb').read() + self.assertEqual(image_data, test_image) + else: + self.assertTrue(os.path.exists(image_path)) + + image_files.add(os.path.basename(image_path)) + self.assertEqual(set(image_files), {'night_logo.png', 'logo.original', 'logo.png', + 'icon.png', 'night_logo.original', 'icon.original'}) + # Test avatars fn = os.path.join(full_data['avatar_dir'], original_avatar_path_id) with open(fn, 'rb') as fb: @@ -391,6 +425,21 @@ class ImportExportTest(ZulipTestCase): self.assertEqual(records[0]['s3_path'], '2/emoji/images/1.png') check_variable_type(records[0]['user_profile_id'], records[0]['realm_id']) + # Test realm logo and icon + records = full_data['realm_icons_dir_records'] + image_files = set() + for record in records: + image_path = os.path.join(full_data['realm_icons_dir'], record["s3_path"]) + if image_path[-9:] == ".original": + image_data = open(image_path, 'rb').read() + self.assertEqual(image_data, test_image) + else: + self.assertTrue(os.path.exists(image_path)) + + image_files.add(os.path.basename(image_path)) + self.assertEqual(set(image_files), {'night_logo.png', 'logo.original', 'logo.png', + 'icon.png', 'night_logo.original', 'icon.original'}) + # Test avatars fn = os.path.join(full_data['avatar_dir'], original_avatar_path_id) with open(fn, 'rb') as file: @@ -958,6 +1007,8 @@ class ImportExportTest(ZulipTestCase): realm = Realm.objects.get(string_id='zulip') self._setup_export_files() + realm.refresh_from_db() + self._export_realm(realm) with patch('logging.info'): @@ -988,6 +1039,29 @@ class ImportExportTest(ZulipTestCase): avatar_file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", avatar_path_id) self.assertTrue(os.path.isfile(avatar_file_path)) + # Test realm icon and logo + upload_path = upload.upload_backend.realm_avatar_and_logo_path(imported_realm) + full_upload_path = os.path.join(settings.LOCAL_UPLOADS_DIR, upload_path) + + with open(get_test_image_file('img.png').name, 'rb') as f: + test_image_data = f.read() + self.assertIsNotNone(test_image_data) + + with open(os.path.join(full_upload_path, "icon.original"), 'rb') as f: + self.assertEqual(f.read(), test_image_data) + self.assertTrue(os.path.isfile(os.path.join(full_upload_path, "icon.png"))) + self.assertEqual(imported_realm.icon_source, Realm.ICON_UPLOADED) + + with open(os.path.join(full_upload_path, "logo.original"), 'rb') as f: + self.assertEqual(f.read(), test_image_data) + self.assertTrue(os.path.isfile(os.path.join(full_upload_path, "logo.png"))) + self.assertEqual(imported_realm.logo_source, Realm.LOGO_UPLOADED) + + with open(os.path.join(full_upload_path, "night_logo.original"), 'rb') as f: + self.assertEqual(f.read(), test_image_data) + self.assertTrue(os.path.isfile(os.path.join(full_upload_path, "night_logo.png"))) + self.assertEqual(imported_realm.night_logo_source, Realm.LOGO_UPLOADED) + @use_s3_backend def test_import_files_from_s3(self) -> None: uploads_bucket, avatar_bucket = create_s3_buckets( @@ -996,6 +1070,8 @@ class ImportExportTest(ZulipTestCase): realm = Realm.objects.get(string_id='zulip') self._setup_export_files() + realm.refresh_from_db() + self._export_realm(realm) with patch('logging.info'): do_import_realm(os.path.join(settings.TEST_WORKER_DIR, 'test-export'), @@ -1030,6 +1106,33 @@ class ImportExportTest(ZulipTestCase): image_data = original_image_key.get_contents_as_string() self.assertEqual(image_data, test_image_data) + # Test realm icon and logo + upload_path = upload.upload_backend.realm_avatar_and_logo_path(imported_realm) + + original_icon_path_id = os.path.join(upload_path, "icon.original") + original_icon_key = avatar_bucket.get_key(original_icon_path_id) + self.assertEqual(original_icon_key.get_contents_as_string(), test_image_data) + resized_icon_path_id = os.path.join(upload_path, "icon.png") + resized_icon_key = avatar_bucket.get_key(resized_icon_path_id) + self.assertEqual(resized_icon_key.key, resized_icon_path_id) + self.assertEqual(imported_realm.icon_source, Realm.ICON_UPLOADED) + + original_logo_path_id = os.path.join(upload_path, "logo.original") + original_logo_key = avatar_bucket.get_key(original_logo_path_id) + self.assertEqual(original_logo_key.get_contents_as_string(), test_image_data) + resized_logo_path_id = os.path.join(upload_path, "logo.png") + resized_logo_key = avatar_bucket.get_key(resized_logo_path_id) + self.assertEqual(resized_logo_key.key, resized_logo_path_id) + self.assertEqual(imported_realm.logo_source, Realm.LOGO_UPLOADED) + + night_logo_original_path_id = os.path.join(upload_path, "night_logo.original") + night_logo_original_key = avatar_bucket.get_key(night_logo_original_path_id) + self.assertEqual(night_logo_original_key.get_contents_as_string(), test_image_data) + resized_night_logo_path_id = os.path.join(upload_path, "night_logo.png") + resized_night_logo_key = avatar_bucket.get_key(resized_night_logo_path_id) + self.assertEqual(resized_night_logo_key.key, resized_night_logo_path_id) + self.assertEqual(imported_realm.night_logo_source, Realm.LOGO_UPLOADED) + def test_get_incoming_message_ids(self) -> None: import_dir = os.path.join(settings.DEPLOY_ROOT, "zerver", "tests", "fixtures", "import_fixtures") message_ids = get_incoming_message_ids(