Mirror
/
zulip
mirror of https://github.com/zulip/zulip.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

upload: Factor out common avatar logic.

This commit is contained in:
Alex Vandiver 2024-06-25 19:03:49 +00:00 committed by Tim Abbott
parent d92993c972
commit 9fb03cb2c7
10 changed files with 176 additions and 227 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
zerver/lib/import_realm.py
View File

@ -34,7 +34,7 @@ from zerver.lib.server_initialization import create_internal_realm, server_initi
from zerver.lib.streams import render_stream_description
from zerver.lib.thumbnail import BadImageError
from zerver.lib.timestamp import datetime_to_timestamp
from zerver.lib.upload import upload_backend
from zerver.lib.upload import ensure_avatar_image, upload_backend
from zerver.lib.upload.base import sanitize_name
from zerver.lib.upload.s3 import get_bucket
from zerver.lib.user_counts import realm_user_count_by_role
@ -790,38 +790,29 @@ def fix_subscriptions_is_user_active_column(
def process_avatars(record: Dict[str, Any]) -> None:
# We need to re-import upload_backend here, because in the
# import-export unit tests, the Zulip settings are overridden for
# specific tests to control the choice of upload backend, and this
# reimport ensures that we use the right choice for the current
# test. Outside the test suite, settings never change after the
# server is started, so this import will have no effect in production.
from zerver.lib.upload import upload_backend
if record["s3_path"].endswith(".original"):
user_profile = get_user_profile_by_id(record["user_profile_id"])
if settings.LOCAL_AVATARS_DIR is not None:
avatar_path = user_avatar_path_from_ids(user_profile.id, record["realm_id"])
medium_file_path = os.path.join(settings.LOCAL_AVATARS_DIR, avatar_path) + "-medium.png"
if os.path.exists(medium_file_path):
# We remove the image here primarily to deal with
# issues when running the import script multiple
# times in development (where one might reuse the
# same realm ID from a previous iteration).
os.remove(medium_file_path)
try:
upload_backend.ensure_avatar_image(user_profile=user_profile, is_medium=True)
if record.get("importer_should_thumbnail"):
upload_backend.ensure_avatar_image(user_profile=user_profile)
except BadImageError:
logging.warning(
"Could not thumbnail avatar image for user %s; ignoring",
user_profile.id,
)
# Delete the record of the avatar to avoid 404s.
do_change_avatar_fields(
user_profile, UserProfile.AVATAR_FROM_GRAVATAR, acting_user=None
)
if not record["s3_path"].endswith(".original"):
return None
user_profile = get_user_profile_by_id(record["user_profile_id"])
if settings.LOCAL_AVATARS_DIR is not None:
avatar_path = user_avatar_path_from_ids(user_profile.id, record["realm_id"])
medium_file_path = os.path.join(settings.LOCAL_AVATARS_DIR, avatar_path) + "-medium.png"
if os.path.exists(medium_file_path):
# We remove the image here primarily to deal with
# issues when running the import script multiple
# times in development (where one might reuse the
# same realm ID from a previous iteration).
os.remove(medium_file_path)
try:
ensure_avatar_image(user_profile=user_profile, medium=True)
if record.get("importer_should_thumbnail"):
ensure_avatar_image(user_profile=user_profile)
except BadImageError:
logging.warning(
"Could not thumbnail avatar image for user %s; ignoring",
user_profile.id,
)
# Delete the record of the avatar to avoid 404s.
do_change_avatar_fields(user_profile, UserProfile.AVATAR_FROM_GRAVATAR, acting_user=None)
def import_uploads(
@ -874,12 +865,7 @@ def import_uploads(
if record["s3_path"].endswith(".original"):
relative_path += ".original"
else:
# TODO: This really should be unconditional. However,
# until we fix the S3 upload backend to use the .png
# path suffix for its normal avatar URLs, we need to
# only do this for the LOCAL_UPLOADS_DIR backend.
if not s3_uploads:
relative_path += ".png"
relative_path = upload_backend.get_avatar_path(relative_path, medium=False)
elif processing_emojis:
# For emojis we follow the function 'upload_emoji_image'
relative_path = RealmEmoji.PATH_ID_TEMPLATE.format(

8
zerver/lib/transfer.py
View File

@ -9,7 +9,7 @@ from django.db import connection
from zerver.lib.avatar_hash import user_avatar_path
from zerver.lib.mime_types import guess_type
from zerver.lib.upload import upload_emoji_image
from zerver.lib.upload import upload_avatar_image, upload_emoji_image
from zerver.lib.upload.s3 import S3UploadBackend, upload_image_to_s3
from zerver.models import Attachment, RealmEmoji, UserProfile
@ -27,10 +27,10 @@ def _transfer_avatar_to_s3(user: UserProfile) -> None:
avatar_path = user_avatar_path(user)
assert settings.LOCAL_UPLOADS_DIR is not None
assert settings.LOCAL_AVATARS_DIR is not None
file_path = os.path.join(settings.LOCAL_AVATARS_DIR, avatar_path) + ".original"
file_path = os.path.join(settings.LOCAL_AVATARS_DIR, avatar_path)
try:
with open(file_path, "rb") as f:
s3backend.upload_avatar_image(f, user, user)
with open(file_path + ".original", "rb") as f:
upload_avatar_image(f, user, user, backend=s3backend)
logging.info("Uploaded avatar for %s in realm %s", user.id, user.realm.name)
except FileNotFoundError:
pass

71
zerver/lib/upload/__init__.py
View File

@ -9,10 +9,11 @@ from django.conf import settings
from django.core.files.uploadedfile import UploadedFile
from django.utils.translation import gettext as _
from zerver.lib.avatar_hash import user_avatar_path
from zerver.lib.exceptions import ErrorCode, JsonableError
from zerver.lib.mime_types import guess_type
from zerver.lib.outgoing_http import OutgoingSession
from zerver.lib.thumbnail import resize_emoji
from zerver.lib.thumbnail import MEDIUM_AVATAR_SIZE, resize_avatar, resize_emoji
from zerver.lib.upload.base import ZulipUploadBackend
from zerver.models import Attachment, Message, Realm, RealmEmoji, ScheduledMessage, UserProfile
@ -140,23 +141,83 @@ def get_avatar_url(hash_key: str, medium: bool = False) -> str:
return upload_backend.get_avatar_url(hash_key, medium)
def write_avatar_images(
file_path: str,
user_profile: UserProfile,
image_data: bytes,
*,
content_type: Optional[str],
backend: Optional[ZulipUploadBackend] = None,
) -> None:
if backend is None:
backend = upload_backend
backend.upload_single_avatar_image(
file_path + ".original",
user_profile=user_profile,
image_data=image_data,
content_type=content_type,
)
backend.upload_single_avatar_image(
backend.get_avatar_path(file_path, medium=False),
user_profile=user_profile,
image_data=resize_avatar(image_data),
content_type="image/png",
)
backend.upload_single_avatar_image(
backend.get_avatar_path(file_path, medium=True),
user_profile=user_profile,
image_data=resize_avatar(image_data, MEDIUM_AVATAR_SIZE),
content_type="image/png",
)
def upload_avatar_image(
user_file: IO[bytes],
acting_user_profile: UserProfile,
target_user_profile: UserProfile,
content_type: Optional[str] = None,
backend: Optional[ZulipUploadBackend] = None,
) -> None:
upload_backend.upload_avatar_image(
user_file, acting_user_profile, target_user_profile, content_type=content_type
if content_type is None:
content_type = guess_type(user_file.name)[0]
file_path = user_avatar_path(target_user_profile)
image_data = user_file.read()
write_avatar_images(
file_path, target_user_profile, image_data, content_type=content_type, backend=backend
)
def copy_avatar(source_profile: UserProfile, target_profile: UserProfile) -> None:
upload_backend.copy_avatar(source_profile, target_profile)
source_file_path = user_avatar_path(source_profile)
target_file_path = user_avatar_path(target_profile)
image_data, content_type = upload_backend.get_avatar_contents(source_file_path)
write_avatar_images(target_file_path, target_profile, image_data, content_type=content_type)
def ensure_avatar_image(user_profile: UserProfile, medium: bool = False) -> None:
file_path = user_avatar_path(user_profile)
image_data, _ = upload_backend.get_avatar_contents(file_path)
if medium:
resized_avatar = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
else:
resized_avatar = resize_avatar(image_data)
upload_backend.upload_single_avatar_image(
upload_backend.get_avatar_path(file_path, medium),
user_profile=user_profile,
image_data=resized_avatar,
content_type="image/png",
)
def delete_avatar_image(user_profile: UserProfile) -> None:
upload_backend.delete_avatar_image(user_profile)
path_id = user_avatar_path(user_profile)
upload_backend.delete_avatar_image(path_id)
# Realm icon and logo uploads

28
zerver/lib/upload/base.py
View File

@ -85,22 +85,26 @@ class ZulipUploadBackend:
def get_avatar_url(self, hash_key: str, medium: bool = False) -> str:
raise NotImplementedError
def upload_avatar_image(
def get_avatar_contents(self, file_path: str) -> Tuple[bytes, str]:
raise NotImplementedError
def get_avatar_path(self, hash_key: str, medium: bool = False) -> str:
if medium:
return f"{hash_key}-medium.png"
else:
return f"{hash_key}.png"
def upload_single_avatar_image(
self,
user_file: IO[bytes],
acting_user_profile: UserProfile,
target_user_profile: UserProfile,
content_type: Optional[str] = None,
file_path: str,
*,
user_profile: UserProfile,
image_data: bytes,
content_type: Optional[str],
) -> None:
raise NotImplementedError
def copy_avatar(self, source_profile: UserProfile, target_profile: UserProfile) -> None:
raise NotImplementedError
def ensure_avatar_image(self, user_profile: UserProfile, is_medium: bool = False) -> None:
raise NotImplementedError
def delete_avatar_image(self, user: UserProfile) -> None:
def delete_avatar_image(self, path_id: str) -> None:
raise NotImplementedError
# Realm icon and logo uploads

77
zerver/lib/upload/local.py
View File

@ -9,8 +9,8 @@ from typing import IO, Any, BinaryIO, Callable, Iterator, Literal, Optional, Tup
from django.conf import settings
from typing_extensions import override
from zerver.lib.avatar_hash import user_avatar_path
from zerver.lib.thumbnail import MEDIUM_AVATAR_SIZE, resize_avatar, resize_logo
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.utils import assert_is_not_none
@ -116,70 +116,35 @@ class LocalUploadBackend(ZulipUploadBackend):
@override
def get_avatar_url(self, hash_key: str, medium: bool = False) -> str:
medium_suffix = "-medium" if medium else ""
return f"/user_avatars/{hash_key}{medium_suffix}.png"
def write_avatar_images(self, file_path: str, image_data: bytes) -> None:
write_local_file("avatars", file_path + ".original", image_data)
resized_data = resize_avatar(image_data)
write_local_file("avatars", file_path + ".png", resized_data)
resized_medium = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
write_local_file("avatars", file_path + "-medium.png", resized_medium)
return "/user_avatars/" + self.get_avatar_path(hash_key, medium)
@override
def upload_avatar_image(
def get_avatar_contents(self, file_path: str) -> Tuple[bytes, str]:
image_data = read_local_file("avatars", file_path + ".original")
content_type = guess_type(file_path)[0]
return image_data, content_type or "application/octet-stream"
@override
def upload_single_avatar_image(
self,
user_file: IO[bytes],
acting_user_profile: UserProfile,
target_user_profile: UserProfile,
content_type: Optional[str] = None,
file_path: str,
*,
user_profile: UserProfile,
image_data: bytes,
content_type: Optional[str],
) -> None:
file_path = user_avatar_path(target_user_profile)
image_data = user_file.read()
self.write_avatar_images(file_path, image_data)
@override
def copy_avatar(self, source_profile: UserProfile, target_profile: UserProfile) -> None:
source_file_path = user_avatar_path(source_profile)
target_file_path = user_avatar_path(target_profile)
image_data = read_local_file("avatars", source_file_path + ".original")
self.write_avatar_images(target_file_path, image_data)
@override
def ensure_avatar_image(self, user_profile: UserProfile, is_medium: bool = False) -> None:
file_extension = "-medium.png" if is_medium else ".png"
file_path = user_avatar_path(user_profile)
output_path = os.path.join(
assert_is_not_none(settings.LOCAL_AVATARS_DIR),
file_path + file_extension,
file_path,
)
if os.path.isfile(output_path):
return
image_path = os.path.join(
assert_is_not_none(settings.LOCAL_AVATARS_DIR),
file_path + ".original",
)
with open(image_path, "rb") as f:
image_data = f.read()
if is_medium:
resized_avatar = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
else:
resized_avatar = resize_avatar(image_data)
write_local_file("avatars", file_path + file_extension, resized_avatar)
if not os.path.isfile(output_path):
write_local_file("avatars", file_path, image_data)
@override
def delete_avatar_image(self, user: UserProfile) -> None:
path_id = user_avatar_path(user)
def delete_avatar_image(self, path_id: str) -> None:
delete_local_file("avatars", path_id + ".original")
delete_local_file("avatars", path_id + ".png")
delete_local_file("avatars", path_id + "-medium.png")
delete_local_file("avatars", self.get_avatar_path(path_id, True))
delete_local_file("avatars", self.get_avatar_path(path_id, False))
@override
def get_realm_icon_url(self, realm_id: int, version: int) -> str:

124
zerver/lib/upload/s3.py
View File

@ -9,12 +9,11 @@ import boto3
import botocore
from botocore.client import Config
from django.conf import settings
from mypy_boto3_s3.service_resource import Bucket, Object
from mypy_boto3_s3.service_resource import Bucket
from typing_extensions import override
from zerver.lib.avatar_hash import user_avatar_path
from zerver.lib.mime_types import guess_type
from zerver.lib.thumbnail import MEDIUM_AVATAR_SIZE, resize_avatar, resize_logo
from zerver.lib.thumbnail import resize_avatar, resize_logo
from zerver.lib.upload.base import (
INLINE_MIME_TYPES,
ZulipUploadBackend,
@ -255,107 +254,48 @@ class S3UploadBackend(ZulipUploadBackend):
item["LastModified"],
)
def write_avatar_images(
@override
def get_avatar_path(self, hash_key: str, medium: bool = False) -> str:
# BUG: The else case should be f"{hash_key}.png".
# See #12852 for details on this bug and how to migrate it.
if medium:
return f"{hash_key}-medium.png"
else:
return hash_key
@override
def get_avatar_url(self, hash_key: str, medium: bool = False) -> str:
return self.get_public_upload_url(self.get_avatar_path(hash_key, medium))
@override
def get_avatar_contents(self, file_path: str) -> Tuple[bytes, str]:
key = self.avatar_bucket.Object(file_path + ".original")
image_data = key.get()["Body"].read()
content_type = key.content_type
return image_data, content_type
@override
def upload_single_avatar_image(
self,
s3_file_name: str,
target_user_profile: UserProfile,
file_path: str,
*,
user_profile: UserProfile,
image_data: bytes,
content_type: Optional[str],
) -> None:
upload_image_to_s3(
self.avatar_bucket,
s3_file_name + ".original",
file_path,
content_type,
target_user_profile,
user_profile,
image_data,
)
# custom 500px wide version
resized_medium = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
upload_image_to_s3(
self.avatar_bucket,
s3_file_name + "-medium.png",
"image/png",
target_user_profile,
resized_medium,
)
resized_data = resize_avatar(image_data)
upload_image_to_s3(
self.avatar_bucket,
s3_file_name,
"image/png",
target_user_profile,
resized_data,
)
# See avatar_url in avatar.py for URL. (That code also handles the case
# that users use gravatar.)
def get_avatar_key(self, file_name: str) -> Object:
key = self.avatar_bucket.Object(file_name)
return key
@override
def get_avatar_url(self, hash_key: str, medium: bool = False) -> str:
medium_suffix = "-medium.png" if medium else ""
return self.get_public_upload_url(f"{hash_key}{medium_suffix}")
@override
def upload_avatar_image(
self,
user_file: IO[bytes],
acting_user_profile: UserProfile,
target_user_profile: UserProfile,
content_type: Optional[str] = None,
) -> None:
if content_type is None:
content_type = guess_type(user_file.name)[0]
s3_file_name = user_avatar_path(target_user_profile)
image_data = user_file.read()
self.write_avatar_images(s3_file_name, target_user_profile, image_data, content_type)
@override
def copy_avatar(self, source_profile: UserProfile, target_profile: UserProfile) -> None:
s3_source_file_name = user_avatar_path(source_profile)
s3_target_file_name = user_avatar_path(target_profile)
key = self.get_avatar_key(s3_source_file_name + ".original")
image_data = key.get()["Body"].read()
content_type = key.content_type
self.write_avatar_images(s3_target_file_name, target_profile, image_data, content_type)
@override
def ensure_avatar_image(self, user_profile: UserProfile, is_medium: bool = False) -> None:
# BUG: The else case should be user_avatar_path(user_profile) + ".png".
# See #12852 for details on this bug and how to migrate it.
file_extension = "-medium.png" if is_medium else ""
file_path = user_avatar_path(user_profile)
s3_file_name = file_path
key = self.avatar_bucket.Object(file_path + ".original")
image_data = key.get()["Body"].read()
if is_medium:
resized_avatar = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
else:
resized_avatar = resize_avatar(image_data)
upload_image_to_s3(
self.avatar_bucket,
s3_file_name + file_extension,
"image/png",
user_profile,
resized_avatar,
)
@override
def delete_avatar_image(self, user: UserProfile) -> None:
path_id = user_avatar_path(user)
def delete_avatar_image(self, path_id: str) -> None:
self.delete_file_from_s3(path_id + ".original", self.avatar_bucket)
self.delete_file_from_s3(path_id + "-medium.png", self.avatar_bucket)
self.delete_file_from_s3(path_id, self.avatar_bucket)
self.delete_file_from_s3(self.get_avatar_path(path_id, True), self.avatar_bucket)
self.delete_file_from_s3(self.get_avatar_path(path_id, False), self.avatar_bucket)
@override
def get_realm_icon_url(self, realm_id: int, version: int) -> str:

7
zerver/migrations/0032_verify_all_medium_avatar_images.py
View File

@ -6,7 +6,7 @@ from django.db import migrations
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps
from zerver.lib.upload import upload_backend
from zerver.lib.upload import ensure_avatar_image
from zerver.models import UserProfile
@ -22,12 +22,11 @@ def patched_user_avatar_path(user_profile: UserProfile) -> str:
return hashlib.sha1(user_key.encode()).hexdigest()
@patch("zerver.lib.upload.s3.user_avatar_path", patched_user_avatar_path)
@patch("zerver.lib.upload.local.user_avatar_path", patched_user_avatar_path)
@patch("zerver.lib.upload.user_avatar_path", patched_user_avatar_path)
def verify_medium_avatar_image(apps: StateApps, schema_editor: BaseDatabaseSchemaEditor) -> None:
user_profile_model = apps.get_model("zerver", "UserProfile")
for user_profile in user_profile_model.objects.filter(avatar_source="U"):
upload_backend.ensure_avatar_image(user_profile, is_medium=True)
ensure_avatar_image(user_profile, medium=True)
class Migration(migrations.Migration):

6
zerver/tests/test_upload.py
View File

@ -1269,16 +1269,14 @@ class AvatarTest(UploadSerializeMixin, ZulipTestCase):
with mock.patch(
"zerver.lib.upload.local.write_local_file"
) as mock_write_local_file:
zerver.lib.upload.upload_backend.ensure_avatar_image(
user_profile, is_medium=True
)
zerver.lib.upload.ensure_avatar_image(user_profile, medium=True)
self.assertFalse(mock_write_local_file.called)
# Confirm that ensure_medium_avatar_url works to recreate
# medium size avatars from the original if needed
os.remove(medium_avatar_disk_path)
self.assertFalse(os.path.exists(medium_avatar_disk_path))
zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile, is_medium=True)
zerver.lib.upload.ensure_avatar_image(user_profile, medium=True)
self.assertTrue(os.path.exists(medium_avatar_disk_path))
# Verify whether the avatar_version gets incremented with every new upload

4
zerver/tests/test_upload_local.py
View File

@ -164,13 +164,13 @@ class LocalStorageTest(UploadSerializeMixin, ZulipTestCase):
image_data = f.read()
resized_avatar = resize_avatar(image_data)
zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile)
zerver.lib.upload.ensure_avatar_image(user_profile)
output_path = os.path.join(settings.LOCAL_AVATARS_DIR, file_path + ".png")
with open(output_path, "rb") as original_file:
self.assertEqual(resized_avatar, original_file.read())
resized_avatar = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile, is_medium=True)
zerver.lib.upload.ensure_avatar_image(user_profile, medium=True)
output_path = os.path.join(settings.LOCAL_AVATARS_DIR, file_path + "-medium.png")
with open(output_path, "rb") as original_file:
self.assertEqual(resized_avatar, original_file.read())

14
zerver/tests/test_upload_s3.py
View File

@ -277,9 +277,7 @@ class S3Test(ZulipTestCase):
medium_path_id = path_id + "-medium.png"
with get_test_image_file("img.png") as image_file:
zerver.lib.upload.upload_backend.upload_avatar_image(
image_file, user_profile, user_profile
)
zerver.lib.upload.upload_avatar_image(image_file, user_profile, user_profile)
test_image_data = read_test_image_file("img.png")
test_medium_image_data = resize_avatar(test_image_data, MEDIUM_AVATAR_SIZE)
@ -294,7 +292,7 @@ class S3Test(ZulipTestCase):
self.assertEqual(medium_image_data, test_medium_image_data)
bucket.Object(medium_image_key.key).delete()
zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile, is_medium=True)
zerver.lib.upload.ensure_avatar_image(user_profile, medium=True)
medium_image_key = bucket.Object(medium_path_id)
self.assertEqual(medium_image_key.key, medium_path_id)
@ -356,19 +354,17 @@ class S3Test(ZulipTestCase):
medium_file_path = base_file_path + "-medium.png"
with get_test_image_file("img.png") as image_file:
zerver.lib.upload.upload_backend.upload_avatar_image(
image_file, user_profile, user_profile
)
zerver.lib.upload.upload_avatar_image(image_file, user_profile, user_profile)
key = bucket.Object(original_file_path)
image_data = key.get()["Body"].read()
zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile)
zerver.lib.upload.ensure_avatar_image(user_profile)
resized_avatar = resize_avatar(image_data)
key = bucket.Object(file_path)
self.assertEqual(resized_avatar, key.get()["Body"].read())
zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile, is_medium=True)
zerver.lib.upload.ensure_avatar_image(user_profile, medium=True)
resized_avatar = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
key = bucket.Object(medium_file_path)
self.assertEqual(resized_avatar, key.get()["Body"].read())