From 887d20795f702fd0231093b4249cb8d9c460f1bc Mon Sep 17 00:00:00 2001
From: Tim Abbott <tabbott@zulipchat.com>
Date: Tue, 21 Aug 2018 15:03:00 -0700
Subject: [PATCH] message flags: Add where_starred helper and use it.

The previous query ended up doing a scan of all a user's UserMessage
rows, not just the ones tracked in the `starred` index.
---
 zerver/lib/message.py | 13 ++++---------
 zerver/models.py      | 19 ++++++++++++++-----
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/zerver/lib/message.py b/zerver/lib/message.py
index eebde87eb0..bb90b3db6c 100644
--- a/zerver/lib/message.py
+++ b/zerver/lib/message.py
@@ -706,18 +706,13 @@ def get_muted_stream_ids(user_profile: UserProfile) -> List[int]:
     return muted_stream_ids
 
 def get_starred_message_ids(user_profile: UserProfile) -> List[int]:
-    rows = UserMessage.objects.filter(
+    return list(UserMessage.objects.filter(
         user_profile=user_profile,
-        flags=UserMessage.flags.starred
-    ).values(
-        'message_id'
+    ).extra(
+        where=[UserMessage.where_starred()]
     ).order_by(
         'message_id'
-    )
-    starred_message_ids = [
-        row['message_id']
-        for row in rows]
-    return starred_message_ids
+    ).values_list('message_id', flat=True)[0:10000])
 
 def get_raw_unread_data(user_profile: UserProfile) -> RawUnreadMessagesResult:
 
diff --git a/zerver/models.py b/zerver/models.py
index 1afaa89004..c343175f26 100644
--- a/zerver/models.py
+++ b/zerver/models.py
@@ -1543,13 +1543,22 @@ class AbstractUserMessage(models.Model):
 
     @staticmethod
     def where_unread() -> str:
-        # Use this for Django ORM queries where we are getting lots
-        # of rows.  This custom SQL plays nice with our partial indexes.
-        # Grep the code for example usage.
-        #
-        # This optimization is only helpful for checking a flag being False.
+        # Use this for Django ORM queries to access unread message.
+        # This custom SQL plays nice with our partial indexes.  Grep
+        # the code for example usage.
         return 'flags & 1 = 0'
 
+    @staticmethod
+    def where_starred() -> str:
+        # Use this for Django ORM queries to access starred messages.
+        # This custom SQL plays nice with our partial indexes.  Grep
+        # the code for example usage.
+        #
+        # The key detail is that e.g.
+        #   UserMessage.objects.filter(user_profile=user_profile, flags=UserMessage.flags.starred)
+        # will generate a query involving `flags & 2 = 2`, which doesn't match our index.
+        return 'flags & 2 <> 0'
+
     def flags_list(self) -> List[str]:
         flags = int(self.flags)
         return self.flags_list_for_flags(flags)