Refactoring before implementing the updated spaces summary. (#10527)

This should have no user-visible changes, but refactors some pieces of
the SpaceSummaryHandler before adding support for the updated
MSC2946.

4 files changed, 198 insertions, 136 deletions

diff --git a/changelog.d/10527.misc b/changelog.d/10527.misc
new file mode 100644
index 0000000000..3cf22f9daf
--- /dev/null
+++ b/changelog.d/10527.misc
@@ -0,0 +1 @@
+Prepare for the new spaces summary endpoint (updates to [MSC2946](https://github.com/matrix-org/matrix-doc/pull/2946)).
diff --git a/synapse/federation/federation_client.py b/synapse/federation/federation_client.py
index b7a10da15a..007d1a27dc 100644
--- a/synapse/federation/federation_client.py
+++ b/synapse/federation/federation_client.py
@@ -1290,7 +1290,7 @@ class FederationClient(FederationBase):
         )
 
 
-@attr.s(frozen=True, slots=True)
+@attr.s(frozen=True, slots=True, auto_attribs=True)
 class FederationSpaceSummaryEventResult:
     """Represents a single event in the result of a successful get_space_summary call.
 
@@ -1299,12 +1299,13 @@ class FederationSpaceSummaryEventResult:
     object attributes.
     """
 
-    event_type = attr.ib(type=str)
-    state_key = attr.ib(type=str)
-    via = attr.ib(type=Sequence[str])
+    event_type: str
+    room_id: str
+    state_key: str
+    via: Sequence[str]
 
     # the raw data, including the above keys
-    data = attr.ib(type=JsonDict)
+    data: JsonDict
 
     @classmethod
     def from_json_dict(cls, d: JsonDict) -> "FederationSpaceSummaryEventResult":
@@ -1321,6 +1322,10 @@ class FederationSpaceSummaryEventResult:
         if not isinstance(event_type, str):
             raise ValueError("Invalid event: 'event_type' must be a str")
 
+        room_id = d.get("room_id")
+        if not isinstance(room_id, str):
+            raise ValueError("Invalid event: 'room_id' must be a str")
+
         state_key = d.get("state_key")
         if not isinstance(state_key, str):
             raise ValueError("Invalid event: 'state_key' must be a str")
@@ -1335,15 +1340,15 @@ class FederationSpaceSummaryEventResult:
         if any(not isinstance(v, str) for v in via):
             raise ValueError("Invalid event: 'via' must be a list of strings")
 
-        return cls(event_type, state_key, via, d)
+        return cls(event_type, room_id, state_key, via, d)
 
 
-@attr.s(frozen=True, slots=True)
+@attr.s(frozen=True, slots=True, auto_attribs=True)
 class FederationSpaceSummaryResult:
     """Represents the data returned by a successful get_space_summary call."""
 
-    rooms = attr.ib(type=Sequence[JsonDict])
-    events = attr.ib(type=Sequence[FederationSpaceSummaryEventResult])
+    rooms: Sequence[JsonDict]
+    events: Sequence[FederationSpaceSummaryEventResult]
 
     @classmethod
     def from_json_dict(cls, d: JsonDict) -> "FederationSpaceSummaryResult":
diff --git a/synapse/handlers/space_summary.py b/synapse/handlers/space_summary.py
index 5f7d4602bd..3eb232c83e 100644
--- a/synapse/handlers/space_summary.py
+++ b/synapse/handlers/space_summary.py
@@ -16,7 +16,17 @@ import itertools
 import logging
 import re
 from collections import deque
-from typing import TYPE_CHECKING, Iterable, List, Optional, Sequence, Set, Tuple
+from typing import (
+    TYPE_CHECKING,
+    Collection,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+)
 
 import attr
 
@@ -116,20 +126,22 @@ class SpaceSummaryHandler:
             max_children = max_rooms_per_space if processed_rooms else None
 
             if is_in_room:
-                room, events = await self._summarize_local_room(
+                room_entry = await self._summarize_local_room(
                     requester, None, room_id, suggested_only, max_children
                 )
 
+                events: Collection[JsonDict] = []
+                if room_entry:
+                    rooms_result.append(room_entry.room)
+                    events = room_entry.children
+
                 logger.debug(
                     "Query of local room %s returned events %s",
                     room_id,
                     ["%s->%s" % (ev["room_id"], ev["state_key"]) for ev in events],
                 )
-
-                if room:
-                    rooms_result.append(room)
             else:
-                fed_rooms, fed_events = await self._summarize_remote_room(
+                fed_rooms = await self._summarize_remote_room(
                     queue_entry,
                     suggested_only,
                     max_children,
@@ -141,12 +153,10 @@ class SpaceSummaryHandler:
                 # user is not permitted see.
                 #
                 # Filter the returned results to only what is accessible to the user.
-                room_ids = set()
                 events = []
-                for room in fed_rooms:
-                    fed_room_id = room.get("room_id")
-                    if not fed_room_id or not isinstance(fed_room_id, str):
-                        continue
+                for room_entry in fed_rooms:
+                    room = room_entry.room
+                    fed_room_id = room_entry.room_id
 
                     # The room should only be included in the summary if:
                     #     a. the user is in the room;
@@ -189,21 +199,17 @@ class SpaceSummaryHandler:
                     # The user can see the room, include it!
                     if include_room:
                         rooms_result.append(room)
-                        room_ids.add(fed_room_id)
+                        events.extend(room_entry.children)
 
                     # All rooms returned don't need visiting again (even if the user
                     # didn't have access to them).
                     processed_rooms.add(fed_room_id)
 
-                for event in fed_events:
-                    if event.get("room_id") in room_ids:
-                        events.append(event)
-
                 logger.debug(
                     "Query of %s returned rooms %s, events %s",
                     room_id,
-                    [room.get("room_id") for room in fed_rooms],
-                    ["%s->%s" % (ev["room_id"], ev["state_key"]) for ev in fed_events],
+                    [room_entry.room.get("room_id") for room_entry in fed_rooms],
+                    ["%s->%s" % (ev["room_id"], ev["state_key"]) for ev in events],
                 )
 
             # the room we queried may or may not have been returned, but don't process
@@ -283,20 +289,20 @@ class SpaceSummaryHandler:
                 # already done this room
                 continue
 
-            logger.debug("Processing room %s", room_id)
-
-            room, events = await self._summarize_local_room(
+            room_entry = await self._summarize_local_room(
                 None, origin, room_id, suggested_only, max_rooms_per_space
             )
 
             processed_rooms.add(room_id)
 
-            if room:
-                rooms_result.append(room)
-                events_result.extend(events)
+            if room_entry:
+                rooms_result.append(room_entry.room)
+                events_result.extend(room_entry.children)
 
-            # add any children to the queue
-            room_queue.extend(edge_event["state_key"] for edge_event in events)
+                # add any children to the queue
+                room_queue.extend(
+                    edge_event["state_key"] for edge_event in room_entry.children
+                )
 
         return {"rooms": rooms_result, "events": events_result}
 
@@ -307,7 +313,7 @@ class SpaceSummaryHandler:
         room_id: str,
         suggested_only: bool,
         max_children: Optional[int],
-    ) -> Tuple[Optional[JsonDict], Sequence[JsonDict]]:
+    ) -> Optional["_RoomEntry"]:
         """
         Generate a room entry and a list of event entries for a given room.
 
@@ -326,21 +332,16 @@ class SpaceSummaryHandler:
                 to a server-set limit.
 
         Returns:
-            A tuple of:
-                The room information, if the room should be returned to the
-                user. None, otherwise.
-
-                An iterable of the sorted children events. This may be limited
-                to a maximum size or may include all children.
+            A room entry if the room should be returned. None, otherwise.
         """
         if not await self._is_room_accessible(room_id, requester, origin):
-            return None, ()
+            return None
 
         room_entry = await self._build_room_entry(room_id)
 
         # If the room is not a space, return just the room information.
         if room_entry.get("room_type") != RoomTypes.SPACE:
-            return room_entry, ()
+            return _RoomEntry(room_id, room_entry)
 
         # Otherwise, look for child rooms/spaces.
         child_events = await self._get_child_events(room_id)
@@ -363,7 +364,7 @@ class SpaceSummaryHandler:
                 )
             )
 
-        return room_entry, events_result
+        return _RoomEntry(room_id, room_entry, events_result)
 
     async def _summarize_remote_room(
         self,
@@ -371,7 +372,7 @@ class SpaceSummaryHandler:
         suggested_only: bool,
         max_children: Optional[int],
         exclude_rooms: Iterable[str],
-    ) -> Tuple[Sequence[JsonDict], Sequence[JsonDict]]:
+    ) -> Iterable["_RoomEntry"]:
         """
         Request room entries and a list of event entries for a given room by querying a remote server.
 
@@ -386,11 +387,7 @@ class SpaceSummaryHandler:
                 Rooms IDs which do not need to be summarized.
 
         Returns:
-            A tuple of:
-                An iterable of rooms.
-
-                An iterable of the sorted children events. This may be limited
-                to a maximum size or may include all children.
+            An iterable of room entries.
         """
         room_id = room.room_id
         logger.info("Requesting summary for %s via %s", room_id, room.via)
@@ -414,11 +411,30 @@ class SpaceSummaryHandler:
                 e,
                 exc_info=logger.isEnabledFor(logging.DEBUG),
             )
