Add order_by argument for paginating room events

This can be used to flip betweek topological ordering (the default)
and stream ordering as needed by the caller.

Author: Fizzadar

synapse/storage/databases/main/stream.py

@@ -652,22 +652,31 @@ def f(txn: LoggingTransaction) -> List[_EventDictReturn]:
         return ret
 
     async def get_recent_events_for_room(
-        self, room_id: str, limit: int, end_token: RoomStreamToken
+        self,
+        room_id: str,
+        limit: int,
+        end_token: RoomStreamToken,
+        order_by: str = "topological",
     ) -> Tuple[List[EventBase], RoomStreamToken]:
         """Get the most recent events in the room in topological ordering.
 
         Args:
             room_id
             limit
             end_token: The stream token representing now.
+            order_by: Either 'topological' or 'stream' to indicate the order in
+                which results should be returned.
 
         Returns:
             A list of events and a token pointing to the start of the returned
             events. The events returned are in ascending topological order.
         """
 
         rows, token = await self.get_recent_event_ids_for_room(
-            room_id, limit, end_token
+            room_id,
+            limit,
+            end_token,
+            order_by,
         )
 
         events = await self.get_events_as_list(
@@ -679,14 +688,20 @@ async def get_recent_events_for_room(
         return events, token
 
     async def get_recent_event_ids_for_room(
-        self, room_id: str, limit: int, end_token: RoomStreamToken
+        self,
+        room_id: str,
+        limit: int,
+        end_token: RoomStreamToken,
+        order_by: str = "topological",
     ) -> Tuple[List[_EventDictReturn], RoomStreamToken]:
         """Get the most recent events in the room in topological ordering.
 
         Args:
             room_id
             limit
             end_token: The stream token representing now.
+            order_by: Either 'topological' or 'stream' to indicate the order in
+                which results should be returned.
 
         Returns:
             A list of _EventDictReturn and a token pointing to the start of the
@@ -701,6 +716,7 @@ async def get_recent_event_ids_for_room(
             self._paginate_room_events_txn,
             room_id,
             from_token=end_token,
+            order_by=order_by,
             limit=limit,
         )
 
@@ -1099,6 +1115,7 @@ def _paginate_room_events_txn(
         from_token: RoomStreamToken,
         to_token: Optional[RoomStreamToken] = None,
         direction: str = "b",
+        order_by: str = "topological",
         limit: int = -1,
         event_filter: Optional[Filter] = None,
     ) -> Tuple[List[_EventDictReturn], RoomStreamToken]:
@@ -1111,6 +1128,8 @@ def _paginate_room_events_txn(
             to_token: A token which if given limits the results to only those before
             direction: Either 'b' or 'f' to indicate whether we are paginating
                 forwards or backwards from `from_key`.
+            order_by: Either 'topological' or 'stream' to indicate the order in
+                which results should be returned.
             limit: The maximum number of events to return.
             event_filter: If provided filters the events to
                 those that match the filter.
@@ -1123,6 +1142,7 @@ def _paginate_room_events_txn(
         """
 
         assert int(limit) >= 0
+        assert order_by in ("topological", "stream")
 
         # Tokens really represent positions between elements, but we use
         # the convention of pointing to the event before the gap. Hence
@@ -1133,6 +1153,12 @@ def _paginate_room_events_txn(
         else:
             order = "ASC"
 
+        order_clause = """ORDER BY event.topological_ordering %(order)s, event.stream_ordering %(order)s"""
+        if order_by == "stream":
+            order_clause = """ORDER BY event.stream_ordering %(order)s, event.topological_ordering %(order)s"""
+
+        order_clause = order_clause % {"order": order}
+
         # The bounds for the stream tokens are complicated by the fact
         # that we need to handle the instance_map part of the tokens. We do this
         # by fetching all events between the min stream token and the maximum
@@ -1228,13 +1254,13 @@ def _paginate_room_events_txn(
             FROM events AS event
             %(join_clause)s
             WHERE event.outlier = ? AND event.room_id = ? AND %(bounds)s
-            ORDER BY event.topological_ordering %(order)s,
-            event.stream_ordering %(order)s LIMIT ?
+            %(order_clause)s
+            LIMIT ?
         """ % {
             "select_keywords": select_keywords,
             "join_clause": join_clause,
             "bounds": bounds,
-            "order": order,
+            "order_clause": order_clause,
         }
 
         txn.execute(sql, args)
@@ -1275,6 +1301,7 @@ async def paginate_room_events(
         from_key: RoomStreamToken,
         to_key: Optional[RoomStreamToken] = None,
         direction: str = "b",
+        order_by: str = "topological",
         limit: int = -1,
         event_filter: Optional[Filter] = None,
     ) -> Tuple[List[EventBase], RoomStreamToken]:
@@ -1286,6 +1313,8 @@ async def paginate_room_events(
             to_key: A token which if given limits the results to only those before
             direction: Either 'b' or 'f' to indicate whether we are paginating
                 forwards or backwards from `from_key`.
+            order_by: Either 'topological' or 'stream' to indicate the order in
+                which results should be returned.
             limit: The maximum number of events to return.
             event_filter: If provided filters the events to those that match the filter.
 
@@ -1303,6 +1332,7 @@ async def paginate_room_events(
             from_key,
             to_key,
             direction,
+            order_by,
             limit,
             event_filter,
         )