settings: Do not pre-fetch DM permission group settings.

This commit updates code to not pre-fetch DM permission
group settings using select_related and instead just
fetch the required data from DB when checking permission.

This will increase one query but will help in pre-fetching
the settings for all users and for all type of messages.

Fixes part of #33677.

zerver/actions/message_send.py

@@ -76,7 +76,12 @@ from zerver.lib.thumbnail import get_user_upload_previews, rewrite_thumbnailed_i
 from zerver.lib.timestamp import timestamp_to_datetime
 from zerver.lib.topic import participants_for_topic
 from zerver.lib.url_preview.types import UrlEmbedData
-from zerver.lib.user_groups import is_any_user_in_group, is_user_in_group
+from zerver.lib.user_groups import (
+    check_any_user_has_permission_by_role,
+    check_user_has_permission_by_role,
+    is_any_user_in_group,
+    is_user_in_group,
+)
 from zerver.lib.user_message import UserMessageLite, bulk_insert_ums
 from zerver.lib.users import (
     check_can_access_user,
@@ -100,7 +105,7 @@ from zerver.models import (
     UserTopic,
 )
 from zerver.models.clients import get_client
-from zerver.models.groups import SystemGroups
+from zerver.models.groups import SystemGroups, get_realm_system_groups_name_dict
 from zerver.models.recipients import get_direct_message_group_user_ids
 from zerver.models.scheduled_jobs import NotificationTriggers
 from zerver.models.streams import (
@@ -1580,53 +1585,56 @@ def check_can_send_direct_message(
     if all(user_profile.is_bot or user_profile.id == sender.id for user_profile in recipient_users):
         return
 
-    direct_message_permission_group = realm.direct_message_permission_group
-
-    if (
-        not hasattr(direct_message_permission_group, "named_user_group")
-        or direct_message_permission_group.named_user_group.name != SystemGroups.EVERYONE
-    ):
-        user_ids = [recipient_user.id for recipient_user in recipient_users] + [sender.id]
-        if not is_any_user_in_group(direct_message_permission_group.id, user_ids):
+    system_groups_name_dict = get_realm_system_groups_name_dict(realm.id)
+    if realm.direct_message_permission_group_id in system_groups_name_dict:
+        users = [*recipient_users, sender]
+        if not check_any_user_has_permission_by_role(
+            users, realm.direct_message_permission_group_id, system_groups_name_dict
+        ):
             is_nobody_group = (
-                hasattr(direct_message_permission_group, "named_user_group")
-                and direct_message_permission_group.named_user_group.name == SystemGroups.NOBODY
+                system_groups_name_dict[realm.direct_message_permission_group_id]
+                == SystemGroups.NOBODY
             )
             raise DirectMessagePermissionError(is_nobody_group)
+    else:
+        user_ids = [recipient_user.id for recipient_user in recipient_users] + [sender.id]
+        if not is_any_user_in_group(realm.direct_message_permission_group_id, user_ids):
+            raise DirectMessagePermissionError(is_nobody_group=False)
 
-    direct_message_initiator_group = realm.direct_message_initiator_group
-    if (
-        not hasattr(direct_message_initiator_group, "named_user_group")
-        or direct_message_initiator_group.named_user_group.name != SystemGroups.EVERYONE
-    ):
-        if is_user_in_group(direct_message_initiator_group.id, sender):
+    if realm.direct_message_initiator_group_id in system_groups_name_dict:
+        if check_user_has_permission_by_role(
+            sender, realm.direct_message_initiator_group_id, system_groups_name_dict
+        ):
+            return
+    else:
+        if is_user_in_group(realm.direct_message_initiator_group_id, sender):
             return
 
-        # TODO: This check is inefficient; we should in the future be able to cache
-        # on the Huddle object whether the conversation already exists, likely in the
-        # form of a `first_message_id` field, and be able to save doing this check in the
-        # common case that this is not the first message in a conversation.
-        if recipient.type == Recipient.PERSONAL:
-            recipient_user_profile = recipient_users[0]
-            previous_messages_exist = (
-                Message.objects.filter(
-                    realm=realm,
-                    recipient__type=Recipient.PERSONAL,
-                )
-                .filter(
-                    Q(sender=sender, recipient=recipient)
-                    | Q(sender=recipient_user_profile, recipient_id=sender.recipient_id)
-                )
-                .exists()
-            )
-        else:
-            assert recipient.type == Recipient.DIRECT_MESSAGE_GROUP
-            previous_messages_exist = Message.objects.filter(
+    # TODO: This check is inefficient; we should in the future be able to cache
+    # on the Huddle object whether the conversation already exists, likely in the
+    # form of a `first_message_id` field, and be able to save doing this check in the
+    # common case that this is not the first message in a conversation.
+    if recipient.type == Recipient.PERSONAL:
+        recipient_user_profile = recipient_users[0]
+        previous_messages_exist = (
+            Message.objects.filter(
                 realm=realm,
-                recipient=recipient,
-            ).exists()
-        if not previous_messages_exist:
-            raise DirectMessageInitiationError
+                recipient__type=Recipient.PERSONAL,
+            )
+            .filter(
+                Q(sender=sender, recipient=recipient)
+                | Q(sender=recipient_user_profile, recipient_id=sender.recipient_id)
+            )
+            .exists()
+        )
+    else:
+        assert recipient.type == Recipient.DIRECT_MESSAGE_GROUP
+        previous_messages_exist = Message.objects.filter(
+            realm=realm,
+            recipient=recipient,
+        ).exists()
+    if not previous_messages_exist:
+        raise DirectMessageInitiationError
 
 
 def check_sender_can_access_recipients(

zerver/lib/cache.py

@@ -527,6 +527,10 @@ def active_non_guest_user_ids_cache_key(realm_id: int) -> str:
     return f"active_non_guest_user_ids:{realm_id}"
 
 
+def get_realm_system_groups_cache_key(realm_id: int) -> str:
+    return f"realm_system_groups:{realm_id}"
+
+
 bot_dict_fields: list[str] = [
     "api_key",
     "avatar_source",

zerver/lib/email_mirror.py

@@ -126,10 +126,6 @@ def get_usable_missed_message_address(address: str) -> MissedMessageEmailAddress
             # can send a direct message to a given recipient.
             "user_profile__realm__can_access_all_users_group",
             "user_profile__realm__can_access_all_users_group__named_user_group",
-            "user_profile__realm__direct_message_initiator_group",
-            "user_profile__realm__direct_message_initiator_group__named_user_group",
-            "user_profile__realm__direct_message_permission_group",
-            "user_profile__realm__direct_message_permission_group__named_user_group",
             "message",
             "message__sender",
             "message__recipient",

zerver/lib/user_groups.py

@@ -1209,3 +1209,43 @@ def get_group_setting_value_for_audit_log_data(
         return setting_value
 
     return asdict(setting_value)
+
+
+def check_user_has_permission_by_role(
+    user: UserProfile, setting_group_id: int, system_groups_name_dict: dict[int, str]
+) -> bool:
+    system_group_name = system_groups_name_dict[setting_group_id]
+
+    if system_group_name == SystemGroups.NOBODY:
+        return False
+
+    if system_group_name == SystemGroups.EVERYONE:
+        return True
+
+    if user.is_guest:
+        return False
+
+    if system_group_name == SystemGroups.MEMBERS:
+        return True
+
+    if system_group_name == SystemGroups.OWNERS:
+        return user.is_realm_owner
+
+    if system_group_name == SystemGroups.ADMINISTRATORS:
+        return user.is_realm_admin
+
+    if system_group_name == SystemGroups.MODERATORS:
+        return user.is_moderator
+
+    # Handle full members case.
+    return user.role != UserProfile.ROLE_MEMBER or not user.is_provisional_member
+
+
+def check_any_user_has_permission_by_role(
+    users: list[UserProfile], setting_group_id: int, system_groups_name_dict: dict[int, str]
+) -> bool:
+    for user in users:
+        if check_user_has_permission_by_role(user, setting_group_id, system_groups_name_dict):
+            return True
+
+    return False

zerver/models/groups.py

@@ -4,6 +4,7 @@ from django.utils.timezone import now as timezone_now
 from django.utils.translation import gettext_lazy
 from django_cte import CTEManager
 
+from zerver.lib.cache import cache_with_key, get_realm_system_groups_cache_key
 from zerver.lib.types import GroupPermissionSetting
 from zerver.models.users import UserProfile
 
@@ -186,3 +187,11 @@ class GroupGroupMembership(models.Model):
                 fields=["supergroup", "subgroup"], name="zerver_groupgroupmembership_uniq"
             )
         ]
+
+
+@cache_with_key(get_realm_system_groups_cache_key, timeout=3600 * 24 * 7)
+def get_realm_system_groups_name_dict(realm_id: int) -> dict[int, str]:
+    system_groups = NamedUserGroup.objects.filter(
+        realm_id=realm_id, is_system_group=True
+    ).values_list("id", "name")
+    return dict(system_groups)

zerver/models/users.py

@@ -942,10 +942,6 @@ def base_get_user_queryset() -> QuerySet[UserProfile]:
         "realm",
         "realm__can_access_all_users_group",
         "realm__can_access_all_users_group__named_user_group",
-        "realm__direct_message_initiator_group",
-        "realm__direct_message_initiator_group__named_user_group",
-        "realm__direct_message_permission_group",
-        "realm__direct_message_permission_group__named_user_group",
         "bot_owner",
     )
 

zerver/tests/test_email_mirror.py

@@ -1267,7 +1267,7 @@ class TestMissedMessageEmailMessages(ZulipTestCase):
         incoming_valid_message["To"] = mm_address
         incoming_valid_message["Reply-to"] = self.example_email("othello")
 
-        with self.assert_database_query_count(16):
+        with self.assert_database_query_count(17):
             process_message(incoming_valid_message)
 
         # confirm that Hamlet got the message
@@ -1312,7 +1312,7 @@ class TestMissedMessageEmailMessages(ZulipTestCase):
         incoming_valid_message["To"] = mm_address
         incoming_valid_message["Reply-to"] = self.example_email("cordelia")
 
-        with self.assert_database_query_count(21):
+        with self.assert_database_query_count(22):
             process_message(incoming_valid_message)
 
         # Confirm Iago received the message.

zerver/tests/test_example.py

@@ -363,7 +363,7 @@ class TestQueryCounts(ZulipTestCase):
         hamlet = self.example_user("hamlet")
         cordelia = self.example_user("cordelia")
 
-        with self.assert_database_query_count(15):
+        with self.assert_database_query_count(16):
             self.send_personal_message(
                 from_user=hamlet,
                 to_user=cordelia,

zerver/tests/test_message_send.py

@@ -2671,7 +2671,22 @@ class PersonalMessageSendTest(ZulipTestCase):
             user_group,
             acting_user=None,
         )
-        self.send_personal_message(user_profile, cordelia)
+        with self.assert_database_query_count(16):
+            self.send_personal_message(user_profile, cordelia)
+
+        # Test that query count decreases if setting is set to a system group.
+        members_group = NamedUserGroup.objects.get(
+            name=SystemGroups.MEMBERS, realm=realm, is_system_group=True
+        )
+        do_change_realm_permission_group_setting(
+            realm,
+            "direct_message_initiator_group",
+            members_group,
+            acting_user=None,
+        )
+        othello = self.example_user("othello")
+        with self.assert_database_query_count(15):
+            self.send_personal_message(user_profile, othello)
 
     def test_direct_message_permission_group_setting(self) -> None:
         """
@@ -2743,7 +2758,26 @@ class PersonalMessageSendTest(ZulipTestCase):
             user_group,
             acting_user=None,
         )
-        self.send_personal_message(user_profile, cordelia)
+        cordelia.refresh_from_db()
+
+        with self.assertRaises(DirectMessagePermissionError):
+            self.send_personal_message(cordelia, polonius)
+
+        with self.assert_database_query_count(16):
+            self.send_personal_message(user_profile, cordelia)
+
+        # Test that query count decreases if setting is set to a system group.
+        members_group = NamedUserGroup.objects.get(
+            name=SystemGroups.MEMBERS, realm=realm, is_system_group=True
+        )
+        do_change_realm_permission_group_setting(
+            realm,
+            "direct_message_permission_group",
+            members_group,
+            acting_user=None,
+        )
+        with self.assert_database_query_count(15):
+            self.send_personal_message(user_profile, cordelia)
 
         do_change_realm_permission_group_setting(
             realm,

zerver/tests/test_user_groups.py

@@ -39,6 +39,7 @@ from zerver.lib.test_helpers import most_recent_usermessage
 from zerver.lib.timestamp import datetime_to_timestamp
 from zerver.lib.types import UserGroupMembersData, UserGroupMembersDict
 from zerver.lib.user_groups import (
+    check_user_has_permission_by_role,
     get_direct_user_groups,
     get_recursive_group_members,
     get_recursive_group_members_union_for_groups,
@@ -65,7 +66,7 @@ from zerver.models import (
     UserGroupMembership,
     UserProfile,
 )
-from zerver.models.groups import SystemGroups
+from zerver.models.groups import SystemGroups, get_realm_system_groups_name_dict
 from zerver.models.realms import get_realm
 
 
@@ -576,6 +577,111 @@ class UserGroupTestCase(ZulipTestCase):
             add_subgroups_to_user_group(test_group, [], acting_user=None)
             remove_subgroups_from_user_group(test_group, [], acting_user=None)
 
+    def test_check_user_has_permission_by_role(self) -> None:
+        realm = get_realm("zulip")
+        system_groups_name_dict = get_realm_system_groups_name_dict(realm.id)
+
+        desdemona = self.example_user("desdemona")
+        iago = self.example_user("iago")
+        shiva = self.example_user("shiva")
+        hamlet = self.example_user("hamlet")
+        othello = self.example_user("othello")
+        polonius = self.example_user("polonius")
+
+        nobody_group = NamedUserGroup.objects.get(
+            name=SystemGroups.NOBODY, realm=realm, is_system_group=True
+        )
+        self.assertFalse(
+            check_user_has_permission_by_role(desdemona, nobody_group.id, system_groups_name_dict)
+        )
+
+        owners_group = NamedUserGroup.objects.get(
+            name=SystemGroups.OWNERS, realm=realm, is_system_group=True
+        )
+        self.assertFalse(
+            check_user_has_permission_by_role(iago, owners_group.id, system_groups_name_dict)
+        )
+        self.assertTrue(
+            check_user_has_permission_by_role(desdemona, owners_group.id, system_groups_name_dict)
+        )
+
+        admins_group = NamedUserGroup.objects.get(
+            name=SystemGroups.ADMINISTRATORS, realm=realm, is_system_group=True
+        )
+        self.assertFalse(
+            check_user_has_permission_by_role(shiva, admins_group.id, system_groups_name_dict)
+        )
+        self.assertTrue(
+            check_user_has_permission_by_role(iago, admins_group.id, system_groups_name_dict)
+        )
+
+        moderators_group = NamedUserGroup.objects.get(
+            name=SystemGroups.MODERATORS, realm=realm, is_system_group=True
+        )
+        self.assertFalse(
+            check_user_has_permission_by_role(hamlet, moderators_group.id, system_groups_name_dict)
+        )
+        self.assertTrue(
+            check_user_has_permission_by_role(shiva, moderators_group.id, system_groups_name_dict)
+        )
+
+        members_group = NamedUserGroup.objects.get(
+            name=SystemGroups.MEMBERS, realm=realm, is_system_group=True
+        )
+        self.assertFalse(
+            check_user_has_permission_by_role(polonius, members_group.id, system_groups_name_dict)
+        )
+        self.assertTrue(
+            check_user_has_permission_by_role(hamlet, members_group.id, system_groups_name_dict)
+        )
+
+        full_members_group = NamedUserGroup.objects.get(
+            name=SystemGroups.FULL_MEMBERS, realm=realm, is_system_group=True
+        )
+        do_set_realm_property(realm, "waiting_period_threshold", 10, acting_user=None)
+        hamlet.refresh_from_db()
+        shiva.refresh_from_db()
+        othello.refresh_from_db()
+        polonius.refresh_from_db()
+
+        hamlet.date_joined = timezone_now() - timedelta(days=9)
+        hamlet.save()
+
+        shiva.date_joined = timezone_now() - timedelta(days=9)
+        shiva.save()
+
+        othello.date_joined = timezone_now() - timedelta(days=11)
+        othello.save()
+
+        polonius.date_joined = timezone_now() - timedelta(days=11)
+        polonius.save()
+
+        self.assertFalse(
+            check_user_has_permission_by_role(
+                polonius, full_members_group.id, system_groups_name_dict
+            )
+        )
+        self.assertFalse(
+            check_user_has_permission_by_role(
+                hamlet, full_members_group.id, system_groups_name_dict
+            )
+        )
+        self.assertTrue(
+            check_user_has_permission_by_role(
+                othello, full_members_group.id, system_groups_name_dict
+            )
+        )
+        self.assertTrue(
+            check_user_has_permission_by_role(shiva, full_members_group.id, system_groups_name_dict)
+        )
+
+        everyone_group = NamedUserGroup.objects.get(
+            name=SystemGroups.EVERYONE, realm=realm, is_system_group=True
+        )
+        self.assertTrue(
+            check_user_has_permission_by_role(polonius, everyone_group.id, system_groups_name_dict)
+        )
+
 
 class UserGroupAPITestCase(UserGroupTestCase):
     def test_user_group_create(self) -> None: