def new_method(*args, **kwargs): # type: (*Any, **Any) -> Any zerver.lib.upload.upload_backend = S3UploadBackend() try: return method(*args, **kwargs) finally: zerver.lib.upload.upload_backend = LocalUploadBackend()
def test_avatar_url(self) -> None: """Verifies URL schemes for avatars and realm icons.""" backend = LocalUploadBackend() # type: ZulipUploadBackend self.assertEqual(backend.get_avatar_url("hash", False), "/user_avatars/hash.png?x=x") self.assertEqual(backend.get_avatar_url("hash", True), "/user_avatars/hash-medium.png?x=x") self.assertEqual(backend.get_realm_icon_url(15, 1), "/user_avatars/15/realm/icon.png?version=1") with self.settings(S3_AVATAR_BUCKET="bucket"): backend = S3UploadBackend() self.assertEqual(backend.get_avatar_url("hash", False), "https://bucket.s3.amazonaws.com/hash?x=x") self.assertEqual(backend.get_avatar_url("hash", True), "https://bucket.s3.amazonaws.com/hash-medium.png?x=x") self.assertEqual(backend.get_realm_icon_url(15, 1), "https://bucket.s3.amazonaws.com/15/realm/icon.png?version=1")
def import_uploads_local(import_dir: Path, processing_avatars: bool = False, processing_emojis: bool = False) -> None: records_filename = os.path.join(import_dir, "records.json") with open(records_filename) as records_file: records = ujson.loads(records_file.read()) re_map_foreign_keys_internal(records, 'records', 'realm_id', related_table="realm", id_field=True) if not processing_emojis: re_map_foreign_keys_internal(records, 'records', 'user_profile_id', related_table="user_profile", id_field=True) for record in records: if processing_avatars: # For avatars, we need to rehash the user ID with the # new server's avatar salt avatar_path = user_avatar_path_from_ids(record['user_profile_id'], record['realm_id']) file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", avatar_path) if record['s3_path'].endswith('.original'): file_path += '.original' else: file_path += '.png' elif processing_emojis: # For emojis we follow the function 'upload_emoji_image' emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format( realm_id=record['realm_id'], emoji_file_name=record['file_name']) file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", emoji_path) else: # Should be kept in sync with its equivalent in zerver/lib/uploads in the # function 'upload_message_image' s3_file_name = "/".join([ str(record['realm_id']), random_name(18), sanitize_name(os.path.basename(record['path'])) ]) file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "files", s3_file_name) path_maps['attachment_path'][record['path']] = s3_file_name orig_file_path = os.path.join(import_dir, record['path']) if not os.path.exists(os.path.dirname(file_path)): subprocess.check_call(["mkdir", "-p", os.path.dirname(file_path)]) shutil.copy(orig_file_path, file_path) if processing_avatars: # Ensure that we have medium-size avatar images for every # avatar. TODO: This implementation is hacky, both in that it # does get_user_profile_by_id for each user, and in that it # might be better to require the export to just have these. upload_backend = LocalUploadBackend() for record in records: if record['s3_path'].endswith('.original'): user_profile = get_user_profile_by_id( record['user_profile_id']) avatar_path = user_avatar_path_from_ids( user_profile.id, record['realm_id']) medium_file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", avatar_path) + '-medium.png' if os.path.exists(medium_file_path): # We remove the image here primarily to deal with # issues when running the import script multiple # times in development (where one might reuse the # same realm ID from a previous iteration). os.remove(medium_file_path) upload_backend.ensure_medium_avatar_image( user_profile=user_profile)
def import_uploads_local(import_dir: Path, processing_avatars: bool=False, processing_emojis: bool=False) -> None: records_filename = os.path.join(import_dir, "records.json") with open(records_filename) as records_file: records = ujson.loads(records_file.read()) re_map_foreign_keys_internal(records, 'records', 'realm_id', related_table="realm", id_field=True) if not processing_emojis: re_map_foreign_keys_internal(records, 'records', 'user_profile_id', related_table="user_profile", id_field=True) for record in records: if processing_avatars: # For avatars, we need to rehash the user ID with the # new server's avatar salt avatar_path = user_avatar_path_from_ids(record['user_profile_id'], record['realm_id']) file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", avatar_path) if record['s3_path'].endswith('.original'): file_path += '.original' else: file_path += '.png' elif processing_emojis: # For emojis we follow the function 'upload_emoji_image' emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format( realm_id=record['realm_id'], emoji_file_name=record['file_name']) file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", emoji_path) else: # Should be kept in sync with its equivalent in zerver/lib/uploads in the # function 'upload_message_image' s3_file_name = "/".join([ str(record['realm_id']), random_name(18), sanitize_name(os.path.basename(record['path'])) ]) file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "files", s3_file_name) path_maps['attachment_path'][record['path']] = s3_file_name orig_file_path = os.path.join(import_dir, record['path']) if not os.path.exists(os.path.dirname(file_path)): subprocess.check_call(["mkdir", "-p", os.path.dirname(file_path)]) shutil.copy(orig_file_path, file_path) if processing_avatars: # Ensure that we have medium-size avatar images for every # avatar. TODO: This implementation is hacky, both in that it # does get_user_profile_by_id for each user, and in that it # might be better to require the export to just have these. upload_backend = LocalUploadBackend() for record in records: if record['s3_path'].endswith('.original'): user_profile = get_user_profile_by_id(record['user_profile_id']) avatar_path = user_avatar_path_from_ids(user_profile.id, record['realm_id']) medium_file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", avatar_path) + '-medium.png' if os.path.exists(medium_file_path): # We remove the image here primarily to deal with # issues when running the import script multiple # times in development (where one might reuse the # same realm ID from a previous iteration). os.remove(medium_file_path) upload_backend.ensure_medium_avatar_image(user_profile=user_profile)