retention: Use Message.realm to simplify private message query.

We no longer need to do the inner joins to figure out the message's
realm and split up the cross-realm and regular case - now we just look
at zerver_message.realm directly.

zerver/lib/retention.py

@@ -53,8 +53,6 @@ from zerver.models import (
     Stream,
     SubMessage,
     UserMessage,
-    get_realm,
-    get_user_including_cross_realm,
 )
 
 logger = logging.getLogger("zulip.retention")
@@ -225,24 +223,16 @@ def move_expired_personal_and_huddle_messages_to_archive(
     check_date = timezone_now() - timedelta(days=message_retention_days)
 
     # This function will archive appropriate messages and their related objects.
-    internal_realm = get_realm(settings.SYSTEM_BOT_REALM)
-    cross_realm_bot_ids = [
-        get_user_including_cross_realm(email, internal_realm).id
-        for email in settings.CROSS_REALM_BOT_EMAILS
-    ]
     recipient_types = (Recipient.PERSONAL, Recipient.HUDDLE)
 
-    # Archive expired personal and huddle Messages in the realm, except cross-realm messages.
-    # The condition zerver_userprofile.realm_id = {realm_id} assures the row won't be
-    # a message sent by a cross-realm bot, because cross-realm bots have their own separate realm.
+    # Archive expired personal and huddle Messages in the realm, including cross-realm messages.
     query = SQL(
         """
     INSERT INTO zerver_archivedmessage ({dst_fields}, archive_transaction_id)
         SELECT {src_fields}, {archive_transaction_id}
         FROM zerver_message
         INNER JOIN zerver_recipient ON zerver_recipient.id = zerver_message.recipient_id
-        INNER JOIN zerver_userprofile ON zerver_userprofile.id = zerver_message.sender_id
-        WHERE zerver_userprofile.realm_id = {realm_id}
+        WHERE zerver_message.realm_id = {realm_id}
             AND zerver_recipient.type in {recipient_types}
             AND zerver_message.date_sent < {check_date}
         LIMIT {chunk_size}
@@ -255,42 +245,12 @@ def move_expired_personal_and_huddle_messages_to_archive(
         query,
         type=ArchiveTransaction.RETENTION_POLICY_BASED,
         realm=realm,
-        cross_realm_bot_ids=Literal(tuple(cross_realm_bot_ids)),
         realm_id=Literal(realm.id),
         recipient_types=Literal(recipient_types),
         check_date=Literal(check_date.isoformat()),
         chunk_size=chunk_size,
     )
 
-    # Archive cross-realm personal messages to users in the realm.  We
-    # don't archive cross-realm huddle messages via retention policy,
-    # as we don't support them as a feature in Zulip, and the query to
-    # find and delete them would be a lot of complexity and potential
-    # performance work for a case that doesn't actually happen.
-    query = SQL(
-        """
-    INSERT INTO zerver_archivedmessage ({dst_fields}, archive_transaction_id)
-        SELECT {src_fields}, {archive_transaction_id}
-        FROM zerver_message
-        INNER JOIN zerver_userprofile recipient_profile ON recipient_profile.recipient_id = zerver_message.recipient_id
-        WHERE zerver_message.sender_id IN {cross_realm_bot_ids}
-            AND recipient_profile.realm_id = {realm_id}
-            AND zerver_message.date_sent < {check_date}
-        LIMIT {chunk_size}
-    ON CONFLICT (id) DO UPDATE SET archive_transaction_id = {archive_transaction_id}
-    RETURNING id
-    """
-    )
-    message_count += run_archiving_in_chunks(
-        query,
-        type=ArchiveTransaction.RETENTION_POLICY_BASED,
-        realm=realm,
-        cross_realm_bot_ids=Literal(tuple(cross_realm_bot_ids)),
-        realm_id=Literal(realm.id),
-        check_date=Literal(check_date.isoformat()),
-        chunk_size=chunk_size,
-    )
-
     return message_count
 
 

zerver/tests/test_retention.py

@@ -147,6 +147,19 @@ class ArchiveMessagesTestingBase(RetentionTestingBase):
         assert msg_id is not None
         return msg_id
 
+    def _send_personal_message_to_cross_realm_bot(self) -> int:
+        # Send message from bot to users from different realm.
+        bot_email = "notification-bot@zulip.com"
+        internal_realm = get_realm(settings.SYSTEM_BOT_REALM)
+        zulip_user = self.example_user("hamlet")
+        msg_id = internal_send_private_message(
+            sender=zulip_user,
+            recipient_user=get_system_bot(bot_email, internal_realm.id),
+            content="test message",
+        )
+        assert msg_id is not None
+        return msg_id
+
     def _make_expired_zulip_messages(self, message_quantity: int) -> List[int]:
         msg_ids = list(
             Message.objects.order_by("id")
@@ -298,9 +311,13 @@ class TestArchiveMessagesGeneral(ArchiveMessagesTestingBase):
 
     def test_cross_realm_personal_message_archiving(self) -> None:
         """Check that cross-realm personal messages get correctly archived."""
+
+        # We want to test on a set of cross-realm messages of both kinds -
+        # from a bot to a user, and from a user to a bot.
         msg_ids = [self._send_cross_realm_personal_message() for i in range(1, 7)]
+        msg_ids += [self._send_personal_message_to_cross_realm_bot() for i in range(1, 7)]
         usermsg_ids = self._get_usermessage_ids(msg_ids)
-        # Make the message expired on the recipient's realm:
+        # Make the message expired in the Zulip realm.:
         self._change_messages_date_sent(msg_ids, timezone_now() - timedelta(ZULIP_REALM_DAYS + 1))
 
         archive_messages()