Support pagination tokens from sync/messages in the relations API

synapse/rest/client/relations.py

@@ -152,7 +152,7 @@ async def on_GET(
             events, now, bundle_aggregations=aggregations
         )
 
-        return_value = pagination_chunk.to_dict()
+        return_value = await pagination_chunk.to_dict(self.store)
         return_value["chunk"] = serialized_events
         return_value["original_event"] = original_event
 
@@ -243,7 +243,7 @@ async def on_GET(
                 to_token=to_token,
             )
 
-        return 200, pagination_chunk.to_dict()
+        return 200, await pagination_chunk.to_dict(self.store)
 
 
 class RelationAggregationGroupPaginationServlet(RestServlet):
@@ -335,7 +335,7 @@ async def on_GET(
         now = self.clock.time_msec()
         serialized_events = self._event_serializer.serialize_events(events, now)
 
-        return_value = result.to_dict()
+        return_value = await result.to_dict(self.store)
         return_value["chunk"] = serialized_events
 
         return 200, return_value

synapse/storage/databases/main/relations.py

@@ -49,6 +49,7 @@
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
+    from synapse.storage.databases.main import DataStore
 
 logger = logging.getLogger(__name__)
 
@@ -682,13 +683,15 @@ async def _get_bundled_aggregation_for_event(
 
         annotations = await self.get_aggregation_groups_for_event(event_id, room_id)
         if annotations.chunk:
-            aggregations.annotations = annotations.to_dict()
+            aggregations.annotations = await annotations.to_dict(
+                cast("DataStore", self)
+            )
 
         references = await self.get_relations_for_event(
             event_id, room_id, RelationTypes.REFERENCE, direction="f"
         )
         if references.chunk:
-            aggregations.references = references.to_dict()
+            aggregations.references = await references.to_dict(cast("DataStore", self))
 
         # If this event is the start of a thread, include a summary of the replies.
         if self._msc3440_enabled:

synapse/storage/relations.py

@@ -13,13 +13,16 @@
 # limitations under the License.
 
 import logging
-from typing import Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 import attr
 
 from synapse.api.errors import SynapseError
 from synapse.types import JsonDict
 
+if TYPE_CHECKING:
+    from synapse.storage.databases.main import DataStore
+
 logger = logging.getLogger(__name__)
 
 
@@ -39,14 +42,14 @@ class PaginationChunk:
     next_batch: Optional[Any] = None
     prev_batch: Optional[Any] = None
 
-    def to_dict(self) -> Dict[str, Any]:
+    async def to_dict(self, store: "DataStore") -> Dict[str, Any]:
         d = {"chunk": self.chunk}
 
         if self.next_batch:
-            d["next_batch"] = self.next_batch.to_string()
+            d["next_batch"] = await self.next_batch.to_string(store)
 
         if self.prev_batch:
-            d["prev_batch"] = self.prev_batch.to_string()
+            d["prev_batch"] = await self.prev_batch.to_string(store)
 
         return d
 
@@ -75,7 +78,7 @@ def from_string(string: str) -> "RelationPaginationToken":
         except ValueError:
             raise SynapseError(400, "Invalid relation pagination token")
 
-    def to_string(self) -> str:
+    async def to_string(self, store: "DataStore") -> str:
         return "%d-%d" % (self.topological, self.stream)
 
     def as_tuple(self) -> Tuple[Any, ...]:
@@ -105,7 +108,7 @@ def from_string(string: str) -> "AggregationPaginationToken":
         except ValueError:
             raise SynapseError(400, "Invalid aggregation pagination token")
 
-    def to_string(self) -> str:
+    async def to_string(self, store: "DataStore") -> str:
         return "%d-%d" % (self.count, self.stream)
 
     def as_tuple(self) -> Tuple[Any, ...]: