From fd5a091b30925e23e687f08ca22a5227b08bc1b7 Mon Sep 17 00:00:00 2001 From: Alex Vandiver Date: Fri, 22 Mar 2024 06:04:07 +0000 Subject: [PATCH] messages: Only check the UserMessage row if necessary. For the common case of not needing to reference the UserMessage row later, and for being a stream without private history, the UserMessage row is irrelevant. Convert `has_user_message` to a thunk, and defer loading it unless necessary. --- zerver/lib/message.py | 22 ++++++++++++-------- zerver/lib/push_notifications.py | 4 +++- zerver/tests/test_message_move_stream.py | 26 ++++++++++++------------ zerver/tests/test_message_move_topic.py | 12 +++++------ 4 files changed, 36 insertions(+), 28 deletions(-) diff --git a/zerver/lib/message.py b/zerver/lib/message.py index 6769cb20ba..bae598335a 100644 --- a/zerver/lib/message.py +++ b/zerver/lib/message.py @@ -2,6 +2,7 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta from typing import ( Any, + Callable, Collection, Dict, List, @@ -23,6 +24,7 @@ from django.utils.timezone import now as timezone_now from django.utils.translation import gettext as _ from django_stubs_ext import ValuesQuerySet from psycopg2.sql import SQL +from returns.curry import partial from analytics.lib.counts import COUNT_STATS from analytics.models import RealmCount @@ -318,9 +320,14 @@ def access_message( if get_user_message == "object": user_message = get_usermessage_by_message_id(user_profile, message_id) - has_user_message = user_message is not None + has_user_message = lambda: user_message is not None + elif get_user_message == "exists": + local_exists = UserMessage.objects.filter( + user_profile=user_profile, message_id=message_id + ).exists() + has_user_message = lambda: local_exists else: - has_user_message = UserMessage.objects.filter( + has_user_message = lambda: UserMessage.objects.filter( user_profile=user_profile, message_id=message_id ).exists() @@ -328,7 +335,7 @@ def access_message( if get_user_message is None: return message if get_user_message == "exists": - return (message, has_user_message) + return (message, local_exists) if get_user_message == "object": return (message, user_message) raise JsonableError(_("Invalid message(s)")) @@ -380,7 +387,7 @@ def has_message_access( user_profile: UserProfile, message: Message, *, - has_user_message: bool, + has_user_message: Callable[[], bool], stream: Optional[Stream] = None, is_subscribed: Optional[bool] = None, ) -> bool: @@ -394,7 +401,7 @@ def has_message_access( if message.recipient.type != Recipient.STREAM: # You can only access direct messages you received - return has_user_message + return has_user_message() if stream is None: stream = Stream.objects.get(id=message.recipient.type_id) @@ -421,7 +428,7 @@ def has_message_access( # (1) Have directly received the message. # AND # (2) Be subscribed to the stream. - return has_user_message and is_subscribed_helper() + return has_user_message() and is_subscribed_helper() # is_history_public_to_subscribers, so check if you're subscribed return is_subscribed_helper() @@ -461,12 +468,11 @@ def bulk_access_messages( subscribed_recipient_ids = set(get_subscribed_stream_recipient_ids_for_user(user_profile)) for message in messages: - has_user_message = message.id in user_message_set is_subscribed = message.recipient_id in subscribed_recipient_ids if has_message_access( user_profile, message, - has_user_message=has_user_message, + has_user_message=partial(lambda m: m.id in user_message_set, message), stream=streams.get(message.recipient_id) if stream is None else stream, is_subscribed=is_subscribed, ): diff --git a/zerver/lib/push_notifications.py b/zerver/lib/push_notifications.py index 8fb3ed38fc..55406642f2 100644 --- a/zerver/lib/push_notifications.py +++ b/zerver/lib/push_notifications.py @@ -1366,7 +1366,9 @@ def handle_push_notification(user_profile_id: int, missed_message: Dict[str, Any mentioned_user_group_id = missed_message.get("mentioned_user_group_id") if mentioned_user_group_id is not None: - user_group = UserGroup.objects.get(id=mentioned_user_group_id, realm=user_profile.realm) + user_group = UserGroup.objects.get( + id=mentioned_user_group_id, realm_id=user_profile.realm_id + ) mentioned_user_group_name = user_group.name # Soft reactivate if pushing to a long_term_idle user that is personally mentioned diff --git a/zerver/tests/test_message_move_stream.py b/zerver/tests/test_message_move_stream.py index 2fdbccdb85..880bad4cfe 100644 --- a/zerver/tests/test_message_move_stream.py +++ b/zerver/tests/test_message_move_stream.py @@ -1098,7 +1098,7 @@ class MessageMoveStreamTest(ZulipTestCase): "iago", "test move stream", "new stream", "test" ) - with self.assert_database_query_count(52), self.assert_memcached_count(14): + with self.assert_database_query_count(51), self.assert_memcached_count(14): result = self.client_patch( f"/json/messages/{msg_id}", { @@ -1202,7 +1202,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user, Message.objects.get(id=msg_id_to_test_acesss), - has_user_message=has_user_message, + has_user_message=lambda: has_user_message, stream=stream, is_subscribed=is_subscribed, ), @@ -1612,7 +1612,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user_losing_access, Message.objects.get(id=msg_id), - has_user_message=True, + has_user_message=lambda: True, stream=old_stream, is_subscribed=True, ), @@ -1628,7 +1628,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user_losing_access, Message.objects.get(id=msg_id), - has_user_message=True, + has_user_message=lambda: True, stream=old_stream, is_subscribed=False, ), @@ -1652,7 +1652,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user_losing_access, Message.objects.get(id=msg_id), - has_user_message=True, + has_user_message=lambda: True, stream=new_stream, is_subscribed=False, ), @@ -1686,7 +1686,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user, Message.objects.get(id=first_msg_id), - has_user_message=True, + has_user_message=lambda: True, stream=old_stream, is_subscribed=True, ), @@ -1702,7 +1702,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user, Message.objects.get(id=first_msg_id), - has_user_message=True, + has_user_message=lambda: True, stream=old_stream, is_subscribed=False, ), @@ -1719,7 +1719,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user, Message.objects.get(id=second_msg_id), - has_user_message=False, + has_user_message=lambda: False, stream=old_stream, is_subscribed=False, ), @@ -1745,7 +1745,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user, Message.objects.get(id=first_msg_id), - has_user_message=True, + has_user_message=lambda: True, stream=new_stream, is_subscribed=True, ), @@ -1755,7 +1755,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user, Message.objects.get(id=second_msg_id), - has_user_message=True, + has_user_message=lambda: True, stream=new_stream, is_subscribed=True, ), @@ -1786,7 +1786,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user, Message.objects.get(id=msg_id), - has_user_message=False, + has_user_message=lambda: False, stream=old_stream, is_subscribed=False, ), @@ -1803,7 +1803,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user, Message.objects.get(id=msg_id), - has_user_message=False, + has_user_message=lambda: False, stream=old_stream, is_subscribed=False, ), @@ -1828,7 +1828,7 @@ class MessageMoveStreamTest(ZulipTestCase): has_message_access( user, Message.objects.get(id=msg_id), - has_user_message=True, + has_user_message=lambda: True, stream=new_stream, is_subscribed=True, ), diff --git a/zerver/tests/test_message_move_topic.py b/zerver/tests/test_message_move_topic.py index 318c0894dd..e115224c24 100644 --- a/zerver/tests/test_message_move_topic.py +++ b/zerver/tests/test_message_move_topic.py @@ -238,7 +238,7 @@ class MessageMoveTopicTest(ZulipTestCase): # state + 1/user with a UserTopic row for the events data) # beyond what is typical were there not UserTopic records to # update. Ideally, we'd eliminate the per-user component. - with self.assert_database_query_count(26): + with self.assert_database_query_count(25): check_update_message( user_profile=hamlet, message_id=message_id, @@ -335,7 +335,7 @@ class MessageMoveTopicTest(ZulipTestCase): set_topic_visibility_policy(desdemona, muted_topics, UserTopic.VisibilityPolicy.MUTED) set_topic_visibility_policy(cordelia, muted_topics, UserTopic.VisibilityPolicy.MUTED) - with self.assert_database_query_count(28): + with self.assert_database_query_count(27): check_update_message( user_profile=desdemona, message_id=message_id, @@ -365,7 +365,7 @@ class MessageMoveTopicTest(ZulipTestCase): ] set_topic_visibility_policy(desdemona, muted_topics, UserTopic.VisibilityPolicy.MUTED) set_topic_visibility_policy(cordelia, muted_topics, UserTopic.VisibilityPolicy.MUTED) - with self.assert_database_query_count(34): + with self.assert_database_query_count(33): check_update_message( user_profile=desdemona, message_id=message_id, @@ -398,7 +398,7 @@ class MessageMoveTopicTest(ZulipTestCase): set_topic_visibility_policy(desdemona, muted_topics, UserTopic.VisibilityPolicy.MUTED) set_topic_visibility_policy(cordelia, muted_topics, UserTopic.VisibilityPolicy.MUTED) - with self.assert_database_query_count(28): + with self.assert_database_query_count(27): check_update_message( user_profile=desdemona, message_id=message_id, @@ -421,7 +421,7 @@ class MessageMoveTopicTest(ZulipTestCase): second_message_id = self.send_stream_message( hamlet, stream_name, topic_name="changed topic name", content="Second message" ) - with self.assert_database_query_count(23): + with self.assert_database_query_count(22): check_update_message( user_profile=desdemona, message_id=second_message_id, @@ -520,7 +520,7 @@ class MessageMoveTopicTest(ZulipTestCase): users_to_be_notified_via_muted_topics_event.append(user_topic.user_profile_id) change_all_topic_name = "Topic 1 edited" - with self.assert_database_query_count(31): + with self.assert_database_query_count(30): check_update_message( user_profile=hamlet, message_id=message_id,