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

zerver/models/recipients.py

@@ -0,0 +1,241 @@
+import hashlib
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from django.db import connection, models, transaction
+from django.db.models import QuerySet
+from psycopg2.sql import SQL, Identifier
+from typing_extensions import override
+
+from zerver.lib.display_recipient import get_display_recipient
+
+if TYPE_CHECKING:
+    from zerver.models import Subscription
+
+
+class Recipient(models.Model):
+    """Represents an audience that can potentially receive messages in Zulip.
+
+    This table essentially functions as a generic foreign key that
+    allows Message.recipient_id to be a simple ForeignKey representing
+    the audience for a message, while supporting the different types
+    of audiences Zulip supports for a message.
+
+    Recipient has just two attributes: The enum type, and a type_id,
+    which is the ID of the Stream/DirectMessageGroup object containing
+    all the metadata for the audience. There are 2 recipient types:
+
+    1. Stream message: The type_id is the ID of the associated Stream.
+    2. Direct message: The type_id is the ID of the associated
+       DirectMessageGroup object, which encodes the set of users in
+       the conversation. The set of users is usually retrieved via the
+       Subscription table. See the DirectMessageGroup model for
+       details.
+
+    See also the Subscription model, which stores which UserProfile
+    objects are subscribed to which Recipient objects.
+    """
+
+    id = models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")
+    type_id = models.IntegerField(db_index=True)
+    type = models.PositiveSmallIntegerField(db_index=True)
+    # Valid types are {stream, direct_message_group}
+
+    # The type for stream messages.
+    STREAM = 2
+    # The type for direct messages (both 1:1 and group).
+    DIRECT_MESSAGE_GROUP = 3
+
+    class Meta:
+        unique_together = ("type", "type_id")
+
+    @override
+    def __str__(self) -> str:
+        return f"{self.label()} ({self.type_id}, {self.type})"
+
+    def label(self) -> str:
+        from zerver.models import Stream
+
+        if self.type == Recipient.STREAM:
+            return Stream.objects.get(id=self.type_id).name
+        else:
+            return str(get_display_recipient(self))
+
+
+def get_direct_message_group_user_ids(recipient: Recipient) -> QuerySet["Subscription", int]:
+    from zerver.models import Subscription
+
+    assert recipient.type == Recipient.DIRECT_MESSAGE_GROUP
+
+    return (
+        Subscription.objects.filter(
+            recipient=recipient,
+        )
+        .order_by("user_profile_id")
+        .values_list("user_profile_id", flat=True)
+    )
+
+
+def bulk_get_direct_message_group_user_ids(recipient_ids: list[int]) -> dict[int, set[int]]:
+    """
+    Takes a list of direct_message_group type recipient_ids, returns
+    a dictmapping recipient id to list of user ids in the direct
+    message group.
+
+    We rely on our caller to pass us recipient_ids that correspond
+    to direct_message_group, but technically this function is valid
+    for any typeof subscription.
+    """
+    from zerver.models import Subscription
+
+    if not recipient_ids:
+        return {}
+
+    subscriptions = Subscription.objects.filter(
+        recipient_id__in=recipient_ids,
+    ).only("user_profile_id", "recipient_id")
+
+    result_dict: dict[int, set[int]] = defaultdict(set)
+    for subscription in subscriptions:
+        result_dict[subscription.recipient_id].add(subscription.user_profile_id)
+
+    return result_dict
+
+
+class DirectMessageGroup(models.Model):
+    """
+    Represents a group of individuals who may have a
+    group direct message conversation together.
+
+    The membership of the DirectMessageGroup is stored in the Subscription
+    table just like with Streams - for each user in the DirectMessageGroup,
+    there is a Subscription object tied to the UserProfile and the
+    DirectMessageGroup's recipient object.
+
+    A hash of the list of user IDs is stored in the huddle_hash field
+    below, to support efficiently mapping from a set of users to the
+    corresponding DirectMessageGroup object.
+    """
+
+    # We omit db_index and unique, as that also results in a useless
+    # varchar_pattern_ops index; see Meta, below.
+    huddle_hash = models.CharField(max_length=40)
+    # Foreign key to the Recipient object for this DirectMessageGroup.
+    recipient = models.OneToOneField(Recipient, null=True, on_delete=models.SET_NULL)
+
+    group_size = models.IntegerField()
+
+    # TODO: The model still uses the old "zerver_huddle" database table.
+    # As a part of the migration of "Huddle" to "DirectMessageGroup"
+    # it needs to be renamed to "zerver_directmessagegroup".
+    class Meta:
+        db_table = "zerver_huddle"
+        constraints = [
+            models.UniqueConstraint(fields=["huddle_hash"], name="zerver_huddle_huddle_hash_key")
+        ]
+
+
+def get_direct_message_group_hash(id_list: list[int]) -> str:
+    id_list = sorted(set(id_list))
+    hash_key = ",".join(str(x) for x in id_list)
+    return hashlib.sha1(hash_key.encode()).hexdigest()
+
+
+def get_or_create_direct_message_group(id_list: list[int]) -> DirectMessageGroup:
+    """
+    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.
+    """
+    from zerver.models import Subscription, UserProfile
+
+    assert len(id_list) == len(set(id_list))
+    direct_message_group_hash = get_direct_message_group_hash(id_list)
+    with transaction.atomic(savepoint=False):
+        (direct_message_group, created) = DirectMessageGroup.objects.get_or_create(
+            huddle_hash=direct_message_group_hash,
+            group_size=len(id_list),
+        )
+
+        if not created:
+            return direct_message_group
+
+        recipient = Recipient.objects.create(
+            type_id=direct_message_group.id, type=Recipient.DIRECT_MESSAGE_GROUP
+        )
+
+        # The equivalent effect of doing
+        # direct_message_group.save(update_fields=["recipient"]) is handled
+        # in the following raw query via UPDATE to save us a query.
+        direct_message_group.recipient = recipient
+
+        # In a single raw query, we do 2 required things after creating
+        # a new DirectMessageGroup object:
+        # 1. Update the recipient field to reflect the above assignment in db.
+        # 2. Bulk create subscriptions for the participants.
+        query = SQL(
+            """
+            UPDATE {dmg_table}
+            SET recipient_id = %(recipient_id)s
+            WHERE {dmg_table}.id = %(dmg_id)s;
+
+            INSERT INTO {subscription_table}
+                (
+                    user_profile_id,
+                    is_user_active,
+                    recipient_id,
+                    active,
+                    is_muted,
+                    color,
+                    pin_to_top
+                )
+            SELECT
+                {user_profile_table}.id,
+                {user_profile_table}.is_active,
+                %(recipient_id)s,
+                %(active)s,
+                %(is_muted)s,
+                %(color)s,
+                %(pin_to_top)s
+            FROM {user_profile_table}
+            WHERE {user_profile_table}.id = ANY(%(participant_ids)s);
+        """
+        ).format(
+            dmg_table=Identifier(DirectMessageGroup._meta.db_table),
+            subscription_table=Identifier(Subscription._meta.db_table),
+            user_profile_table=Identifier(UserProfile._meta.db_table),
+        )
+        with connection.cursor() as cursor:
+            cursor.execute(
+                query,
+                {
+                    "recipient_id": recipient.id,
+                    "dmg_id": direct_message_group.id,
+                    "active": Subscription._meta.get_field("active").get_default(),
+                    "is_muted": Subscription._meta.get_field("is_muted").get_default(),
+                    "color": Subscription._meta.get_field("color").get_default(),
+                    "pin_to_top": Subscription._meta.get_field("pin_to_top").get_default(),
+                    "participant_ids": id_list,
+                },
+            )
+
+        return direct_message_group
+
+
+def get_direct_message_group(id_list: list[int]) -> DirectMessageGroup | None:
+    """
+    Takes a list of user IDs and returns the DirectMessageGroup
+    object for the group consisting of these users if exists. If
+    the DirectMessageGroup object does not yet exist, it will
+    return None.
+    """
+    assert len(id_list) == len(set(id_list))
+    try:
+        direct_message_group_hash = get_direct_message_group_hash(id_list)
+        return DirectMessageGroup.objects.get(
+            huddle_hash=direct_message_group_hash,
+            group_size=len(id_list),
+        )
+    except DirectMessageGroup.DoesNotExist:
+        return None