Split a storage function in two so that we don't have to do extra work.

3 files changed, 14 insertions, 8 deletions

diff --git a/synapse/federation/federation_server.py b/synapse/federation/federation_server.py
index 25c0014f97..2b46188c91 100644
--- a/synapse/federation/federation_server.py
+++ b/synapse/federation/federation_server.py
@@ -417,13 +417,13 @@ class FederationServer(FederationBase):
                 pdu.internal_metadata.outlier = True
             elif min_depth and pdu.depth > min_depth:
                 if get_missing and prevs - seen:
-                    latest_tuples = yield self.store.get_latest_events_in_room(
+                    latest = yield self.store.get_latest_event_ids_in_room(
                         pdu.room_id
                     )
 
                     # We add the prev events that we have seen to the latest
                     # list to ensure the remote server doesn't give them to us
-                    latest = set(e_id for e_id, _, _ in latest_tuples)
+                    latest = set(latest)
                     latest |= seen
 
                     missing_events = yield self.get_missing_events(
diff --git a/synapse/state.py b/synapse/state.py
index ba2500d61c..9dddb77d5b 100644
--- a/synapse/state.py
+++ b/synapse/state.py
@@ -86,12 +86,7 @@ class StateHandler(object):
         If `event_type` is specified, then the method returns only the one
         event (or None) with that `event_type` and `state_key`.
         """
-        events = yield self.store.get_latest_events_in_room(room_id)
-
-        event_ids = [
-            e_id
-            for e_id, _, _ in events
-        ]
+        event_ids = yield self.store.get_latest_event_ids_in_room(room_id)
 
         cache = None
         if self._state_cache is not None:
diff --git a/synapse/storage/event_federation.py b/synapse/storage/event_federation.py
index 54a3c9d805..8bbb42c27e 100644
--- a/synapse/storage/event_federation.py
+++ b/synapse/storage/event_federation.py
@@ -96,11 +96,22 @@ class EventFederationStore(SQLBaseStore):
             room_id,
         )
 
+    def get_latest_event_ids_in_room(self, room_id):
+        return self._simple_select_onecol(
+            table="event_forward_extremities",
+            keyvalues={
+                "room_id": room_id,
+            },
+            retcol="event_id",
+            desc="get_latest_events_in_room",
+        )
+
     def _get_latest_events_in_room(self, txn, room_id):
         sql = (
             "SELECT e.event_id, e.depth FROM events as e "
             "INNER JOIN event_forward_extremities as f "
             "ON e.event_id = f.event_id "
+            "AND e.room_id = f.room_id "
             "WHERE f.room_id = ?"
         )