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")
import logging
import os
from mimetypes import guess_type
from django.conf import settings
from django.db import connection
from zerver.lib.avatar_hash import user_avatar_path
from zerver.lib.parallel import run_parallel
from zerver.lib.upload import S3UploadBackend, upload_image_to_s3
from zerver.models import Attachment, RealmEmoji, UserProfile
s3backend = S3UploadBackend()
def transfer_uploads_to_s3(processes: int) -> None:
# TODO: Eventually, we'll want to add realm icon and logo
transfer_avatars_to_s3(processes)
transfer_message_files_to_s3(processes)
transfer_emoji_to_s3(processes)
def transfer_avatars_to_s3(processes: int) -> None:
def _transfer_avatar_to_s3(user: UserProfile) -> int:
avatar_path = user_avatar_path(user)
file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars",
avatar_path) + ".original"
try:
with open(file_path, 'rb') as f:
s3backend.upload_avatar_image(f, user, user)
logging.info("Uploaded avatar for %s in realm %s", user.id,
def import_uploads_s3(bucket_name: str,
import_dir: Path,
processing_avatars: bool = False,
processing_emojis: bool = False) -> None:
upload_backend = S3UploadBackend()
conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY)
bucket = conn.get_bucket(bucket_name, validate=True)
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)
re_map_foreign_keys_internal(records,
'records',
'user_profile_id',
related_table="user_profile",
id_field=True)
for record in records:
key = Key(bucket)
if processing_avatars:
# For avatars, we need to rehash the user's email with the
# new server's avatar salt
avatar_path = user_avatar_path_from_ids(record['user_profile_id'],
record['realm_id'])
key.key = avatar_path
if record['s3_path'].endswith('.original'):
key.key += '.original'
if 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'])
key.key = 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']))
])
key.key = s3_file_name
path_maps['attachment_path'][record['path']] = s3_file_name
user_profile_id = int(record['user_profile_id'])
# Support email gateway bot and other cross-realm messages
if user_profile_id in id_maps["user_profile"]:
logging.info("Uploaded by ID mapped user: %s!" %
(user_profile_id, ))
user_profile_id = id_maps["user_profile"][user_profile_id]
user_profile = get_user_profile_by_id(user_profile_id)
key.set_metadata("user_profile_id", str(user_profile.id))
key.set_metadata("realm_id", str(user_profile.realm_id))
key.set_metadata("orig_last_modified", record['last_modified'])
headers = {'Content-Type': record['content_type']}
key.set_contents_from_filename(os.path.join(import_dir,
record['path']),
headers=headers)
if processing_avatars:
# TODO: Ideally, we'd do this in a separate pass, after
# all the avatars have been uploaded, since we may end up
# unnecssarily resizing images just before the medium-size
# image in the export is uploaded. See the local uplods
# code path for more notes.
upload_backend.ensure_medium_avatar_image(
user_profile=user_profile)
def import_uploads_s3(bucket_name: str, import_dir: Path, processing_avatars: bool=False,
processing_emojis: bool=False) -> None:
upload_backend = S3UploadBackend()
conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY)
bucket = conn.get_bucket(bucket_name, validate=True)
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)
timestamp = datetime_to_timestamp(timezone_now())
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:
key = Key(bucket)
if processing_avatars:
# For avatars, we need to rehash the user's email with the
# new server's avatar salt
avatar_path = user_avatar_path_from_ids(record['user_profile_id'], record['realm_id'])
key.key = avatar_path
if record['s3_path'].endswith('.original'):
key.key += '.original'
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'])
key.key = emoji_path
record['last_modified'] = timestamp
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']))
])
key.key = s3_file_name
path_maps['attachment_path'][record['s3_path']] = s3_file_name
# 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:
user_profile_id = int(record['user_profile_id'])
# Support email gateway bot and other cross-realm messages
if user_profile_id in id_maps["user_profile"]:
logging.info("Uploaded by ID mapped user: %s!" % (user_profile_id,))
user_profile_id = id_maps["user_profile"][user_profile_id]
user_profile = get_user_profile_by_id(user_profile_id)
key.set_metadata("user_profile_id", str(user_profile.id))
key.set_metadata("orig_last_modified", record['last_modified'])
key.set_metadata("realm_id", str(record['realm_id']))
# Zulip exports will always have a content-type, but third-party exports might not.
content_type = record.get("content_type", guess_type(record['s3_path'])[0])
headers = {'Content-Type': content_type}
key.set_contents_from_filename(os.path.join(import_dir, record['path']), headers=headers)
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 = S3UploadBackend()
for record in records:
if record['s3_path'].endswith('.original'):
user_profile = get_user_profile_by_id(record['user_profile_id'])
upload_backend.ensure_medium_avatar_image(user_profile=user_profile)
