diff --git a/zerver/lib/display_recipient.py b/zerver/lib/display_recipient.py index c1f9792d5d..d9b5f2da56 100644 --- a/zerver/lib/display_recipient.py +++ b/zerver/lib/display_recipient.py @@ -147,7 +147,7 @@ def bulk_fetch_display_recipients( # Find all user ids whose UserProfiles we will need to fetch: user_ids_to_fetch: Set[int] = set() - huddle_user_ids: Dict[int, List[int]] = {} + huddle_user_ids: Dict[int, Set[int]] = {} huddle_user_ids = bulk_get_huddle_user_ids( [recipient for recipient in recipients if recipient.type == Recipient.HUDDLE] ) diff --git a/zerver/models.py b/zerver/models.py index e3202812c7..e33bb5a14a 100644 --- a/zerver/models.py +++ b/zerver/models.py @@ -3,6 +3,7 @@ import datetime import hashlib import secrets import time +from collections import defaultdict from contextlib import suppress from datetime import timedelta from email.headerregistry import Address @@ -2898,7 +2899,7 @@ def get_huddle_user_ids(recipient: Recipient) -> ValuesQuerySet["Subscription", ) -def bulk_get_huddle_user_ids(recipients: List[Recipient]) -> Dict[int, List[int]]: +def bulk_get_huddle_user_ids(recipients: List[Recipient]) -> Dict[int, Set[int]]: """ Takes a list of huddle-type recipients, returns a dict mapping recipient id to list of user ids in the huddle. @@ -2909,15 +2910,11 @@ def bulk_get_huddle_user_ids(recipients: List[Recipient]) -> Dict[int, List[int] subscriptions = Subscription.objects.filter( recipient__in=recipients, - ).order_by("user_profile_id") + ).only("user_profile_id", "recipient_id") - result_dict: Dict[int, List[int]] = {} - for recipient in recipients: - result_dict[recipient.id] = [ - subscription.user_profile_id - for subscription in subscriptions - if subscription.recipient_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 diff --git a/zerver/tests/test_messages.py b/zerver/tests/test_messages.py index 3fb3ece1f1..bcbcdf0ea3 100644 --- a/zerver/tests/test_messages.py +++ b/zerver/tests/test_messages.py @@ -90,9 +90,9 @@ class TestBulkGetHuddleUserIds(ZulipTestCase): messages = Message.objects.filter(id__in=message_ids).order_by("id") first_huddle_recipient = messages[0].recipient - first_huddle_user_ids = list(get_huddle_user_ids(first_huddle_recipient)) + first_huddle_user_ids = set(get_huddle_user_ids(first_huddle_recipient)) second_huddle_recipient = messages[1].recipient - second_huddle_user_ids = list(get_huddle_user_ids(second_huddle_recipient)) + second_huddle_user_ids = set(get_huddle_user_ids(second_huddle_recipient)) huddle_user_ids = bulk_get_huddle_user_ids( [first_huddle_recipient, second_huddle_recipient]