upstream-zulip/zerver/migrations/0810_delete_unused_anonymous_groups.py

from django.db import connection, migrations, transaction
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps
from psycopg2.sql import SQL, Literal


def delete_unused_anonymous_groups(
    apps: StateApps, schema_editor: BaseDatabaseSchemaEditor
) -> None:
    Realm = apps.get_model("zerver", "Realm")

    query = """
        WITH user_group_ids AS (
            SELECT table_cols.value
            FROM zerver_stream
            CROSS JOIN LATERAL (
                VALUES
                    (can_add_subscribers_group_id),
                    (can_administer_channel_group_id),
                    (can_create_topic_group_id),
                    (can_delete_any_message_group_id),
                    (can_delete_own_message_group_id),
                    (can_move_messages_out_of_channel_group_id),
                    (can_move_messages_within_channel_group_id),
                    (can_remove_subscribers_group_id),
                    (can_send_message_group_id),
                    (can_subscribe_group_id),
                    (can_resolve_topics_group_id)
            ) AS table_cols(value)
            WHERE realm_id = %s

            UNION ALL

            SELECT table_cols.value
            FROM zerver_namedusergroup
            CROSS JOIN LATERAL (
                VALUES
                    (can_add_members_group_id),
                    (can_join_group_id),
                    (can_leave_group_id),
                    (can_manage_group_id),
                    (can_mention_group_id),
                    (can_remove_members_group_id)
            ) AS table_cols(value)
            WHERE realm_id = %s

            UNION ALL

            SELECT realm_cols.value
            FROM zerver_realm
            CROSS JOIN LATERAL (
                VALUES
                    (create_multiuse_invite_group_id),
                    (can_access_all_users_group_id),
                    (can_add_subscribers_group_id),
                    (can_add_custom_emoji_group_id),
                    (can_create_bots_group_id),
                    (can_create_groups_id),
                    (can_create_public_channel_group_id),
                    (can_create_private_channel_group_id),
                    (can_create_web_public_channel_group_id),
                    (can_create_write_only_bots_group_id),
                    (can_delete_any_message_group_id),
                    (can_delete_own_message_group_id),
                    (can_invite_users_group_id),
                    (can_manage_all_groups_id),
                    (can_manage_billing_group_id),
                    (can_mention_many_users_group_id),
                    (can_move_messages_between_channels_group_id),
                    (can_move_messages_between_topics_group_id),
                    (can_resolve_topics_group_id),
                    (can_set_delete_message_policy_group_id),
                    (can_set_topics_policy_group_id),
                    (can_summarize_topics_group_id),
                    (direct_message_initiator_group_id),
                    (direct_message_permission_group_id),
                    (workplace_users_group_id)
            ) AS realm_cols(value)
            WHERE id = %s
        )

        SELECT ug.id
        FROM zerver_usergroup ug
        LEFT JOIN zerver_namedusergroup ng
          ON ng.usergroup_ptr_id = ug.id
        WHERE ug.realm_id = %s
          AND ng.usergroup_ptr_id IS NULL

        EXCEPT

        SELECT value FROM user_group_ids
    """

    BATCH_SIZE = 1000

    with connection.cursor() as cursor:
        for realm_id in Realm.objects.values_list("id", flat=True).iterator():
            cursor.execute(query, [realm_id, realm_id, realm_id, realm_id])
            unused_ids = [row[0] for row in cursor.fetchall()]

            while unused_ids:
                unused_id_str = SQL(",").join(map(Literal, unused_ids[:BATCH_SIZE]))
                unused_ids = unused_ids[BATCH_SIZE:]
                with transaction.atomic():
                    cursor.execute(
                        SQL(
                            "DELETE FROM zerver_usergroupmembership WHERE user_group_id IN ({unused_ids})"
                        ).format(unused_ids=unused_id_str)
                    )
                    cursor.execute(
                        SQL(
                            "DELETE FROM zerver_groupgroupmembership WHERE supergroup_id IN ({unused_ids})"
                        ).format(unused_ids=unused_id_str)
                    )
                    cursor.execute(
                        SQL("DELETE FROM zerver_usergroup WHERE id IN ({unused_ids})").format(
                            unused_ids=unused_id_str
                        )
                    )


class Migration(migrations.Migration):
    atomic = False

    dependencies = [
        ("zerver", "0809_remove_message_nulls_last_topic_indexes"),
    ]

    operations = [
        migrations.RunPython(
            delete_unused_anonymous_groups,
            reverse_code=migrations.RunPython.noop,
            elidable=True,
        ),
    ]