export/upload: Refactor tarball upload logic to upload_backend.

The conditional block containing the tarball upload logic for both S3
and local uploads was deconstructed and moved to the more appropriate
location within `zerver/lib/upload.py`.

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,))
 

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)

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: