CodeCow/zulip
1
0
Fork 0
mirror of https://github.com/zulip/zulip.git synced 2026-06-21 21:32:29 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

import: Support importing realm icon and logo.

Browse Source 
Fixes #11216
This commit is contained in:
Vishnu Ks 2019-07-19 22:45:23 +05:30 committed by Tim Abbott
parent af3a37b58b
commit 2ea53a347a
3 changed files with 191 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

82
zerver/lib/export.py
View File

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

38
zerver/lib/import_realm.py
View File

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

105
zerver/tests/test_import_export.py
View File

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