upstream-zulip/zerver/models/recipients.py

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