upstream-zulip/zerver/migrations/0718_fix_topics_for_direct_messages.py

# Generated by Django 5.1.8 on 2025-05-23 09:25

from django.db import migrations
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps
from django.db.models import Max


def fix_topic_for_direct_messages(apps: StateApps, schema_editor: BaseDatabaseSchemaEditor) -> None:
    Message = apps.get_model("zerver", "Message")
    ArchivedMessage = apps.get_model("zerver", "ArchivedMessage")
    ScheduledMessage = apps.get_model("zerver", "ScheduledMessage")

    RECIPIENT_STREAM = 2

    DM_TOPIC = "\x07"

    BATCH_SIZE = 10000
    for message_model in [Message, ArchivedMessage, ScheduledMessage]:
        lower_bound = 1

        max_id = message_model.objects.aggregate(Max("id"))["id__max"]
        if max_id is None:
            continue

        while lower_bound <= max_id:
            upper_bound = lower_bound + BATCH_SIZE - 1
            print(f"Processing batch {lower_bound} to {upper_bound} for {message_model.__name__}")

            message_model.objects.filter(
                id__range=(lower_bound, upper_bound),
            ).exclude(recipient__type=RECIPIENT_STREAM).update(subject=DM_TOPIC)

            lower_bound += BATCH_SIZE


class Migration(migrations.Migration):
    atomic = False

    dependencies = [
        ("zerver", "0717_stream_topics_policy"),
    ]

    operations = [
        migrations.RunPython(
            fix_topic_for_direct_messages,
            reverse_code=migrations.RunPython.noop,
            elidable=True,
        ),
        # We need to run an ANALYZE for the query planner statistics
        # to no longer think most messages have "" as the topic.
        migrations.RunSQL("ANALYZE zerver_message"),
    ]