-            return (), ()
+            return ()
+
+        # Group the events by their room.
+        children_by_room: Dict[str, List[JsonDict]] = {}
+        for ev in res.events:
+            if ev.event_type == EventTypes.SpaceChild:
+                children_by_room.setdefault(ev.room_id, []).append(ev.data)
+
+        # Generate the final results.
+        results = []
+        for fed_room in res.rooms:
+            fed_room_id = fed_room.get("room_id")
+            if not fed_room_id or not isinstance(fed_room_id, str):
+                continue
 
-        return res.rooms, tuple(
-            ev.data for ev in res.events if ev.event_type == EventTypes.SpaceChild
-        )
+            results.append(
+                _RoomEntry(
+                    fed_room_id,
+                    fed_room,
+                    children_by_room.get(fed_room_id, []),
+                )
+            )
+
+        return results
 
     async def _is_room_accessible(
         self, room_id: str, requester: Optional[str], origin: Optional[str]
@@ -606,10 +622,21 @@ class SpaceSummaryHandler:
         return sorted(filter(_has_valid_via, events), key=_child_events_comparison_key)
 
 
-@attr.s(frozen=True, slots=True)
+@attr.s(frozen=True, slots=True, auto_attribs=True)
 class _RoomQueueEntry:
-    room_id = attr.ib(type=str)
-    via = attr.ib(type=Sequence[str])
+    room_id: str
+    via: Sequence[str]
+
+
+@attr.s(frozen=True, slots=True, auto_attribs=True)
+class _RoomEntry:
+    room_id: str
+    # The room summary for this room.
+    room: JsonDict
+    # An iterable of the sorted, stripped children events for children of this room.
+    #
+    # This may not include all children.
+    children: Collection[JsonDict] = ()
 
 
 def _has_valid_via(e: EventBase) -> bool:
diff --git a/tests/handlers/test_space_summary.py b/tests/handlers/test_space_summary.py
index 3f73ad7f94..f982a8c8b4 100644
--- a/tests/handlers/test_space_summary.py
+++ b/tests/handlers/test_space_summary.py
@@ -26,7 +26,7 @@ from synapse.api.constants import (
 from synapse.api.errors import AuthError
 from synapse.api.room_versions import RoomVersions
 from synapse.events import make_event_from_dict
-from synapse.handlers.space_summary import _child_events_comparison_key
+from synapse.handlers.space_summary import _child_events_comparison_key, _RoomEntry
 from synapse.rest import admin
 from synapse.rest.client.v1 import login, room
 from synapse.server import HomeServer
@@ -351,26 +351,30 @@ class SpaceSummaryTestCase(unittest.HomeserverTestCase):
             #   events before child events).
 
             # Note that these entries are brief, but should contain enough info.
-            rooms = [
-                {
-                    "room_id": subspace,
-                    "world_readable": True,
-                    "room_type": RoomTypes.SPACE,
-                },
-                {
-                    "room_id": subroom,
-                    "world_readable": True,
-                },
-            ]
-            event_content = {"via": [fed_hostname]}
-            events = [
-                {
-                    "room_id": subspace,
-                    "state_key": subroom,
-                    "content": event_content,
-                },
+            return [
+                _RoomEntry(
+                    subspace,
+                    {
+                        "room_id": subspace,
+                        "world_readable": True,
+                        "room_type": RoomTypes.SPACE,
+                    },
+                    [
+                        {
+                            "room_id": subspace,
+                            "state_key": subroom,
+                            "content": {"via": [fed_hostname]},
+                        }
+                    ],
+                ),
+                _RoomEntry(
+                    subroom,
+                    {
+                        "room_id": subroom,
+                        "world_readable": True,
+                    },
+                ),
             ]
-            return rooms, events
 
         # Add a room to the space which is on another server.
         self._add_child(self.space, subspace, self.token)
@@ -436,70 +440,95 @@ class SpaceSummaryTestCase(unittest.HomeserverTestCase):
         ):
             # Note that these entries are brief, but should contain enough info.
             rooms = [
-                {
-                    "room_id": public_room,
-                    "world_readable": False,
-                    "join_rules": JoinRules.PUBLIC,
-                },
-                {
-                    "room_id": knock_room,
-                    "world_readable": False,
-                    "join_rules": JoinRules.KNOCK,
-                },
-                {
-                    "room_id": not_invited_room,
-                    "world_readable": False,
-                    "join_rules": JoinRules.INVITE,
-                },
-                {
-                    "room_id": invited_room,
-                    "world_readable": False,
-                    "join_rules": JoinRules.INVITE,
-                },
-                {
-                    "room_id": restricted_room,
-                    "world_readable": False,
-                    "join_rules": JoinRules.MSC3083_RESTRICTED,
-                    "allowed_spaces": [],
-                },
-                {
-                    "room_id": restricted_accessible_room,
-                    "world_readable": False,
-                    "join_rules": JoinRules.MSC3083_RESTRICTED,
-                    "allowed_spaces": [self.room],
-                },
-                {
-                    "room_id": world_readable_room,
-                    "world_readable": True,
-                    "join_rules": JoinRules.INVITE,
-                },
-                {
-                    "room_id": joined_room,
-                    "world_readable": False,
-                    "join_rules": JoinRules.INVITE,
-                },
-            ]
-
-            # Place each room in the sub-space.
-            event_content = {"via": [fed_hostname]}
-            events = [
-                {
-                    "room_id": subspace,
-                    "state_key": room["room_id"],
-                    "content": event_content,
-                }
-                for room in rooms
+                _RoomEntry(
+                    public_room,
+                    {
+                        "room_id": public_room,
+                        "world_readable": False,
+                        "join_rules": JoinRules.PUBLIC,
+                    },
+                ),
+                _RoomEntry(
+                    knock_room,
+                    {
+                        "room_id": knock_room,
+                        "world_readable": False,
+                        "join_rules": JoinRules.KNOCK,
+                    },
+                ),
+                _RoomEntry(
+                    not_invited_room,
+                    {
+                        "room_id": not_invited_room,
+                        "world_readable": False,
+                        "join_rules": JoinRules.INVITE,
+                    },
+                ),
+                _RoomEntry(
+                    invited_room,
+                    {
+                        "room_id": invited_room,
+                        "world_readable": False,
+                        "join_rules": JoinRules.INVITE,
+                    },
+                ),
+                _RoomEntry(
+                    restricted_room,
+                    {
+                        "room_id": restricted_room,
+                        "world_readable": False,
+                        "join_rules": JoinRules.MSC3083_RESTRICTED,
+                        "allowed_spaces": [],
+                    },
+                ),
+                _RoomEntry(
+                    restricted_accessible_room,
+                    {
+                        "room_id": restricted_accessible_room,
+                        "world_readable": False,
+                        "join_rules": JoinRules.MSC3083_RESTRICTED,
+                        "allowed_spaces": [self.room],
+                    },
+                ),
+                _RoomEntry(
+                    world_readable_room,
+                    {
+                        "room_id": world_readable_room,
+                        "world_readable": True,
+                        "join_rules": JoinRules.INVITE,
+                    },
+                ),
+                _RoomEntry(
+                    joined_room,
+                    {
+                        "room_id": joined_room,
+                        "world_readable": False,
+                        "join_rules": JoinRules.INVITE,
+                    },
+                ),
             ]
 
             # Also include the subspace.
             rooms.insert(
                 0,
-                {
-                    "room_id": subspace,
-                    "world_readable": True,
-                },
+                _RoomEntry(
+                    subspace,
+                    {
+                        "room_id": subspace,
+                        "world_readable": True,
+                    },
+                    # Place each room in the sub-space.
+                    [
+                        {
+                            "room_id": subspace,
+                            "state_key": room.room_id,
+                            "content": {"via": [fed_hostname]},
+                        }
+                        for room in rooms
+                    ],
+                ),
             )
-            return rooms, events
+            return rooms
 
         # Add a room to the space which is on another server.
         self._add_child(self.space, subspace, self.token)