snapshot of https://github.com/zulip/zulip.git @ 48c69911c7dd4e03b20552eb3bd77ad8151c1357 (2026-09-15)

zerver/migrations/0419_backfill_message_realm.py

@@ -0,0 +1,114 @@
+from django.conf import settings
+from django.db import migrations, transaction
+from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+from django.db.migrations.state import StateApps
+from django.db.models import Exists, Max, OuterRef, Subquery
+
+
+def backfill_message_realm(apps: StateApps, schema_editor: BaseDatabaseSchemaEditor) -> None:
+    RECIPIENT_PERSONAL = 1
+    RECIPIENT_STREAM = 2
+    RECIPIENT_HUDDLE = 3
+
+    Message = apps.get_model("zerver", "Message")
+    ArchivedMessage = apps.get_model("zerver", "ArchivedMessage")
+    Recipient = apps.get_model("zerver", "Recipient")
+    Subscription = apps.get_model("zerver", "Subscription")
+    Stream = apps.get_model("zerver", "Stream")
+    UserProfile = apps.get_model("zerver", "UserProfile")
+    Huddle = apps.get_model("zerver", "Huddle")
+
+    print()
+
+    print("Deleting dangling Recipient objects and their messages, which are inaccessible.")
+    Recipient.objects.annotate(
+        has_object=Exists(UserProfile.objects.filter(id=OuterRef("type_id")))
+    ).filter(type=RECIPIENT_PERSONAL, has_object=False).delete()
+    Recipient.objects.annotate(
+        has_object=Exists(Stream.objects.filter(id=OuterRef("type_id")))
+    ).filter(type=RECIPIENT_STREAM, has_object=False).delete()
+    Recipient.objects.annotate(
+        has_object=Exists(Huddle.objects.filter(id=OuterRef("type_id")))
+    ).filter(type=RECIPIENT_HUDDLE, has_object=False).delete()
+
+    BATCH_SIZE = 10000
+    for message_model in [Message, ArchivedMessage]:
+        lower_bound = 1
+
+        max_id = message_model.objects.aggregate(Max("id"))["id__max"]
+        if max_id is None:
+            continue
+
+        while lower_bound <= max_id:
+            # Django's range() function is inclusive on both ends.
+            upper_bound = lower_bound + BATCH_SIZE - 1
+            print(f"Processing batch {lower_bound} to {upper_bound} for {message_model.__name__}")
+
+            with transaction.atomic():
+                message_model.objects.filter(
+                    id__range=(lower_bound, upper_bound),
+                    recipient__type=RECIPIENT_STREAM,
+                ).update(
+                    realm=Subquery(
+                        Recipient.objects.filter(pk=OuterRef("recipient")).values("stream__realm")
+                    )
+                )
+
+                # Private message to cross-realm bots are a special case, and the .realm
+                # of the message should be realm of the sender.
+                message_model.objects.filter(
+                    id__range=(lower_bound, upper_bound),
+                    recipient__type=RECIPIENT_PERSONAL,
+                    recipient__userprofile__delivery_email__in=settings.CROSS_REALM_BOT_EMAILS,
+                ).update(
+                    realm=Subquery(
+                        UserProfile.objects.filter(pk=OuterRef("sender")).values("realm")
+                    )
+                )
+
+                message_model.objects.filter(
+                    id__range=(lower_bound, upper_bound),
+                    recipient__type=RECIPIENT_PERSONAL,
+                ).exclude(
+                    recipient__userprofile__delivery_email__in=settings.CROSS_REALM_BOT_EMAILS
+                ).update(
+                    realm=Subquery(
+                        Recipient.objects.filter(pk=OuterRef("recipient")).values(
+                            "userprofile__realm"
+                        )
+                    )
+                )
+
+                # Huddles don't have a direct way of finding their
+                # realm, so we have to go through the Subscription
+                # table. For huddles including a cross-realm bot, all
+                # of the other users will be in the same realm, so
+                # just find any of those users to get the message's realm.
+                message_model.objects.filter(
+                    id__range=(lower_bound, upper_bound),
+                    recipient__type=RECIPIENT_HUDDLE,
+                ).update(
+                    realm=Subquery(
+                        Subscription.objects.filter(recipient=OuterRef("recipient"))
+                        .exclude(user_profile__delivery_email__in=settings.CROSS_REALM_BOT_EMAILS)
+                        .values("user_profile__realm")[:1]
+                    )
+                )
+
+            lower_bound += BATCH_SIZE
+
+
+class Migration(migrations.Migration):
+    atomic = False
+
+    dependencies = [
+        ("zerver", "0418_archivedmessage_realm_message_realm"),
+    ]
+
+    operations = [
+        migrations.RunPython(
+            backfill_message_realm,
+            reverse_code=migrations.RunPython.noop,
+            elidable=True,
+        ),
+    ]