From faa98317bc2ad025f2df2dd9c73c456b1a460fbe Mon Sep 17 00:00:00 2001
From: Prakhar Pratyush <prakhar@zulip.com>
Date: Mon, 25 Sep 2023 12:39:53 +0530
Subject: [PATCH] user_topics: Refactor build_topic_mute_checker to support
 each policy.

This prep commit renames the 'build_topic_mute_checker' function
to 'build_get_topic_visibility_policy' and updates it to support
all the visibility policies.

The function prefetches the visibility policies the user has
configured for various topics and prepares a dict named
'topic_to_visibility_policy' to be used later on.
---
 zerver/lib/message.py     |  7 ++++---
 zerver/lib/user_topics.py | 29 +++++++++++++++++++----------
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/zerver/lib/message.py b/zerver/lib/message.py
index 5cde4d16fd..bd1a3b8f8e 100644
--- a/zerver/lib/message.py
+++ b/zerver/lib/message.py
@@ -53,7 +53,7 @@ from zerver.lib.topic import DB_TOPIC_NAME, MESSAGE__TOPIC, TOPIC_LINKS, TOPIC_N
 from zerver.lib.types import DisplayRecipientT, EditHistoryEvent, UserDisplayRecipient
 from zerver.lib.url_preview.types import UrlEmbedData
 from zerver.lib.user_groups import is_user_in_group
-from zerver.lib.user_topics import build_topic_mute_checker, topic_has_visibility_policy
+from zerver.lib.user_topics import build_get_topic_visibility_policy, topic_has_visibility_policy
 from zerver.models import (
     MAX_TOPIC_NAME_LENGTH,
     Message,
@@ -1111,13 +1111,14 @@ def extract_unread_data_from_um_rows(
     muted_stream_ids = get_muted_stream_ids(user_profile)
     raw_unread_messages["muted_stream_ids"] = muted_stream_ids
 
-    topic_mute_checker = build_topic_mute_checker(user_profile)
+    get_topic_visibility_policy = build_get_topic_visibility_policy(user_profile)
 
     def is_row_muted(stream_id: int, recipient_id: int, topic: str) -> bool:
         if stream_id in muted_stream_ids:
             return True
 
-        if topic_mute_checker(recipient_id, topic):
+        visibility_policy = get_topic_visibility_policy(recipient_id, topic)
+        if visibility_policy == UserTopic.VisibilityPolicy.MUTED:
             return True
 
         # Messages sent by muted users are never unread, so we don't
diff --git a/zerver/lib/user_topics.py b/zerver/lib/user_topics.py
index e99f96e5da..554a0bb871 100644
--- a/zerver/lib/user_topics.py
+++ b/zerver/lib/user_topics.py
@@ -1,6 +1,7 @@
 import datetime
 import logging
-from typing import Callable, List, Optional, Tuple, TypedDict
+from collections import defaultdict
+from typing import Callable, Dict, List, Optional, Tuple, TypedDict
 
 from django.db import transaction
 from django.db.models import QuerySet
@@ -233,24 +234,32 @@ def exclude_topic_mutes(
     return [*conditions, condition]
 
 
-def build_topic_mute_checker(user_profile: UserProfile) -> Callable[[int, str], bool]:
-    rows = UserTopic.objects.filter(
-        user_profile=user_profile, visibility_policy=UserTopic.VisibilityPolicy.MUTED
-    ).values(
+def build_get_topic_visibility_policy(
+    user_profile: UserProfile,
+) -> Callable[[int, str], int]:
+    """Prefetch the visibility policies the user has configured for
+    various topics.
+
+    The prefetching helps to avoid the db queries later in the loop
+    to determine the user's visibility policy for a topic.
+    """
+    rows = UserTopic.objects.filter(user_profile=user_profile).values(
         "recipient_id",
         "topic_name",
+        "visibility_policy",
     )
 
-    tups = set()
+    topic_to_visiblity_policy: Dict[Tuple[int, str], int] = defaultdict(int)
     for row in rows:
         recipient_id = row["recipient_id"]
         topic_name = row["topic_name"]
-        tups.add((recipient_id, topic_name.lower()))
+        visibility_policy = row["visibility_policy"]
+        topic_to_visiblity_policy[(recipient_id, topic_name)] = visibility_policy
 
-    def is_muted(recipient_id: int, topic: str) -> bool:
-        return (recipient_id, topic.lower()) in tups
+    def get_topic_visibility_policy(recipient_id: int, topic: str) -> int:
+        return topic_to_visiblity_policy[(recipient_id, topic.lower())]
 
-    return is_muted
+    return get_topic_visibility_policy
 
 
 def get_users_with_user_topic_visibility_policy(