import: Move 'process_uploads' and 'process_emojis' to import_util.

zerver/data_import/import_util.py

@@ -2,11 +2,12 @@ import random
 import requests
 import shutil
 import logging
+import os
 
 from typing import List, Dict, Any
 from django.forms.models import model_to_dict
 
-from zerver.models import Realm
+from zerver.models import Realm, RealmEmoji
 from zerver.lib.actions import STREAM_ASSIGNMENT_COLORS as stream_colors
 from zerver.lib.avatar_hash import user_avatar_path_from_ids
 from zerver.lib.parallel import run_parallel
@@ -112,3 +113,89 @@ def process_avatars(avatar_list: List[ZerverFieldsT], avatar_dir: str, realm_id:
 
     logging.info('######### GETTING AVATARS FINISHED #########\n')
     return avatar_list + avatar_original_list
+
+def process_uploads(upload_list: List[ZerverFieldsT], upload_dir: str,
+                    threads: int) -> List[ZerverFieldsT]:
+    """
+    This function downloads the uploads and saves it in the realm's upload directory.
+    Required parameters:
+
+    1. upload_list: List of uploads to be mapped in uploads records.json file
+    2. upload_dir: Folder where the downloaded uploads are saved
+    """
+    def get_uploads(upload: List[str]) -> int:
+        upload_url = upload[0]
+        upload_path = upload[1]
+        upload_path = os.path.join(upload_dir, upload_path)
+
+        response = requests.get(upload_url, stream=True)
+        os.makedirs(os.path.dirname(upload_path), exist_ok=True)
+        with open(upload_path, 'wb') as upload_file:
+            shutil.copyfileobj(response.raw, upload_file)
+        return 0
+
+    logging.info('######### GETTING ATTACHMENTS #########\n')
+    logging.info('DOWNLOADING ATTACHMENTS .......\n')
+    upload_url_list = []
+    for upload in upload_list:
+        upload_url = upload['path']
+        upload_s3_path = upload['s3_path']
+        upload_url_list.append([upload_url, upload_s3_path])
+        upload['path'] = upload_s3_path
+
+    # Run downloads parallely
+    output = []
+    for (status, job) in run_parallel(get_uploads, upload_url_list, threads=threads):
+        output.append(job)
+
+    logging.info('######### GETTING ATTACHMENTS FINISHED #########\n')
+    return upload_list
+
+def process_emojis(zerver_realmemoji: List[ZerverFieldsT], emoji_dir: str,
+                   emoji_url_map: ZerverFieldsT, threads: int) -> List[ZerverFieldsT]:
+    """
+    This function downloads the custom emojis and saves in the output emoji folder.
+    Required parameters:
+
+    1. zerver_realmemoji: List of all RealmEmoji objects to be imported
+    2. emoji_dir: Folder where the downloaded emojis are saved
+    3. emoji_url_map: Maps emoji name to its url
+    """
+    def get_emojis(upload: List[str]) -> int:
+        emoji_url = upload[0]
+        emoji_path = upload[1]
+        upload_emoji_path = os.path.join(emoji_dir, emoji_path)
+
+        response = requests.get(emoji_url, stream=True)
+        os.makedirs(os.path.dirname(upload_emoji_path), exist_ok=True)
+        with open(upload_emoji_path, 'wb') as emoji_file:
+            shutil.copyfileobj(response.raw, emoji_file)
+        return 0
+
+    emoji_records = []
+    upload_emoji_list = []
+    logging.info('######### GETTING EMOJIS #########\n')
+    logging.info('DOWNLOADING EMOJIS .......\n')
+    for emoji in zerver_realmemoji:
+        emoji_url = emoji_url_map[emoji['name']]
+        emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format(
+            realm_id=emoji['realm'],
+            emoji_file_name=emoji['name'])
+
+        upload_emoji_list.append([emoji_url, emoji_path])
+
+        emoji_record = dict(emoji)
+        emoji_record['path'] = emoji_path
+        emoji_record['s3_path'] = emoji_path
+        emoji_record['realm_id'] = emoji_record['realm']
+        emoji_record.pop('realm')
+
+        emoji_records.append(emoji_record)
+
+    # Run downloads parallely
+    output = []
+    for (status, job) in run_parallel(get_emojis, upload_emoji_list, threads=threads):
+        output.append(job)
+
+    logging.info('######### GETTING EMOJIS FINISHED #########\n')
+    return emoji_records