From 7bf0793c945d25a2f5d45c18302ae63f65e7cf9f Mon Sep 17 00:00:00 2001 From: Aman Agrawal Date: Mon, 17 Apr 2023 13:02:07 +0000 Subject: [PATCH] scheduled_messages: Move database related function to /actions. This would match the pattern we use for other modules and also shrink the giant message_send.py. --- zerver/actions/message_send.py | 97 --------------------- zerver/actions/scheduled_messages.py | 121 +++++++++++++++++++++++++++ zerver/views/message_send.py | 2 +- zerver/views/scheduled_messages.py | 14 +--- 4 files changed, 124 insertions(+), 110 deletions(-) create mode 100644 zerver/actions/scheduled_messages.py diff --git a/zerver/actions/message_send.py b/zerver/actions/message_send.py index fdcf765f66..dd9edfcb12 100644 --- a/zerver/actions/message_send.py +++ b/zerver/actions/message_send.py @@ -60,7 +60,6 @@ from zerver.lib.notification_data import ( ) from zerver.lib.queue import queue_json_publish from zerver.lib.recipient_users import recipient_for_user_profiles -from zerver.lib.scheduled_messages import access_scheduled_message from zerver.lib.stream_subscription import ( get_subscriptions_for_send_message, num_subscribers_for_stream_id, @@ -79,7 +78,6 @@ from zerver.models import ( Message, Realm, Recipient, - ScheduledMessage, Stream, UserMessage, UserPresence, @@ -459,63 +457,6 @@ def get_service_bot_events( return event_dict -def do_schedule_messages(send_message_requests: Sequence[SendMessageRequest]) -> List[int]: - scheduled_messages: List[ScheduledMessage] = [] - - for send_request in send_message_requests: - scheduled_message = ScheduledMessage() - scheduled_message.sender = send_request.message.sender - scheduled_message.recipient = send_request.message.recipient - topic_name = send_request.message.topic_name() - scheduled_message.set_topic_name(topic_name=topic_name) - rendering_result = render_markdown( - send_request.message, send_request.message.content, send_request.realm - ) - scheduled_message.content = send_request.message.content - scheduled_message.rendered_content = rendering_result.rendered_content - scheduled_message.sending_client = send_request.message.sending_client - scheduled_message.stream = send_request.stream - scheduled_message.realm = send_request.realm - assert send_request.deliver_at is not None - scheduled_message.scheduled_timestamp = send_request.deliver_at - if send_request.delivery_type == "send_later": - scheduled_message.delivery_type = ScheduledMessage.SEND_LATER - elif send_request.delivery_type == "remind": - scheduled_message.delivery_type = ScheduledMessage.REMIND - - scheduled_messages.append(scheduled_message) - - ScheduledMessage.objects.bulk_create(scheduled_messages) - return [scheduled_message.id for scheduled_message in scheduled_messages] - - -def edit_scheduled_message( - scheduled_message_id: int, send_request: SendMessageRequest, sender: UserProfile -) -> int: - with transaction.atomic(): - scheduled_message_object = access_scheduled_message(sender, scheduled_message_id) - - # Handles the race between us initiating this transaction and user sending us the edit request. - if scheduled_message_object.delivered is True: - raise JsonableError(_("Scheduled message was already sent")) - - # Only override fields that user can change. - scheduled_message_object.recipient = send_request.message.recipient - topic_name = send_request.message.topic_name() - scheduled_message_object.set_topic_name(topic_name=topic_name) - rendering_result = render_markdown( - send_request.message, send_request.message.content, send_request.realm - ) - scheduled_message_object.content = send_request.message.content - scheduled_message_object.rendered_content = rendering_result.rendered_content - scheduled_message_object.sending_client = send_request.message.sending_client - scheduled_message_object.stream = send_request.stream - assert send_request.deliver_at is not None - scheduled_message_object.scheduled_timestamp = send_request.deliver_at - scheduled_message_object.save() - return scheduled_message_id - - def build_message_send_dict( message: Message, stream: Optional[Stream] = None, @@ -1177,44 +1118,6 @@ def check_send_message( return do_send_messages([message])[0] -def check_schedule_message( - sender: UserProfile, - client: Client, - recipient_type_name: str, - message_to: Union[Sequence[str], Sequence[int]], - topic_name: Optional[str], - message_content: str, - scheduled_message_id: Optional[int], - delivery_type: str, - deliver_at: datetime.datetime, - realm: Optional[Realm] = None, - forwarder_user_profile: Optional[UserProfile] = None, -) -> int: - addressee = Addressee.legacy_build(sender, recipient_type_name, message_to, topic_name) - - send_request = check_message( - sender, - client, - addressee, - message_content, - realm=realm, - forwarder_user_profile=forwarder_user_profile, - ) - send_request.deliver_at = deliver_at - send_request.delivery_type = delivery_type - - recipient = send_request.message.recipient - if delivery_type == "remind" and ( - recipient.type != Recipient.STREAM and recipient.type_id != sender.id - ): - raise JsonableError(_("Reminders can only be set for streams.")) - - if scheduled_message_id is not None: - return edit_scheduled_message(scheduled_message_id, send_request, sender) - - return do_schedule_messages([send_request])[0] - - def send_rate_limited_pm_notification_to_bot_owner( sender: UserProfile, realm: Realm, content: str ) -> None: diff --git a/zerver/actions/scheduled_messages.py b/zerver/actions/scheduled_messages.py new file mode 100644 index 0000000000..66166ad881 --- /dev/null +++ b/zerver/actions/scheduled_messages.py @@ -0,0 +1,121 @@ +import datetime +from typing import List, Optional, Sequence, Union + +from django.db import transaction +from django.utils.translation import gettext as _ + +from zerver.actions.message_send import check_message +from zerver.lib.addressee import Addressee +from zerver.lib.exceptions import JsonableError +from zerver.lib.message import SendMessageRequest, render_markdown +from zerver.lib.scheduled_messages import access_scheduled_message +from zerver.models import Client, Realm, Recipient, ScheduledMessage, UserProfile +from zerver.tornado.django_api import send_event + + +def check_schedule_message( + sender: UserProfile, + client: Client, + recipient_type_name: str, + message_to: Union[Sequence[str], Sequence[int]], + topic_name: Optional[str], + message_content: str, + scheduled_message_id: Optional[int], + delivery_type: str, + deliver_at: datetime.datetime, + realm: Optional[Realm] = None, + forwarder_user_profile: Optional[UserProfile] = None, +) -> int: + addressee = Addressee.legacy_build(sender, recipient_type_name, message_to, topic_name) + + send_request = check_message( + sender, + client, + addressee, + message_content, + realm=realm, + forwarder_user_profile=forwarder_user_profile, + ) + send_request.deliver_at = deliver_at + send_request.delivery_type = delivery_type + + recipient = send_request.message.recipient + if delivery_type == "remind" and ( + recipient.type != Recipient.STREAM and recipient.type_id != sender.id + ): + raise JsonableError(_("Reminders can only be set for streams.")) + + if scheduled_message_id is not None: + return edit_scheduled_message(scheduled_message_id, send_request, sender) + + return do_schedule_messages([send_request])[0] + + +def do_schedule_messages(send_message_requests: Sequence[SendMessageRequest]) -> List[int]: + scheduled_messages: List[ScheduledMessage] = [] + + for send_request in send_message_requests: + scheduled_message = ScheduledMessage() + scheduled_message.sender = send_request.message.sender + scheduled_message.recipient = send_request.message.recipient + topic_name = send_request.message.topic_name() + scheduled_message.set_topic_name(topic_name=topic_name) + rendering_result = render_markdown( + send_request.message, send_request.message.content, send_request.realm + ) + scheduled_message.content = send_request.message.content + scheduled_message.rendered_content = rendering_result.rendered_content + scheduled_message.sending_client = send_request.message.sending_client + scheduled_message.stream = send_request.stream + scheduled_message.realm = send_request.realm + assert send_request.deliver_at is not None + scheduled_message.scheduled_timestamp = send_request.deliver_at + if send_request.delivery_type == "send_later": + scheduled_message.delivery_type = ScheduledMessage.SEND_LATER + elif send_request.delivery_type == "remind": + scheduled_message.delivery_type = ScheduledMessage.REMIND + + scheduled_messages.append(scheduled_message) + + ScheduledMessage.objects.bulk_create(scheduled_messages) + return [scheduled_message.id for scheduled_message in scheduled_messages] + + +def edit_scheduled_message( + scheduled_message_id: int, send_request: SendMessageRequest, sender: UserProfile +) -> int: + with transaction.atomic(): + scheduled_message_object = access_scheduled_message(sender, scheduled_message_id) + + # Handles the race between us initiating this transaction and user sending us the edit request. + if scheduled_message_object.delivered is True: + raise JsonableError(_("Scheduled message was already sent")) + + # Only override fields that user can change. + scheduled_message_object.recipient = send_request.message.recipient + topic_name = send_request.message.topic_name() + scheduled_message_object.set_topic_name(topic_name=topic_name) + rendering_result = render_markdown( + send_request.message, send_request.message.content, send_request.realm + ) + scheduled_message_object.content = send_request.message.content + scheduled_message_object.rendered_content = rendering_result.rendered_content + scheduled_message_object.sending_client = send_request.message.sending_client + scheduled_message_object.stream = send_request.stream + assert send_request.deliver_at is not None + scheduled_message_object.scheduled_timestamp = send_request.deliver_at + scheduled_message_object.save() + return scheduled_message_id + + +def delete_scheduled_message(user_profile: UserProfile, scheduled_message_id: int) -> None: + scheduled_message_object = access_scheduled_message(user_profile, scheduled_message_id) + scheduled_message_id = scheduled_message_object.id + scheduled_message_object.delete() + + event = { + "type": "scheduled_message", + "op": "remove", + "scheduled_message_id": scheduled_message_id, + } + send_event(user_profile.realm, event, [user_profile.id]) diff --git a/zerver/views/message_send.py b/zerver/views/message_send.py index 07fcf2f704..9fbb256b35 100644 --- a/zerver/views/message_send.py +++ b/zerver/views/message_send.py @@ -10,7 +10,6 @@ from django.utils.timezone import now as timezone_now from django.utils.translation import gettext as _ from zerver.actions.message_send import ( - check_schedule_message, check_send_message, compute_irc_user_fullname, compute_jabber_user_fullname, @@ -18,6 +17,7 @@ from zerver.actions.message_send import ( extract_private_recipients, extract_stream_indicator, ) +from zerver.actions.scheduled_messages import check_schedule_message from zerver.lib.exceptions import JsonableError from zerver.lib.message import render_markdown from zerver.lib.request import REQ, RequestNotes, has_request_variables diff --git a/zerver/views/scheduled_messages.py b/zerver/views/scheduled_messages.py index 489b3754c8..247e586a55 100644 --- a/zerver/views/scheduled_messages.py +++ b/zerver/views/scheduled_messages.py @@ -2,11 +2,10 @@ from typing import List, TypedDict from django.http import HttpRequest, HttpResponse +from zerver.actions.scheduled_messages import delete_scheduled_message from zerver.lib.request import has_request_variables from zerver.lib.response import json_success -from zerver.lib.scheduled_messages import access_scheduled_message from zerver.models import ScheduledMessage, UserProfile, get_recipient_ids -from zerver.tornado.django_api import send_event class ScheduledMessageDict(TypedDict): @@ -49,14 +48,5 @@ def fetch_scheduled_messages(request: HttpRequest, user_profile: UserProfile) -> def delete_scheduled_messages( request: HttpRequest, user_profile: UserProfile, scheduled_message_id: int ) -> HttpResponse: - scheduled_message_object = access_scheduled_message(user_profile, scheduled_message_id) - scheduled_message_id = scheduled_message_object.id - scheduled_message_object.delete() - - event = { - "type": "scheduled_message", - "op": "remove", - "scheduled_message_id": scheduled_message_id, - } - send_event(user_profile.realm, event, [user_profile.id]) + delete_scheduled_message(user_profile, scheduled_message_id) return json_success(request)