Add local support for the new spaces summary endpoint

synapse/handlers/space_summary.py

@@ -38,9 +38,11 @@
     Membership,
     RoomTypes,
 )
+from synapse.api.errors import Codes, SynapseError
 from synapse.events import EventBase
 from synapse.events.utils import format_event_for_client_v2
 from synapse.types import JsonDict
+from synapse.util.stringutils import random_string
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
@@ -67,6 +69,12 @@ def __init__(self, hs: "HomeServer"):
         self._server_name = hs.hostname
         self._federation_client = hs.get_federation_client()
 
+        # TODO Allow for multiple workers to share this data.
+        # TODO Expire pagination tokens.
+        self._pagination_sessions: Dict[
+            str, Tuple[Deque[_RoomQueueEntry], Set[str]]
+        ] = {}
+
     async def get_space_summary(
         self,
         requester: str,
@@ -278,13 +286,19 @@ async def get_room_hierarchy(
         # world-readable)
         await self._auth.check_user_in_room_or_world_readable(room_id, requester)
 
-        # TODO Handle pulling previously persisted state by the pagination token.
+        # If this is continuing a previous session, pull the persisted data.
+        if from_token:
+            if from_token not in self._pagination_sessions:
+                raise SynapseError(400, "Unknown pagination token", Codes.INVALID_PARAM)
 
-        # the queue of rooms to process
-        room_queue = deque((_RoomQueueEntry(room_id, ()),))
+            # TODO Assert that other parameters have not changed.
+            (room_queue, processed_rooms) = self._pagination_sessions[from_token]
+        else:
+            # the queue of rooms to process
+            room_queue = deque((_RoomQueueEntry(room_id, ()),))
 
-        # Rooms we have already processed.
-        processed_rooms: Set[str] = set()
+            # Rooms we have already processed.
+            processed_rooms = set()
 
         rooms_result: List[JsonDict] = []
 
@@ -333,9 +347,13 @@ async def get_room_hierarchy(
                 # TODO Federation.
                 pass
 
-        result = {"rooms": rooms_result}
+        result: JsonDict = {"rooms": rooms_result}
 
-        # TODO Handle adding a pagination token (and persisting state).
+        # If there's additional data, generate a pagination token (and persist state).
+        if room_queue:
+            next_token = random_string(24)
+            result["next_token"] = next_token
+            self._pagination_sessions[next_token] = (room_queue, processed_rooms)
 
         return result
 

tests/handlers/test_space_summary.py

@@ -123,12 +123,17 @@ def prepare(self, reactor, clock, hs: HomeServer):
         self.room = self.helper.create_room_as(self.user, tok=self.token)
         self._add_child(self.space, self.room, self.token)
 
-    def _add_child(self, space_id: str, room_id: str, token: str) -> None:
+    def _add_child(
+        self, space_id: str, room_id: str, token: str, order: Optional[str] = None
+    ) -> None:
         """Add a child room to a space."""
+        content = {"via": [self.hs.hostname]}
+        if order is not None:
+            content["order"] = order
         self.helper.send_state(
             space_id,
             event_type=EventTypes.SpaceChild,
-            body={"via": [self.hs.hostname]},
+            body=content,
             tok=token,
             state_key=room_id,
         )
@@ -389,6 +394,39 @@ def test_complex_space(self):
         )
         self._assert_hierarchy(result, expected)
 
+    def test_pagination(self):
+        """Test simple pagination works."""
+        room_ids = []
+        for i in range(1, 10):
+            room = self.helper.create_room_as(self.user, tok=self.token)
+            self._add_child(self.space, room, self.token, order=str(i))
+            room_ids.append(room)
+        # The room created initially doesn't have an order, so comes last.
+        room_ids.append(self.room)
+
+        result = self.get_success(
+            self.handler.get_room_hierarchy(self.user, self.space, limit=7)
+        )
+        # The result should have the space and all of the links, plus some of the
+        # rooms and a pagination token.
+        expected = [(self.space, room_ids)] + [
+            (room_id, ()) for room_id in room_ids[:6]
+        ]
+        self._assert_hierarchy(result, expected)
+        self.assertIn("next_token", result)
+
+        # Check the next page.
+        result = self.get_success(
+            self.handler.get_room_hierarchy(
+                self.user, self.space, limit=5, from_token=result["next_token"]
+            )
+        )
+        # The result should have the space and the room in it, along with a link
+        # from space -> room.
+        expected = [(room_id, ()) for room_id in room_ids[6:]]
+        self._assert_hierarchy(result, expected)
+        self.assertNotIn("next_token", result)
+
     def test_fed_complex(self):
         """
         Return data over federation and ensure that it is handled properly.