email digests: Query streams for messages up front.

This should save us many hops to the database when
we process users in bulk.

zerver/lib/digest.py

@@ -20,6 +20,7 @@ from zerver.models import (
     Realm,
     RealmAuditLog,
     Recipient,
+    Stream,
     Subscription,
     UserActivityInterval,
     UserProfile,
@@ -58,11 +59,12 @@ class DigestTopic:
     def diversity(self) -> int:
         return len(self.human_senders)
 
-    def teaser_data(self, user_profile: UserProfile) -> Dict[str, Any]:
+    def teaser_data(self, user: UserProfile, stream_map: Dict[int, Stream]) -> Dict[str, Any]:
         teaser_count = self.num_human_messages - len(self.sample_messages)
         first_few_messages = build_message_list(
-            user_profile,
-            self.sample_messages,
+            user=user,
+            messages=self.sample_messages,
+            stream_map=stream_map,
         )
         return {
             "participants": sorted(self.human_senders),
@@ -233,6 +235,14 @@ def get_user_stream_map(user_ids: List[int]) -> Dict[int, Set[int]]:
 
     return dct
 
+def get_slim_stream_map(stream_ids: Set[int]) -> Dict[int, Stream]:
+    # This can be passed to build_message_list.
+    streams = Stream.objects.filter(
+        id__in=stream_ids,
+    ).only('id', 'name')
+
+    return {stream.id: stream for stream in streams}
+
 def bulk_get_digest_context(users: List[UserProfile], cutoff: float) -> Dict[int, Dict[str, Any]]:
     # Convert from epoch seconds to a datetime object.
     cutoff_date = datetime.datetime.fromtimestamp(int(cutoff), tz=datetime.timezone.utc)
@@ -257,6 +267,8 @@ def bulk_get_digest_context(users: List[UserProfile], cutoff: float) -> Dict[int
     # for each user, we filter to just the streams they care about.
     recent_topics = get_recent_topics(sorted(list(all_stream_ids)), cutoff_date)
 
+    stream_map = get_slim_stream_map(all_stream_ids)
+
     for user in users:
         stream_ids = user_stream_map[user.id]
 
@@ -270,7 +282,7 @@ def bulk_get_digest_context(users: List[UserProfile], cutoff: float) -> Dict[int
 
         # Get context data for hot conversations.
         context["hot_conversations"] = [
-            hot_topic.teaser_data(user)
+            hot_topic.teaser_data(user, stream_map)
             for hot_topic in hot_topics
         ]
 

zerver/lib/email_notifications.py

@@ -166,7 +166,11 @@ def fix_spoilers_in_text(content: str, language: str) -> str:
             output.append(line)
     return '\n'.join(output)
 
-def build_message_list(user: UserProfile, messages: List[Message]) -> List[Dict[str, Any]]:
+def build_message_list(
+    user: UserProfile,
+    messages: List[Message],
+    stream_map: Dict[int, Stream],  # only needs id, name
+) -> List[Dict[str, Any]]:
     """
     Builds the message list object for the missed message email template.
     The messages are collapsed into per-recipient and per-sender blocks, like
@@ -239,7 +243,12 @@ def build_message_list(user: UserProfile, messages: List[Message]) -> List[Dict[
             header = "You and {}".format(", ".join(other_recipients))
             header_html = f"<a style='color: #ffffff;' href='{narrow_link}'>{header}</a>"
         else:
-            stream = Stream.objects.only('id', 'name').get(id=message.recipient.type_id)
+            stream_id = message.recipient.type_id
+            stream = stream_map.get(stream_id, None)
+            if stream is None:
+                # Some of our callers don't populate stream_map, so
+                # we just populate the stream from the database.
+                stream = Stream.objects.only('id', 'name').get(id=stream_id)
             narrow_link = get_narrow_url(user, message, stream=stream)
             header = f"{stream.name} > {message.topic_name()}"
             stream_link = stream_narrow_url(user.realm, stream)
@@ -456,7 +465,11 @@ def do_send_missedmessage_events_reply_in_zulip(user_profile: UserProfile,
         )
     else:
         context.update(
-            messages=build_message_list(user_profile, [m['message'] for m in missed_messages]),
+            messages=build_message_list(
+                user=user_profile,
+                messages=[m['message'] for m in missed_messages],
+                stream_map={},
+            ),
             sender_str=", ".join(sender.full_name for sender in senders),
             realm_str=user_profile.realm.name,
             show_message_content=True,

zerver/tests/test_digest.py

@@ -68,7 +68,7 @@ class TestDigestEmailMessages(ZulipTestCase):
         with queries_captured() as queries:
             handle_digest_email(othello.id, cutoff)
 
-        self.assert_length(queries, 10)
+        self.assert_length(queries, 9)
 
         self.assertEqual(mock_send_future_email.call_count, 1)
         kwargs = mock_send_future_email.call_args[1]
@@ -124,7 +124,7 @@ class TestDigestEmailMessages(ZulipTestCase):
         with queries_captured() as queries:
             handle_digest_email(polonius.id, cutoff)
 
-        self.assert_length(queries, 10)
+        self.assert_length(queries, 9)
 
         self.assertEqual(mock_send_future_email.call_count, 1)
         kwargs = mock_send_future_email.call_args[1]
@@ -182,7 +182,7 @@ class TestDigestEmailMessages(ZulipTestCase):
                 with cache_tries_captured() as cache_tries:
                     bulk_handle_digest_email(digest_user_ids, cutoff)
 
-            self.assert_length(queries, 30)
+            self.assert_length(queries, 15)
             self.assert_length(cache_tries, 0)
 
         self.assertEqual(mock_send_future_email.call_count, len(digest_users))