diff --git a/zerver/lib/export.py b/zerver/lib/export.py index da6e6d8422..aab6bf3fe0 100644 --- a/zerver/lib/export.py +++ b/zerver/lib/export.py @@ -21,7 +21,6 @@ import ujson import subprocess import tempfile import shutil -import sys from scripts.lib.zulip_tools import overwrite_symlink from zerver.lib.avatar_hash import user_avatar_path_from_ids from analytics.models import RealmCount, UserCount, StreamCount @@ -33,8 +32,7 @@ from zerver.models import UserProfile, Realm, Client, Huddle, Stream, \ RealmAuditLog, UserHotspot, MutedTopic, Service, UserGroup, \ UserGroupMembership, BotStorageData, BotConfigData from zerver.lib.parallel import run_parallel -from zerver.lib.utils import generate_random_token -from zerver.lib.upload import random_name, get_bucket +from zerver.lib.upload import upload_backend from typing import Any, Callable, Dict, List, Optional, Set, Tuple, \ Union @@ -1689,39 +1687,11 @@ def export_realm_wrapper(realm: Realm, output_dir: str, if not upload: return None - def percent_callback(complete: Any, total: Any) -> None: - sys.stdout.write('.') - sys.stdout.flush() - - print("Uploading export tarball...") # We upload to the `avatars` bucket because that's world-readable # without additional configuration. We'll likely want to change - # that in the future, after moving the below code into - # `zerver/lib/upload.py`. - if settings.LOCAL_UPLOADS_DIR is not None: - path = os.path.join( - 'exports', - str(realm.id), - random_name(18), - os.path.basename(tarball_path), - ) - abs_path = os.path.join(settings.LOCAL_UPLOADS_DIR, 'avatars', path) - os.makedirs(os.path.dirname(abs_path), exist_ok=True) - shutil.copy(tarball_path, abs_path) - public_url = realm.uri + '/user_avatars/' + path - else: - conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY) - # We use the avatar bucket, because it's world-readable. - bucket = get_bucket(conn, settings.S3_AVATAR_BUCKET) - key = Key(bucket) - key.key = os.path.join("exports", generate_random_token(32), os.path.basename(tarball_path)) - key.set_contents_from_filename(tarball_path, cb=percent_callback, num_cb=40) - - public_url = 'https://{bucket}.{host}/{key}'.format( - host=conn.server_name(), - bucket=bucket.name, - key=key.key) - + # that in the future. + print("Uploading export tarball...") + public_url = upload_backend.upload_export_tarball(realm, tarball_path) print() print("Uploaded to %s" % (public_url,)) diff --git a/zerver/lib/upload.py b/zerver/lib/upload.py index 314441c890..9bab2bd94d 100644 --- a/zerver/lib/upload.py +++ b/zerver/lib/upload.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional, Tuple +from typing import Dict, Optional, Tuple, Any from django.utils.translation import ugettext as _ from django.conf import settings @@ -19,6 +19,8 @@ from zerver.models import get_user_profile_by_id from zerver.models import Attachment from zerver.models import Realm, RealmEmoji, UserProfile, Message +from zerver.lib.utils import generate_random_token + import urllib import base64 import os @@ -29,6 +31,8 @@ from PIL.GifImagePlugin import GifImageFile import io import random import logging +import shutil +import sys DEFAULT_AVATAR_SIZE = 100 MEDIUM_AVATAR_SIZE = 500 @@ -238,6 +242,9 @@ class ZulipUploadBackend: def get_emoji_url(self, emoji_file_name: str, realm_id: int) -> str: raise NotImplementedError() + def upload_export_tarball(self, realm: Realm, tarball_path: str) -> str: + raise NotImplementedError() + ### S3 @@ -573,6 +580,24 @@ class S3UploadBackend(ZulipUploadBackend): emoji_file_name=emoji_file_name) return "https://%s.s3.amazonaws.com/%s" % (bucket, emoji_path) + def upload_export_tarball(self, realm: Optional[Realm], tarball_path: str) -> str: + def percent_callback(complete: Any, total: Any) -> None: + sys.stdout.write('.') + sys.stdout.flush() + + conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY) + # We use the avatar bucket, because it's world-readable. + bucket = get_bucket(conn, settings.S3_AVATAR_BUCKET) + key = Key(bucket) + key.key = os.path.join("exports", generate_random_token(32), os.path.basename(tarball_path)) + key.set_contents_from_filename(tarball_path, cb=percent_callback, num_cb=40) + + public_url = 'https://{bucket}.{host}/{key}'.format( + host=conn.server_name(), + bucket=bucket.name, + key=key.key) + return public_url + ### Local @@ -750,6 +775,19 @@ class LocalUploadBackend(ZulipUploadBackend): "/user_avatars", RealmEmoji.PATH_ID_TEMPLATE.format(realm_id=realm_id, emoji_file_name=emoji_file_name)) + def upload_export_tarball(self, realm: Realm, tarball_path: str) -> str: + path = os.path.join( + 'exports', + str(realm.id), + random_name(18), + os.path.basename(tarball_path), + ) + abs_path = os.path.join(settings.LOCAL_UPLOADS_DIR, 'avatars', path) + os.makedirs(os.path.dirname(abs_path), exist_ok=True) + shutil.copy(tarball_path, abs_path) + public_url = realm.uri + '/user_avatars/' + path + return public_url + # Common and wrappers if settings.LOCAL_UPLOADS_DIR is not None: upload_backend = LocalUploadBackend() # type: ZulipUploadBackend @@ -810,3 +848,6 @@ def upload_message_image_from_request(request: HttpRequest, user_file: File, uploaded_file_name, uploaded_file_size, content_type = get_file_info(request, user_file) return upload_message_file(uploaded_file_name, uploaded_file_size, content_type, user_file.read(), user_profile) + +def upload_export_tarball(realm: Realm, tarball_path: str) -> str: + return upload_backend.upload_export_tarball(realm, tarball_path) diff --git a/zerver/tests/test_upload.py b/zerver/tests/test_upload.py index 1feb9a8a96..fef78078d4 100644 --- a/zerver/tests/test_upload.py +++ b/zerver/tests/test_upload.py @@ -23,7 +23,8 @@ from zerver.lib.upload import sanitize_name, S3UploadBackend, \ upload_message_file, upload_emoji_image, delete_message_image, LocalUploadBackend, \ ZulipUploadBackend, MEDIUM_AVATAR_SIZE, resize_avatar, \ resize_emoji, BadImageError, get_realm_for_filename, \ - DEFAULT_AVATAR_SIZE, DEFAULT_EMOJI_SIZE, exif_rotate + DEFAULT_AVATAR_SIZE, DEFAULT_EMOJI_SIZE, exif_rotate, \ + upload_export_tarball import zerver.lib.upload from zerver.models import Attachment, get_user, \ Message, UserProfile, Realm, \ @@ -41,6 +42,8 @@ from zerver.lib.cache import get_realm_used_upload_space_cache_key, cache_get from zerver.lib.create_user import copy_user_settings from zerver.lib.users import get_api_key +from scripts.lib.zulip_tools import get_or_create_dev_uuid_var_path + import urllib import ujson from PIL import Image @@ -1440,6 +1443,28 @@ class LocalStorageTest(UploadSerializeMixin, ZulipTestCase): expected_url = "/user_avatars/{emoji_path}".format(emoji_path=emoji_path) self.assertEqual(expected_url, url) + def test_tarball_upload_local(self) -> None: + user_profile = self.example_user("iago") + self.assertTrue(user_profile.is_realm_admin) + + tarball_path = os.path.join(get_or_create_dev_uuid_var_path('test-backend'), + 'tarball.tar.gz') + with open(tarball_path, 'w') as f: + f.write('dummy') + + uri = upload_export_tarball(user_profile.realm, tarball_path) + self.assertTrue(os.path.isfile(os.path.join(settings.LOCAL_UPLOADS_DIR, + 'avatars', + tarball_path))) + + result = re.search(re.compile(r"([A-Za-z0-9\-_]{24})"), uri) + if result is not None: + random_name = result.group(1) + expected_url = "http://zulip.testserver/user_avatars/exports/1/{random_name}/tarball.tar.gz".format( + random_name=random_name, + ) + self.assertEqual(expected_url, uri) + def tearDown(self) -> None: destroy_uploads() @@ -1709,6 +1734,29 @@ class S3Test(ZulipTestCase): expected_url = "https://{bucket}.s3.amazonaws.com/{path}".format(bucket=bucket, path=path) self.assertEqual(expected_url, url) + @use_s3_backend + def test_tarball_upload(self) -> None: + bucket = create_s3_buckets(settings.S3_AVATAR_BUCKET)[0] + + user_profile = self.example_user("iago") + self.assertTrue(user_profile.is_realm_admin) + + tarball_path = os.path.join(get_or_create_dev_uuid_var_path('test-backend'), + 'tarball.tar.gz') + with open(tarball_path, 'w') as f: + f.write('dummy') + + uri = upload_export_tarball(user_profile.realm, tarball_path) + + result = re.search(re.compile(r"([0-9a-fA-F]{32})"), uri) + if result is not None: + hex_value = result.group(1) + expected_url = "https://{bucket}.s3.amazonaws.com:443/exports/{hex_value}/{path}".format( + bucket=bucket.name, + hex_value=hex_value, + path=os.path.basename(tarball_path)) + self.assertEqual(uri, expected_url) + class UploadTitleTests(TestCase): def test_upload_titles(self) -> None: