mirror of https://github.com/zulip/zulip.git
113 lines
4.5 KiB
Python
113 lines
4.5 KiB
Python
import os
|
|
import shutil
|
|
|
|
import boto3
|
|
from django.conf import settings
|
|
from django.db import migrations, models
|
|
from django.db.backends.postgresql.schema import DatabaseSchemaEditor
|
|
from django.db.migrations.state import StateApps
|
|
|
|
|
|
class Uploader:
|
|
def __init__(self) -> None:
|
|
self.old_orig_image_path_template = "{realm_id}/emoji/{emoji_file_name}.original"
|
|
self.old_path_template = "{realm_id}/emoji/{emoji_file_name}"
|
|
self.new_orig_image_path_template = "{realm_id}/emoji/images/{emoji_file_name}.original"
|
|
self.new_path_template = "{realm_id}/emoji/images/{emoji_file_name}"
|
|
|
|
def copy_files(self, src_path: str, dst_path: str) -> None:
|
|
raise NotImplementedError()
|
|
|
|
def ensure_emoji_images(self, realm_id: int, old_filename: str, new_filename: str) -> None:
|
|
# Copy original image file.
|
|
old_file_path = self.old_orig_image_path_template.format(realm_id=realm_id,
|
|
emoji_file_name=old_filename)
|
|
new_file_path = self.new_orig_image_path_template.format(realm_id=realm_id,
|
|
emoji_file_name=new_filename)
|
|
self.copy_files(old_file_path, new_file_path)
|
|
|
|
# Copy resized image file.
|
|
old_file_path = self.old_path_template.format(realm_id=realm_id,
|
|
emoji_file_name=old_filename)
|
|
new_file_path = self.new_path_template.format(realm_id=realm_id,
|
|
emoji_file_name=new_filename)
|
|
self.copy_files(old_file_path, new_file_path)
|
|
|
|
class LocalUploader(Uploader):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
@staticmethod
|
|
def mkdirs(path: str) -> None:
|
|
dirname = os.path.dirname(path)
|
|
if not os.path.isdir(dirname):
|
|
os.makedirs(dirname)
|
|
|
|
def copy_files(self, src_path: str, dst_path: str) -> None:
|
|
src_path = os.path.join(settings.LOCAL_UPLOADS_DIR, 'avatars', src_path)
|
|
self.mkdirs(src_path)
|
|
dst_path = os.path.join(settings.LOCAL_UPLOADS_DIR, 'avatars', dst_path)
|
|
self.mkdirs(dst_path)
|
|
shutil.copyfile(src_path, dst_path)
|
|
|
|
class S3Uploader(Uploader):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
session = boto3.Session(settings.S3_KEY, settings.S3_SECRET_KEY)
|
|
self.bucket_name = settings.S3_AVATAR_BUCKET
|
|
self.bucket = session.resource('s3').Bucket(self.bucket_name)
|
|
|
|
def copy_files(self, src_key: str, dst_key: str) -> None:
|
|
source = dict(Bucket=self.bucket_name, Key=src_key)
|
|
self.bucket.copy(source, dst_key)
|
|
|
|
def get_uploader() -> Uploader:
|
|
if settings.LOCAL_UPLOADS_DIR is None:
|
|
return S3Uploader()
|
|
return LocalUploader()
|
|
|
|
def get_emoji_file_name(emoji_file_name: str, new_name: str) -> str:
|
|
_, image_ext = os.path.splitext(emoji_file_name)
|
|
return ''.join((new_name, image_ext))
|
|
|
|
def migrate_realm_emoji_image_files(apps: StateApps, schema_editor: DatabaseSchemaEditor) -> None:
|
|
RealmEmoji = apps.get_model('zerver', 'RealmEmoji')
|
|
uploader = get_uploader()
|
|
for realm_emoji in RealmEmoji.objects.all():
|
|
old_file_name = realm_emoji.file_name
|
|
new_file_name = get_emoji_file_name(old_file_name, str(realm_emoji.id))
|
|
uploader.ensure_emoji_images(realm_emoji.realm_id, old_file_name, new_file_name)
|
|
realm_emoji.file_name = new_file_name
|
|
realm_emoji.save(update_fields=['file_name'])
|
|
|
|
def reversal(apps: StateApps, schema_editor: DatabaseSchemaEditor) -> None:
|
|
# Ensures that migration can be re-run in case of a failure.
|
|
RealmEmoji = apps.get_model('zerver', 'RealmEmoji')
|
|
for realm_emoji in RealmEmoji.objects.all():
|
|
corrupt_file_name = realm_emoji.file_name
|
|
correct_file_name = get_emoji_file_name(corrupt_file_name, realm_emoji.name)
|
|
realm_emoji.file_name = correct_file_name
|
|
realm_emoji.save(update_fields=['file_name'])
|
|
|
|
class Migration(migrations.Migration):
|
|
|
|
dependencies = [
|
|
('zerver', '0148_max_invites_forget_default'),
|
|
]
|
|
|
|
operations = [
|
|
migrations.AlterUniqueTogether(
|
|
name='realmemoji',
|
|
unique_together=set(),
|
|
),
|
|
migrations.AlterField(
|
|
model_name='realmemoji',
|
|
name='file_name',
|
|
field=models.TextField(db_index=True, null=True, blank=True),
|
|
),
|
|
migrations.RunPython(
|
|
migrate_realm_emoji_image_files,
|
|
reverse_code=reversal,
|
|
elidable=True),
|
|
]
|