From 51f1dc257d1f2a76d9bbdb7c2a7908eea0a10dcf Mon Sep 17 00:00:00 2001 From: Anders Kaseorg Date: Thu, 14 Dec 2023 18:18:20 -0800 Subject: [PATCH] models: Extract zerver.models.recipients. Signed-off-by: Anders Kaseorg --- zerver/actions/message_send.py | 2 +- zerver/data_import/rocketchat.py | 2 +- zerver/lib/display_recipient.py | 3 +- zerver/lib/import_realm.py | 2 +- zerver/lib/recipient_users.py | 3 +- zerver/models/__init__.py | 154 +----------------- zerver/models/recipients.py | 167 ++++++++++++++++++++ zerver/tests/test_import_export.py | 2 +- zerver/tests/test_message_send.py | 2 +- zerver/tests/test_messages.py | 9 +- zerver/tests/test_typing.py | 3 +- zilencer/management/commands/populate_db.py | 2 +- 12 files changed, 185 insertions(+), 166 deletions(-) create mode 100644 zerver/models/recipients.py diff --git a/zerver/actions/message_send.py b/zerver/actions/message_send.py index 5c62a83c35..29ccf9458d 100644 --- a/zerver/actions/message_send.py +++ b/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 diff --git a/zerver/data_import/rocketchat.py b/zerver/data_import/rocketchat.py index d724ab1dda..8013ecdde2 100644 --- a/zerver/data_import/rocketchat.py +++ b/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. # diff --git a/zerver/lib/display_recipient.py b/zerver/lib/display_recipient.py index 2e398b45c6..b3488c4dbc 100644 --- a/zerver/lib/display_recipient.py +++ b/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 {} diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py index 6347214856..d496d7e8d6 100644 --- a/zerver/lib/import_realm.py +++ b/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 = [ diff --git a/zerver/lib/recipient_users.py b/zerver/lib/recipient_users.py index cd4764d886..84364350bd 100644 --- a/zerver/lib/recipient_users.py +++ b/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 diff --git a/zerver/models/__init__.py b/zerver/models/__init__.py index 29bc1adfde..ca675fe115 100644 --- a/zerver/models/__init__.py +++ b/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 diff --git a/zerver/models/recipients.py b/zerver/models/recipients.py new file mode 100644 index 0000000000..28ba2f1a29 --- /dev/null +++ b/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 diff --git a/zerver/tests/test_import_export.py b/zerver/tests/test_import_export.py index c0b6a00c3f..7d3596eb4b 100644 --- a/zerver/tests/test_import_export.py +++ b/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 diff --git a/zerver/tests/test_message_send.py b/zerver/tests/test_message_send.py index b8b268c0b9..fd1e604a18 100644 --- a/zerver/tests/test_message_send.py +++ b/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 diff --git a/zerver/tests/test_messages.py b/zerver/tests/test_messages.py index 2fb797158c..5b08732aaf 100644 --- a/zerver/tests/test_messages.py +++ b/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): diff --git a/zerver/tests/test_typing.py b/zerver/tests/test_typing.py index 701b57d22f..ca01c2bf67 100644 --- a/zerver/tests/test_typing.py +++ b/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): diff --git a/zilencer/management/commands/populate_db.py b/zilencer/management/commands/populate_db.py index 3cde92da31..eb3269b065 100644 --- a/zilencer/management/commands/populate_db.py +++ b/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