diff --git a/zerver/lib/transfer.py b/zerver/lib/transfer.py index 8d88673577..a147496d43 100644 --- a/zerver/lib/transfer.py +++ b/zerver/lib/transfer.py @@ -1,6 +1,7 @@ import logging import os from concurrent.futures import ProcessPoolExecutor, as_completed +from glob import glob import bmemcached import magic @@ -83,6 +84,35 @@ def _transfer_message_files_to_s3(attachment: Attachment) -> None: storage_class=settings.S3_UPLOADS_STORAGE_CLASS, ) logging.info("Uploaded message file in path %s", file_path) + thumbnail_dir = os.path.join(settings.LOCAL_FILES_DIR, "thumbnail", attachment.path_id) + if os.path.isdir(thumbnail_dir): + thumbnails = 0 + for thumbnail_path in glob(os.path.join(thumbnail_dir, "*")): + with open(thumbnail_path, "rb") as f: + # This relies on the thumbnails having guessable + # content-type from their path, in order to avoid + # having to fetch the ImageAttachment inside the + # ProcessPoolExecutor. We also have no clean way + # to prefetch those rows via select_related in the + # outer query, as they match on `path_id`, which + # is not supported as a foreign key. + guessed_type = guess_type(thumbnail_path)[0] + upload_content_to_s3( + s3backend.uploads_bucket, + os.path.join( + "thumbnail", attachment.path_id, os.path.basename(thumbnail_path) + ), + guessed_type, + None, + f.read(), + storage_class=settings.S3_UPLOADS_STORAGE_CLASS, + ) + thumbnails += 1 + logging.info( + "Uploaded %d thumbnails into %s", + thumbnails, + os.path.join("thumbnail", attachment.path_id), + ) except FileNotFoundError: # nocoverage pass diff --git a/zerver/tests/test_transfer.py b/zerver/tests/test_transfer.py index aa9a2383a2..606a92227c 100644 --- a/zerver/tests/test_transfer.py +++ b/zerver/tests/test_transfer.py @@ -13,7 +13,7 @@ from zerver.lib.test_helpers import ( get_test_image_file, read_test_image_file, ) -from zerver.lib.thumbnail import resize_emoji +from zerver.lib.thumbnail import ThumbnailFormat, resize_emoji from zerver.lib.transfer import ( transfer_avatars_to_s3, transfer_emoji_to_s3, @@ -69,13 +69,27 @@ class TransferUploadsToS3Test(ZulipTestCase): upload_message_attachment("dummy1.txt", "text/plain", b"zulip1!", hamlet) upload_message_attachment("dummy2.txt", "text/plain", b"zulip2!", othello) + with ( + self.thumbnail_formats(ThumbnailFormat("webp", 100, 75, animated=False)), + self.captureOnCommitCallbacks(execute=True), + ): + access_path, _ = upload_message_attachment( + "img.png", "image/png", read_test_image_file("img.png"), hamlet + ) + self.assertTrue(access_path.startswith("/user_uploads/")) + image_path_id = access_path.removeprefix("/user_uploads/") + assert settings.LOCAL_FILES_DIR is not None + thumbnail_path = os.path.join( + settings.LOCAL_FILES_DIR, "thumbnail", image_path_id, "100x75.webp" + ) + self.assertTrue(os.path.exists(thumbnail_path)) with self.assertLogs(level="INFO"): transfer_message_files_to_s3(1) attachments = Attachment.objects.all().order_by("id") - self.assert_length(list(bucket.objects.all()), 2) + self.assert_length(list(bucket.objects.all()), 4) s3_dummy1 = bucket.Object(attachments[0].path_id).get() self.assertEqual(s3_dummy1["Body"].read(), b"zulip1!") @@ -91,6 +105,26 @@ class TransferUploadsToS3Test(ZulipTestCase): {"realm_id": str(attachments[1].realm_id), "user_profile_id": str(othello.id)}, ) + s3_image = bucket.Object(attachments[2].path_id).get() + self.assertEqual( + s3_image["Body"].read(), + read_test_image_file("img.png"), + ) + self.assertEqual( + s3_image["Metadata"], + {"realm_id": str(attachments[2].realm_id), "user_profile_id": str(hamlet.id)}, + ) + + s3_image_thumbnail = bucket.Object( + os.path.join("thumbnail", attachments[2].path_id, "100x75.webp") + ).get() + self.assertEqual(s3_image_thumbnail["Metadata"], {}) + with open(thumbnail_path, "rb") as thumbnail_file: + self.assertEqual( + s3_image_thumbnail["Body"].read(), + thumbnail_file.read(), + ) + @mock_aws def test_transfer_emoji_to_s3(self) -> None: bucket = create_s3_buckets(settings.S3_AVATAR_BUCKET)[0]