upload: Factor out common code into zerver.lib.upload.

zerver/lib/upload/__init__.py

@@ -20,7 +20,7 @@ from zerver.lib.thumbnail import (
     resize_avatar,
     resize_emoji,
 )
-from zerver.lib.upload.base import ZulipUploadBackend
+from zerver.lib.upload.base import ZulipUploadBackend, create_attachment, sanitize_name
 from zerver.models import Attachment, Message, Realm, RealmEmoji, ScheduledMessage, UserProfile
 
 
@@ -83,14 +83,27 @@ def upload_message_attachment(
     user_profile: UserProfile,
     target_realm: Optional[Realm] = None,
 ) -> str:
-    return upload_backend.upload_message_attachment(
+    if target_realm is None:
-        uploaded_file_name,
+        target_realm = user_profile.realm
+    path_id = upload_backend.generate_message_upload_path(
+        str(target_realm.id), sanitize_name(uploaded_file_name)
+    )
+    upload_backend.upload_message_attachment(
+        path_id,
         uploaded_file_size,
         content_type,
         file_data,
         user_profile,
-        target_realm=target_realm,
+        target_realm,
     )
+    create_attachment(
+        uploaded_file_name,
+        path_id,
+        user_profile,
+        target_realm,
+        uploaded_file_size,
+    )
+    return f"/user_uploads/{path_id}"
 
 
 def claim_attachment(

zerver/lib/upload/base.py

@@ -59,13 +59,13 @@ class ZulipUploadBackend:
 
     def upload_message_attachment(
         self,
-        uploaded_file_name: str,
+        path_id: str,
         uploaded_file_size: int,
         content_type: Optional[str],
         file_data: bytes,
         user_profile: UserProfile,
-        target_realm: Optional[Realm] = None,
+        target_realm: Realm,
-    ) -> str:
+    ) -> None:
         raise NotImplementedError
 
     def save_attachment_contents(self, path_id: str, filehandle: BinaryIO) -> None:

zerver/lib/upload/local.py

@@ -12,7 +12,7 @@ from typing_extensions import override
 from zerver.lib.mime_types import guess_type
 from zerver.lib.thumbnail import resize_avatar, resize_logo
 from zerver.lib.timestamp import timestamp_to_datetime
-from zerver.lib.upload.base import ZulipUploadBackend, create_attachment, sanitize_name
+from zerver.lib.upload.base import ZulipUploadBackend
 from zerver.lib.utils import assert_is_not_none
 from zerver.models import Realm, RealmEmoji, UserProfile
 
@@ -65,35 +65,28 @@ class LocalUploadBackend(ZulipUploadBackend):
         return "/user_avatars/"
 
     @override
-    def generate_message_upload_path(self, realm_id: str, uploaded_file_name: str) -> str:
+    def generate_message_upload_path(self, realm_id: str, sanitized_file_name: str) -> str:
         # Split into 256 subdirectories to prevent directories from getting too big
         return "/".join(
             [
                 realm_id,
                 format(random.randint(0, 255), "x"),
                 secrets.token_urlsafe(18),
-                sanitize_name(uploaded_file_name),
+                sanitized_file_name,
             ]
         )
 
     @override
     def upload_message_attachment(
         self,
-        uploaded_file_name: str,
+        path_id: str,
         uploaded_file_size: int,
         content_type: Optional[str],
         file_data: bytes,
         user_profile: UserProfile,
-        target_realm: Optional[Realm] = None,
+        target_realm: Realm,
-    ) -> str:
+    ) -> None:
-        if target_realm is None:
+        write_local_file("files", path_id, file_data)
-            target_realm = user_profile.realm
-
-        path = self.generate_message_upload_path(str(target_realm.id), uploaded_file_name)
-
-        write_local_file("files", path, file_data)
-        create_attachment(uploaded_file_name, path, user_profile, target_realm, uploaded_file_size)
-        return "/user_uploads/" + path
 
     @override
     def save_attachment_contents(self, path_id: str, filehandle: BinaryIO) -> None:

zerver/lib/upload/s3.py

@@ -14,12 +14,7 @@ from typing_extensions import override
 
 from zerver.lib.mime_types import guess_type
 from zerver.lib.thumbnail import resize_avatar, resize_logo
-from zerver.lib.upload.base import (
+from zerver.lib.upload.base import INLINE_MIME_TYPES, ZulipUploadBackend
-    INLINE_MIME_TYPES,
-    ZulipUploadBackend,
-    create_attachment,
-    sanitize_name,
-)
 from zerver.models import Realm, RealmEmoji, UserProfile
 
 # Duration that the signed upload URLs that we redirect to when
@@ -196,44 +191,34 @@ class S3UploadBackend(ZulipUploadBackend):
         return urljoin(self.public_upload_url_base, key)
 
     @override
-    def generate_message_upload_path(self, realm_id: str, uploaded_file_name: str) -> str:
+    def generate_message_upload_path(self, realm_id: str, sanitized_file_name: str) -> str:
         return "/".join(
             [
                 realm_id,
                 secrets.token_urlsafe(18),
-                sanitize_name(uploaded_file_name),
+                sanitized_file_name,
             ]
         )
 
     @override
     def upload_message_attachment(
         self,
-        uploaded_file_name: str,
+        path_id: str,
         uploaded_file_size: int,
         content_type: Optional[str],
         file_data: bytes,
         user_profile: UserProfile,
-        target_realm: Optional[Realm] = None,
+        target_realm: Realm,
-    ) -> str:
+    ) -> None:
-        if target_realm is None:
-            target_realm = user_profile.realm
-        s3_file_name = self.generate_message_upload_path(str(target_realm.id), uploaded_file_name)
-        url = f"/user_uploads/{s3_file_name}"
-
         upload_image_to_s3(
             self.uploads_bucket,
-            s3_file_name,
+            path_id,
             content_type,
             user_profile,
             file_data,
             settings.S3_UPLOADS_STORAGE_CLASS,
         )
 
-        create_attachment(
-            uploaded_file_name, s3_file_name, user_profile, target_realm, uploaded_file_size
-        )
-        return url
-
     @override
     def save_attachment_contents(self, path_id: str, filehandle: BinaryIO) -> None:
         for chunk in self.uploads_bucket.Object(path_id).get()["Body"]: