diff --git a/zerver/lib/actions.py b/zerver/lib/actions.py index 98585a05ba..c127279112 100644 --- a/zerver/lib/actions.py +++ b/zerver/lib/actions.py @@ -4600,7 +4600,7 @@ def do_delete_messages(realm: Realm, messages: Iterable[Message]) -> None: (event, message_id_to_notifiable_users.get(message.id, [])) ) - move_messages_to_archive(message_ids) + move_messages_to_archive(message_ids, realm=realm) for event, users_to_notify in events_and_users_to_notify: # TODO: Figure out some kind of bulk event that we could send just one of? send_event(realm, event, users_to_notify) diff --git a/zerver/lib/retention.py b/zerver/lib/retention.py index 5420bf48be..6660b1d48f 100644 --- a/zerver/lib/retention.py +++ b/zerver/lib/retention.py @@ -383,7 +383,8 @@ def get_realms_and_streams_for_archiving() -> List[Tuple[Realm, List[Stream]]]: return [(realm_id_to_realm[realm_id], realm_id_to_streams_list[realm_id]) for realm_id in realm_id_to_realm] -def move_messages_to_archive(message_ids: List[int], chunk_size: int=MESSAGE_BATCH_SIZE) -> None: +def move_messages_to_archive(message_ids: List[int], realm: Optional[Realm]=None, + chunk_size: int=MESSAGE_BATCH_SIZE) -> None: query = SQL(""" INSERT INTO zerver_archivedmessage ({dst_fields}, archive_transaction_id) SELECT {src_fields}, {archive_transaction_id} @@ -397,6 +398,7 @@ def move_messages_to_archive(message_ids: List[int], chunk_size: int=MESSAGE_BAT query, type=ArchiveTransaction.MANUAL, message_ids=Literal(tuple(message_ids)), + realm=realm, chunk_size=chunk_size, ) diff --git a/zerver/tests/test_retention.py b/zerver/tests/test_retention.py index 352519ea1d..d3f6457a8c 100644 --- a/zerver/tests/test_retention.py +++ b/zerver/tests/test_retention.py @@ -511,6 +511,19 @@ class MoveMessageToArchiveGeneral(MoveMessageToArchiveBase): restore_all_data_from_archive() self._verify_restored_data(msg_ids, usermsg_ids) + def test_move_messages_to_archive_with_realm_argument(self) -> None: + realm = get_realm("zulip") + msg_ids = [self.send_personal_message(self.sender, self.recipient) + for i in range(0, 3)] + usermsg_ids = self._get_usermessage_ids(msg_ids) + + self._assert_archive_empty() + move_messages_to_archive(message_ids=msg_ids, realm=realm) + self._verify_archive_data(msg_ids, usermsg_ids) + + archive_transaction = ArchiveTransaction.objects.last() + self.assertEqual(archive_transaction.realm, realm) + def test_stream_messages_archiving(self) -> None: msg_ids = [self.send_stream_message(self.sender, "Verona") for i in range(0, 3)] @@ -535,6 +548,31 @@ class MoveMessageToArchiveGeneral(MoveMessageToArchiveBase): with self.assertRaises(Message.DoesNotExist): move_messages_to_archive(message_ids=msg_ids) + def test_archiving_messages_multiple_realms(self) -> None: + """ + Verifies that move_messages_to_archive works correctly + if called on messages in multiple realms. + """ + iago = self.example_user("iago") + othello = self.example_user("othello") + + cordelia = self.lear_user("cordelia") + king = self.lear_user("king") + + zulip_msg_ids = [self.send_personal_message(iago, othello) + for i in range(0, 3)] + leary_msg_ids = [self.send_personal_message(cordelia, king) + for i in range(0, 3)] + msg_ids = zulip_msg_ids + leary_msg_ids + usermsg_ids = self._get_usermessage_ids(msg_ids) + + self._assert_archive_empty() + move_messages_to_archive(message_ids=msg_ids) + self._verify_archive_data(msg_ids, usermsg_ids) + + restore_all_data_from_archive() + self._verify_restored_data(msg_ids, usermsg_ids) + def test_archiving_messages_with_attachment(self) -> None: self._create_attachments() realm_id = get_realm("zulip").id