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

zerver/migrations/0798_remove_userprofile_recipient_and_personal_recipients.py

@@ -0,0 +1,314 @@
+import hashlib
+
+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
+from psycopg2.sql import SQL, Identifier
+
+PERSONAL_RECIPIENT_TYPE = 1
+DIRECT_MESSAGE_GROUP_RECIPIENT_TYPE = 3
+BATCH_SIZE = 5000
+
+# Tables whose rows can point at a personal recipient, with the column
+# identifying the message's sender.  Copied, like the backfill below,
+# from migration 0790.
+TABLES = [
+    ("zerver_message", "sender_id"),
+    ("zerver_archivedmessage", "sender_id"),
+    ("zerver_scheduledmessage", "sender_id"),
+    ("zerver_draft", "user_profile_id"),
+]
+
+
+def get_direct_message_group_hash(id_list: list[int]) -> str:
+    """
+    Takes a list of user IDs and returns a hash for the group consisting of
+    these users. This is used to create a unique identifier for direct message groups.
+
+    This function is copied from zerver/model/recipients.py to avoid import issues.
+    """
+    sorted_ids = sorted(id_list)
+    hash_key = ",".join(str(x) for x in sorted_ids)
+    return hashlib.sha1(hash_key.encode()).hexdigest()
+
+
+def get_or_create_direct_message_group(apps: StateApps, id_list: list[int]) -> int:
+    """
+    Takes a list of user IDs and returns the DirectMessageGroup object for
+    the group consisting of these users. If the DirectMessageGroup object
+    does not yet exist, it will be transparently created.
+    """
+
+    direct_message_group_hash = get_direct_message_group_hash(id_list)
+    with transaction.atomic(savepoint=False), connection.cursor() as cursor:
+        # The "do update" is mostly superfluous here, but required so
+        # that we get the id back if the row already existed.
+        cursor.execute(
+            """
+            INSERT INTO zerver_huddle (huddle_hash, group_size)
+            VALUES (%s, %s)
+            ON CONFLICT (huddle_hash) DO UPDATE SET group_size = %s
+            RETURNING id, recipient_id
+            """,
+            [direct_message_group_hash, len(id_list), len(id_list)],
+        )
+        direct_message_group_id, recipient_id = cursor.fetchone()
+        if recipient_id is not None:
+            return recipient_id
+
+        cursor.execute(
+            """
+            INSERT INTO zerver_recipient (type_id, type)
+            VALUES (%s, %s)
+            RETURNING id
+            """,
+            [direct_message_group_id, DIRECT_MESSAGE_GROUP_RECIPIENT_TYPE],
+        )
+        recipient_id = cursor.fetchone()[0]
+
+        cursor.execute(
+            """
+            UPDATE zerver_huddle SET recipient_id = %s WHERE id = %s
+            """,
+            [recipient_id, direct_message_group_id],
+        )
+
+        cursor.execute(
+            """
+            INSERT INTO zerver_subscription
+                (recipient_id, user_profile_id, is_user_active,
+                 active, is_muted, color, pin_to_top,
+                 desktop_notifications, audible_notifications,
+                 push_notifications, email_notifications,
+                 wildcard_mentions_notify)
+            SELECT %s, id, is_active,
+                   true, false, '#c2c2c2', false,
+                   NULL, NULL, NULL, NULL, NULL
+            FROM zerver_userprofile
+            WHERE id IN %s
+            """,
+            [recipient_id, tuple(id_list)],
+        )
+        return recipient_id
+
+
+def backfill_update_recipient_type_of_personal_messages_to_dm_group(
+    apps: StateApps, schema_editor: BaseDatabaseSchemaEditor
+) -> None:
+    """Convert messages still attached to personal recipients into direct
+    message groups.
+
+    Copied from migration 0790, which is meant to have already done
+    this. A deployment that ran an intermediate `main` commit can reach
+    0798 with such messages still present, which would make the personal
+    recipient DELETE below fail; rerunning the conversion here repairs
+    that.  It only selects personal recipients that still have
+    referencing rows, so it is a no-op once 0790 has run cleanly.
+    """
+    print()
+
+    with connection.cursor() as cursor:
+        # We use a temp table name distinct from the personal_recipients_to_process
+        # table in migration 0790, which does the same work. An install or upgrade
+        # applies both migrations in a single `migrate` run on one
+        # connection, so reusing the name here would collide.
+        #
+        # Only include personal recipients that have rows in at least
+        # one of the four tables we need to convert.
+        cursor.execute(
+            "CREATE TEMP TABLE personal_recipients_to_process_0798 AS "
+            "SELECT r.id, type_id FROM zerver_recipient r "
+            "JOIN zerver_userprofile u ON u.id = r.type_id "
+            "WHERE r.type = %s AND ("
+            "  EXISTS (SELECT 1 FROM zerver_message WHERE recipient_id = r.id)"
+            "  OR EXISTS (SELECT 1 FROM zerver_archivedmessage WHERE recipient_id = r.id)"
+            "  OR EXISTS (SELECT 1 FROM zerver_scheduledmessage WHERE recipient_id = r.id)"
+            "  OR EXISTS (SELECT 1 FROM zerver_draft WHERE recipient_id = r.id)"
+            ") ORDER BY u.realm_id, r.id",
+            [PERSONAL_RECIPIENT_TYPE],
+        )
+        cursor.execute(
+            "CREATE INDEX personal_recipients_to_process_0798_id"
+            " ON personal_recipients_to_process_0798 (id)"
+        )
+        cursor.execute("SELECT COUNT(*) FROM personal_recipients_to_process_0798")
+        (total,) = cursor.fetchone()
+        if total == 0:
+            return
+
+        print(f"Processing {total} personal recipients...")
+
+        processed = 0
+        while True:
+            cursor.execute(
+                "SELECT id, type_id FROM personal_recipients_to_process_0798 LIMIT %s",
+                # Use the batch size matching 0790:
+                [500],
+            )
+            batch = cursor.fetchall()
+            if not batch:
+                break
+
+            for personal_rec_id, receiving_user_id in batch:
+                # Collect distinct senders per table for this recipient.
+                sender_tables: dict[int, list[tuple[str, str]]] = {}
+                for table, col in TABLES:
+                    cursor.execute(
+                        SQL("SELECT DISTINCT {} FROM {} WHERE recipient_id = %s").format(
+                            Identifier(col), Identifier(table)
+                        ),
+                        [personal_rec_id],
+                    )
+                    for (sender_id,) in cursor.fetchall():
+                        sender_tables.setdefault(sender_id, []).append((table, col))
+
+                # For each sender, create the DMG and update rows atomically.
+                for sender_id, tables_to_update in sender_tables.items():
+                    if sender_id == receiving_user_id:
+                        id_list = [sender_id]
+                    else:
+                        id_list = [sender_id, receiving_user_id]
+
+                    with transaction.atomic():
+                        group_recipient_id = get_or_create_direct_message_group(apps, id_list)
+                        for table, col in tables_to_update:
+                            cursor.execute(
+                                SQL(
+                                    "UPDATE {} SET recipient_id = %s"
+                                    " WHERE {} = %s AND recipient_id = %s"
+                                ).format(Identifier(table), Identifier(col)),
+                                [group_recipient_id, sender_id, personal_rec_id],
+                            )
+
+            cursor.execute(
+                "DELETE FROM personal_recipients_to_process_0798 WHERE id = ANY(%s)",
+                [[row[0] for row in batch]],
+            )
+            processed += len(batch)
+            print(f"  Processed {processed}/{total} personal recipients.")
+
+        print("Migration completed.")
+
+
+def delete_orphan_personal_recipients(
+    apps: StateApps, schema_editor: BaseDatabaseSchemaEditor
+) -> None:
+    """Delete personal recipients whose owning UserProfile is gone.
+
+    A personal recipient's type_id is its owner's user id, but it is a
+    plain integer, not a foreign key, so a user removed outside the
+    normal codepath leaves its personal recipient -- and the messages
+    pointing at it -- behind. The backfill above skips these (its join
+    to zerver_userprofile finds no owner), so they would block the
+    recipient DELETE below.
+
+    Deleting the recipient cascades to its messages, archived messages,
+    scheduled messages, and subscriptions; drafts are left with a null
+    recipient, matching that on_delete behavior. This mirrors the orphan
+    cleanup that 0803 does for stream and direct-message-group
+    recipients.
+    """
+    Recipient = apps.get_model("zerver", "Recipient")
+    UserProfile = apps.get_model("zerver", "UserProfile")
+
+    orphan_ids = list(
+        Recipient.objects.filter(type=PERSONAL_RECIPIENT_TYPE)
+        .filter(~Exists(UserProfile.objects.filter(id=OuterRef("type_id"))))
+        .values_list("id", flat=True)
+    )
+    print()
+    print(f"Deleting {len(orphan_ids)} orphan personal recipients...")
+    for chunk_start in range(0, len(orphan_ids), BATCH_SIZE):
+        chunk = orphan_ids[chunk_start : chunk_start + BATCH_SIZE]
+        Recipient.objects.filter(id__in=chunk).delete()
+        print(f"  Deleted {chunk_start + len(chunk)}/{len(orphan_ids)} orphan personal recipients.")
+
+
+def delete_personal_recipient_data(
+    apps: StateApps, schema_editor: BaseDatabaseSchemaEditor
+) -> None:
+    """Delete Subscription rows and then Recipient rows for personal recipients.
+
+    Done in batches to avoid long-running transactions on large
+    deployments. Both loops are idempotent: re-running simply continues
+    deleting whatever rows remain.
+    """
+    with connection.cursor() as cursor:
+        while True:
+            cursor.execute(
+                SQL(
+                    """
+                    DELETE FROM zerver_subscription
+                    WHERE id IN (
+                        SELECT s.id FROM zerver_subscription s
+                        JOIN zerver_recipient r ON s.recipient_id = r.id
+                        WHERE r.type = %s
+                        LIMIT %s
+                    )
+                    """
+                ),
+                [PERSONAL_RECIPIENT_TYPE, BATCH_SIZE],
+            )
+            if cursor.rowcount == 0:
+                break
+
+        while True:
+            cursor.execute(
+                SQL(
+                    """
+                    DELETE FROM zerver_recipient
+                    WHERE id IN (
+                        SELECT id FROM zerver_recipient
+                        WHERE type = %s
+                        LIMIT %s
+                    )
+                    """
+                ),
+                [PERSONAL_RECIPIENT_TYPE, BATCH_SIZE],
+            )
+            if cursor.rowcount == 0:
+                break
+
+
+class Migration(migrations.Migration):
+    atomic = False
+
+    dependencies = [
+        ("zerver", "0797_userprofile_is_deleted"),
+    ]
+
+    operations = [
+        migrations.RunPython(
+            backfill_update_recipient_type_of_personal_messages_to_dm_group,
+            reverse_code=migrations.RunPython.noop,
+            elidable=True,
+        ),
+        migrations.RunPython(
+            delete_orphan_personal_recipients,
+            reverse_code=migrations.RunPython.noop,
+            elidable=True,
+        ),
+        # Drop the column with DROP COLUMN IF EXISTS, via
+        # SeparateDatabaseAndState, so that re-running the migration after an
+        # interrupted run does not fail on an already-dropped column.
+        migrations.SeparateDatabaseAndState(
+            database_operations=[
+                migrations.RunSQL(
+                    "ALTER TABLE zerver_userprofile DROP COLUMN IF EXISTS recipient_id",
+                    reverse_sql=migrations.RunSQL.noop,
+                ),
+            ],
+            state_operations=[
+                migrations.RemoveField(
+                    model_name="userprofile",
+                    name="recipient",
+                ),
+            ],
+        ),
+        migrations.RunPython(
+            delete_personal_recipient_data,
+            reverse_code=migrations.RunPython.noop,
+            elidable=True,
+        ),
+    ]