models: Extract zerver.models.recipients.

Signed-off-by: Anders Kaseorg <anders@zulip.com>

zerver/actions/message_send.py

@@ -103,12 +103,12 @@ from zerver.models import (
     UserProfile,
     UserTopic,
     get_client,
-    get_huddle_user_ids,
     get_stream,
     get_stream_by_id_in_realm,
     query_for_ids,
 )
 from zerver.models.groups import SystemGroups
+from zerver.models.recipients import get_huddle_user_ids
 from zerver.models.users import get_system_bot, get_user_by_delivery_email, is_cross_realm_bot_email
 from zerver.tornado.django_api import send_event
 

zerver/data_import/rocketchat.py

@@ -886,7 +886,7 @@ def map_receiver_id_to_recipient_id(
             user_id_to_recipient_id[recipient["type_id"]] = recipient["id"]
 
 
-# This is inspired by get_huddle_hash from zerver/models/__init__.py. It
+# This is inspired by get_huddle_hash from zerver/models/recipients.py. It
 # expects strings identifying Rocket.Chat users, like
 # `LdBZ7kPxtKESyHPEe`, not integer IDs.
 #

zerver/lib/display_recipient.py

@@ -134,7 +134,8 @@ def bulk_fetch_user_display_recipients(
     Returns dict mapping recipient_id to corresponding display_recipient
     """
 
-    from zerver.models import Recipient, bulk_get_huddle_user_ids
+    from zerver.models import Recipient
+    from zerver.models.recipients import bulk_get_huddle_user_ids
 
     if len(recipient_tuples) == 0:
         return {}

zerver/lib/import_realm.py

@@ -76,10 +76,10 @@ from zerver.models import (
     UserProfile,
     UserStatus,
     UserTopic,
-    get_huddle_hash,
 )
 from zerver.models.groups import SystemGroups
 from zerver.models.realms import get_realm
+from zerver.models.recipients import get_huddle_hash
 from zerver.models.users import get_system_bot, get_user_profile_by_id
 
 realm_tables = [

zerver/lib/recipient_users.py

@@ -3,7 +3,8 @@ from typing import Dict, Optional, Sequence
 from django.core.exceptions import ValidationError
 from django.utils.translation import gettext as _
 
-from zerver.models import Recipient, UserProfile, get_or_create_huddle
+from zerver.models import Recipient, UserProfile
+from zerver.models.recipients import get_or_create_huddle
 from zerver.models.users import is_cross_realm_bot_email
 
 

zerver/models/__init__.py

@@ -4,7 +4,6 @@
 import hashlib
 import secrets
 import time
-from collections import defaultdict
 from datetime import datetime, timedelta, timezone
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, TypedDict, TypeVar, Union
 
@@ -18,7 +17,7 @@ from django.contrib.postgres.indexes import GinIndex
 from django.contrib.postgres.search import SearchVectorField
 from django.core.exceptions import ValidationError
 from django.core.serializers.json import DjangoJSONEncoder
-from django.db import models, transaction
+from django.db import models
 from django.db.backends.base.base import BaseDatabaseWrapper
 from django.db.models import CASCADE, Exists, F, OuterRef, Q, QuerySet
 from django.db.models.functions import Lower, Upper
@@ -43,7 +42,7 @@ from zerver.lib.cache import (
     realm_alert_words_automaton_cache_key,
     realm_alert_words_cache_key,
 )
-from zerver.lib.display_recipient import get_display_recipient, get_recipient_ids
+from zerver.lib.display_recipient import get_recipient_ids
 from zerver.lib.exceptions import RateLimitedError
 from zerver.lib.timestamp import datetime_to_timestamp
 from zerver.lib.types import (
@@ -79,6 +78,8 @@ from zerver.models.realm_playgrounds import RealmPlayground as RealmPlayground
 from zerver.models.realms import Realm as Realm
 from zerver.models.realms import RealmAuthenticationMethod as RealmAuthenticationMethod
 from zerver.models.realms import RealmDomain as RealmDomain
+from zerver.models.recipients import Huddle as Huddle
+from zerver.models.recipients import Recipient as Recipient
 from zerver.models.users import RealmUserDefault as RealmUserDefault
 from zerver.models.users import UserBaseSettings as UserBaseSettings
 from zerver.models.users import UserProfile as UserProfile
@@ -138,64 +139,6 @@ def query_for_ids(
     return query
 
 
-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 UserProfile/Stream/Huddle object containing
-    all the metadata for the audience. There are 3 recipient types:
-
-    1. 1:1 direct message: The type_id is the ID of the UserProfile
-       who will receive any message to this Recipient. The sender
-       of such a message is represented separately.
-    2. Stream message: The type_id is the ID of the associated Stream.
-    3. Group direct message: In Zulip, group direct messages are
-       represented by Huddle objects, which encode the set of users
-       in the conversation. The type_id is the ID of the associated Huddle
-       object; the set of users is usually retrieved via the Subscription
-       table. See the Huddle model for details.
-
-    See also the Subscription model, which stores which UserProfile
-    objects are subscribed to which Recipient objects.
-    """
-
-    type_id = models.IntegerField(db_index=True)
-    type = models.PositiveSmallIntegerField(db_index=True)
-    # Valid types are {personal, stream, huddle}
-
-    # The type for 1:1 direct messages.
-    PERSONAL = 1
-    # The type for stream messages.
-    STREAM = 2
-    # The type group direct messages.
-    HUDDLE = 3
-
-    class Meta:
-        unique_together = ("type", "type_id")
-
-    # N.B. If we used Django's choice=... we would get this for free (kinda)
-    _type_names = {PERSONAL: "personal", STREAM: "stream", HUDDLE: "huddle"}
-
-    @override
-    def __str__(self) -> str:
-        return f"{self.label()} ({self.type_id}, {self.type})"
-
-    def label(self) -> str:
-        if self.type == Recipient.STREAM:
-            return Stream.objects.get(id=self.type_id).name
-        else:
-            return str(get_display_recipient(self))
-
-    def type_name(self) -> str:
-        # Raises KeyError if invalid
-        return self._type_names[self.type]
-
-
 class PreregistrationRealm(models.Model):
     """Data on a partially created realm entered by a user who has
     completed the "new organization" form. Used to transfer the user's
@@ -794,41 +737,6 @@ def bulk_get_streams(realm: Realm, stream_names: Set[str]) -> Dict[str, Any]:
     return {stream.name.lower(): stream for stream in streams}
 
 
-def get_huddle_user_ids(recipient: Recipient) -> ValuesQuerySet["Subscription", int]:
-    assert recipient.type == Recipient.HUDDLE
-
-    return (
-        Subscription.objects.filter(
-            recipient=recipient,
-        )
-        .order_by("user_profile_id")
-        .values_list("user_profile_id", flat=True)
-    )
-
-
-def bulk_get_huddle_user_ids(recipient_ids: List[int]) -> Dict[int, Set[int]]:
-    """
-    Takes a list of huddle-type recipient_ids, returns a dict
-    mapping recipient id to list of user ids in the huddle.
-
-    We rely on our caller to pass us recipient_ids that correspond
-    to huddles, but technically this function is valid for any type
-    of 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 AbstractMessage(models.Model):
     sender = models.ForeignKey(UserProfile, on_delete=CASCADE)
 
@@ -1802,60 +1710,6 @@ class Subscription(models.Model):
     ]
 
 
-class Huddle(models.Model):
-    """
-    Represents a group of individuals who may have a
-    group direct message conversation together.
-
-    The membership of the Huddle is stored in the Subscription table just like with
-    Streams - for each user in the Huddle, there is a Subscription object
-    tied to the UserProfile and the Huddle'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 Huddle object.
-    """
-
-    # TODO: We should consider whether using
-    # CommaSeparatedIntegerField would be better.
-    huddle_hash = models.CharField(max_length=40, db_index=True, unique=True)
-    # Foreign key to the Recipient object for this Huddle.
-    recipient = models.ForeignKey(Recipient, null=True, on_delete=models.SET_NULL)
-
-
-def get_huddle_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_huddle(id_list: List[int]) -> Huddle:
-    """
-    Takes a list of user IDs and returns the Huddle object for the
-    group consisting of these users. If the Huddle object does not
-    yet exist, it will be transparently created.
-    """
-    huddle_hash = get_huddle_hash(id_list)
-    with transaction.atomic():
-        (huddle, created) = Huddle.objects.get_or_create(huddle_hash=huddle_hash)
-        if created:
-            recipient = Recipient.objects.create(type_id=huddle.id, type=Recipient.HUDDLE)
-            huddle.recipient = recipient
-            huddle.save(update_fields=["recipient"])
-            subs_to_create = [
-                Subscription(
-                    recipient=recipient,
-                    user_profile_id=user_profile_id,
-                    is_user_active=is_active,
-                )
-                for user_profile_id, is_active in UserProfile.objects.filter(id__in=id_list)
-                .distinct("id")
-                .values_list("id", "is_active")
-            ]
-            Subscription.objects.bulk_create(subs_to_create)
-        return huddle
-
-
 class UserActivity(models.Model):
     """Data table recording the last time each user hit Zulip endpoints
     via which Clients; unlike UserPresence, these data are not exposed

zerver/models/recipients.py

@@ -0,0 +1,167 @@
+import hashlib
+from collections import defaultdict
+from typing import TYPE_CHECKING, Dict, List, Set
+
+from django.db import models, transaction
+from django_stubs_ext import ValuesQuerySet
+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 UserProfile/Stream/Huddle object containing
+    all the metadata for the audience. There are 3 recipient types:
+
+    1. 1:1 direct message: The type_id is the ID of the UserProfile
+       who will receive any message to this Recipient. The sender
+       of such a message is represented separately.
+    2. Stream message: The type_id is the ID of the associated Stream.
+    3. Group direct message: In Zulip, group direct messages are
+       represented by Huddle objects, which encode the set of users
+       in the conversation. The type_id is the ID of the associated Huddle
+       object; the set of users is usually retrieved via the Subscription
+       table. See the Huddle model for details.
+
+    See also the Subscription model, which stores which UserProfile
+    objects are subscribed to which Recipient objects.
+    """
+
+    type_id = models.IntegerField(db_index=True)
+    type = models.PositiveSmallIntegerField(db_index=True)
+    # Valid types are {personal, stream, huddle}
+
+    # The type for 1:1 direct messages.
+    PERSONAL = 1
+    # The type for stream messages.
+    STREAM = 2
+    # The type group direct messages.
+    HUDDLE = 3
+
+    class Meta:
+        unique_together = ("type", "type_id")
+
+    # N.B. If we used Django's choice=... we would get this for free (kinda)
+    _type_names = {PERSONAL: "personal", STREAM: "stream", HUDDLE: "huddle"}
+
+    @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 type_name(self) -> str:
+        # Raises KeyError if invalid
+        return self._type_names[self.type]
+
+
+def get_huddle_user_ids(recipient: Recipient) -> ValuesQuerySet["Subscription", int]:
+    from zerver.models import Subscription
+
+    assert recipient.type == Recipient.HUDDLE
+
+    return (
+        Subscription.objects.filter(
+            recipient=recipient,
+        )
+        .order_by("user_profile_id")
+        .values_list("user_profile_id", flat=True)
+    )
+
+
+def bulk_get_huddle_user_ids(recipient_ids: List[int]) -> Dict[int, Set[int]]:
+    """
+    Takes a list of huddle-type recipient_ids, returns a dict
+    mapping recipient id to list of user ids in the huddle.
+
+    We rely on our caller to pass us recipient_ids that correspond
+    to huddles, but technically this function is valid for any type
+    of 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 Huddle(models.Model):
+    """
+    Represents a group of individuals who may have a
+    group direct message conversation together.
+
+    The membership of the Huddle is stored in the Subscription table just like with
+    Streams - for each user in the Huddle, there is a Subscription object
+    tied to the UserProfile and the Huddle'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 Huddle object.
+    """
+
+    # TODO: We should consider whether using
+    # CommaSeparatedIntegerField would be better.
+    huddle_hash = models.CharField(max_length=40, db_index=True, unique=True)
+    # Foreign key to the Recipient object for this Huddle.
+    recipient = models.ForeignKey(Recipient, null=True, on_delete=models.SET_NULL)
+
+
+def get_huddle_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_huddle(id_list: List[int]) -> Huddle:
+    """
+    Takes a list of user IDs and returns the Huddle object for the
+    group consisting of these users. If the Huddle object does not
+    yet exist, it will be transparently created.
+    """
+    from zerver.models import Subscription, UserProfile
+
+    huddle_hash = get_huddle_hash(id_list)
+    with transaction.atomic():
+        (huddle, created) = Huddle.objects.get_or_create(huddle_hash=huddle_hash)
+        if created:
+            recipient = Recipient.objects.create(type_id=huddle.id, type=Recipient.HUDDLE)
+            huddle.recipient = recipient
+            huddle.save(update_fields=["recipient"])
+            subs_to_create = [
+                Subscription(
+                    recipient=recipient,
+                    user_profile_id=user_profile_id,
+                    is_user_active=is_active,
+                )
+                for user_profile_id, is_active in UserProfile.objects.filter(id__in=id_list)
+                .distinct("id")
+                .values_list("id", "is_active")
+            ]
+            Subscription.objects.bulk_create(subs_to_create)
+        return huddle

zerver/tests/test_import_export.py

@@ -85,11 +85,11 @@ from zerver.models import (
     UserTopic,
     get_active_streams,
     get_client,
-    get_huddle_hash,
     get_stream,
 )
 from zerver.models.groups import SystemGroups
 from zerver.models.realms import get_realm
+from zerver.models.recipients import get_huddle_hash
 from zerver.models.users import get_system_bot, get_user_by_delivery_email
 
 

zerver/tests/test_message_send.py

@@ -56,12 +56,12 @@ from zerver.models import (
     UserGroup,
     UserMessage,
     UserProfile,
-    get_or_create_huddle,
     get_stream,
 )
 from zerver.models.constants import MAX_TOPIC_NAME_LENGTH
 from zerver.models.groups import SystemGroups
 from zerver.models.realms import get_realm
+from zerver.models.recipients import get_or_create_huddle
 from zerver.models.users import get_system_bot, get_user
 from zerver.views.message_send import InvalidMirrorInputError
 

zerver/tests/test_messages.py

@@ -5,13 +5,8 @@ from django.utils.timezone import now as timezone_now
 
 from zerver.actions.message_send import get_active_presence_idle_user_ids
 from zerver.lib.test_classes import ZulipTestCase
-from zerver.models import (
-    Message,
-    UserPresence,
-    UserProfile,
-    bulk_get_huddle_user_ids,
-    get_huddle_user_ids,
-)
+from zerver.models import Message, UserPresence, UserProfile
+from zerver.models.recipients import bulk_get_huddle_user_ids, get_huddle_user_ids
 
 
 class MissedMessageTest(ZulipTestCase):

zerver/tests/test_typing.py

@@ -1,7 +1,8 @@
 import orjson
 
 from zerver.lib.test_classes import ZulipTestCase
-from zerver.models import Huddle, get_huddle_hash
+from zerver.models import Huddle
+from zerver.models.recipients import get_huddle_hash
 
 
 class TypingValidateOperatorTest(ZulipTestCase):

zilencer/management/commands/populate_db.py

@@ -66,10 +66,10 @@ from zerver.models import (
     UserProfile,
     flush_alert_word,
     get_client,
-    get_or_create_huddle,
     get_stream,
 )
 from zerver.models.realms import get_realm
+from zerver.models.recipients import get_or_create_huddle
 from zerver.models.users import get_user, get_user_by_delivery_email, get_user_profile_by_id
 from zilencer.models import RemoteRealm, RemoteZulipServer
 from zilencer.views import update_remote_realm_data_for_server