Fetch membership counts all at once

1 files changed, 11 insertions, 22 deletions

diff --git a/synapse/storage/roommember.py b/synapse/storage/roommember.py
index 4bd1669458..7617913326 100644
--- a/synapse/storage/roommember.py
+++ b/synapse/storage/roommember.py
@@ -142,26 +142,9 @@ class RoomMemberWorkerStore(EventsWorkerStore):
 
         return self.runInteraction("get_room_summary", _get_room_summary_txn)
 
-    def _get_user_count_in_room_txn(self, txn, room_id, membership):
+    def _get_user_counts_in_room_txn(self, txn, room_id):
         """
-        See get_user_count_in_room.
-        """
-        sql = (
-            "SELECT count(*) FROM room_memberships as m"
-            " INNER JOIN current_state_events as c"
-            " ON m.event_id = c.event_id "
-            " AND m.room_id = c.room_id "
-            " AND m.user_id = c.state_key"
-            " WHERE c.type = 'm.room.member' AND c.room_id = ? AND m.membership = ?"
-        )
-
-        txn.execute(sql, (room_id, membership))
-        row = txn.fetchone()
-        return row[0]
-
-    def get_user_count_in_room(self, room_id, membership):
-        """
-        Get the user count in a room with a particular membership.
+        Get the user count in a room by membership.
 
         Args:
             room_id (str)
@@ -170,9 +153,15 @@ class RoomMemberWorkerStore(EventsWorkerStore):
         Returns:
             Deferred[int]
         """
-        return self.runInteraction(
-            "get_users_in_room", self._get_user_count_in_room_txn, room_id, membership
-        )
+        sql = """
+        SELECT m.membership, count(*) FROM room_memberships as m
+            INNER JOIN current_state_events as c USING(event_id)
+            WHERE c.type = 'm.room.member' AND c.room_id = ?
+            GROUP BY m.membership
+        """
+
+        txn.execute(sql, (room_id,))
+        return {row[0]: row[1] for row in txn}
 
     @cached()
     def get_invited_rooms_for_user(self, user_id):