upstream-zulip/zerver/migrations/0770_recreate_missing_realmemoji.py

from typing import Any

import orjson
from django.db import connection, migrations, transaction
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps
from django.db.models import Exists, OuterRef

ROLE_REALM_OWNER = 100
REALM_EMOJI_ADDED = 226


def recreate_missing_realmemoji(apps: StateApps, schema_editor: BaseDatabaseSchemaEditor) -> None:
    RealmEmoji = apps.get_model("zerver", "RealmEmoji")
    RealmAuditLog = apps.get_model("zerver", "RealmAuditLog")
    Realm = apps.get_model("zerver", "Realm")
    UserProfile = apps.get_model("zerver", "UserProfile")

    def get_first_human_owner_user(realm: Any) -> Any:
        first_active_human_owner = UserProfile.objects.filter(
            realm=realm, is_bot=False, role=ROLE_REALM_OWNER, is_active=True
        ).first()
        if first_active_human_owner is not None:
            return first_active_human_owner

        fallback_owner_user = UserProfile.objects.filter(
            realm=realm, is_bot=False, role=ROLE_REALM_OWNER, is_active=False
        ).first()
        if fallback_owner_user is None:
            raise AssertionError(f"No human owner users in realm {realm.string_id}")
        return fallback_owner_user

    emoji_add_log_queryset = RealmAuditLog.objects.filter(
        realm=OuterRef("pk"),
        event_type=REALM_EMOJI_ADDED,
    )
    for realm in (
        Realm.objects.annotate(has_emoji_add_log=Exists(emoji_add_log_queryset))
        .filter(has_emoji_add_log=True)
        .iterator()
    ):
        with transaction.atomic():
            first_owner = None
            sql = """
                WITH distinct_emoji AS (
                    SELECT DISTINCT emoji.value AS e
                    FROM zerver_realmauditlog
                    CROSS JOIN LATERAL jsonb_each(
                        zerver_realmauditlog.extra_data->'realm_emoji'
                    ) AS emoji(key, value)
                    WHERE zerver_realmauditlog.realm_id = %s
                    AND zerver_realmauditlog.event_type = %s
                )
                SELECT distinct_emoji.e FROM distinct_emoji
                LEFT JOIN zerver_realmemoji
                    ON zerver_realmemoji.id = (e->>'id')::integer
                    OR (zerver_realmemoji.name = e->>'name'
                        AND zerver_realmemoji.realm_id = %s)
                WHERE zerver_realmemoji.id IS NULL
            """
            with connection.cursor() as cursor:
                cursor.execute(sql, [realm.id, REALM_EMOJI_ADDED, realm.id])
                for (raw_emoji,) in cursor.fetchall():
                    orphaned_emoji = orjson.loads(raw_emoji)
                    print(
                        f"Creating missing emoji {orphaned_emoji['id']} = {realm.string_id} / {orphaned_emoji['name']}"
                    )
                    file_name = orphaned_emoji["source_url"][
                        1 + orphaned_emoji["source_url"].rfind("/") :
                    ]
                    if first_owner is None:
                        first_owner = get_first_human_owner_user(realm)

                    RealmEmoji.objects.create(
                        id=int(orphaned_emoji["id"]),
                        realm_id=realm.id,
                        name=orphaned_emoji["name"],
                        author_id=first_owner.id,
                        file_name=file_name,
                        deactivated=orphaned_emoji["deactivated"],
                        is_animated=orphaned_emoji["still_url"] is not None,
                    )


class Migration(migrations.Migration):
    atomic = False

    dependencies = [
        ("zerver", "0769_rename_profile_field_twitter_to_x"),
    ]

    operations = [
        migrations.RunPython(
            recreate_missing_realmemoji,
            elidable=True,
        ),
    ]