Introduce an RdataSafeValue type and correct some minor type annotation mistakes. (#20071)

Introduce a `SafeRdataValue` type that represents what is safe to push
over the wire in RDATA commands.

Spawns from a mistake when working on
https://github.com/element-hq/synapse/pull/19556

Would have caught the issue that
https://github.com/element-hq/synapse/pull/19556/commits/9bcd110080849575ce11d9b50806423867b5c1b5
fixes.

At the same time, correct some mistaken type annotations that flared up
after this correction.

---------

Signed-off-by: Olivier 'reivilibre <oliverw@matrix.org>

changelog.d/20071.misc

@@ -0,0 +1 @@
+Introduce an `RdataSafeValue` type and correct some minor type annotation mistakes.

synapse/replication/tcp/resource.py

@@ -22,7 +22,7 @@
 
 import logging
 import random
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Sequence
 
 from prometheus_client import Counter
 
@@ -317,7 +317,7 @@ class ReplicationStreamer:
 
 
 def _batch_updates(
-    updates: list[tuple[Token, StreamRow]],
+    updates: Sequence[tuple[Token, StreamRow]],
 ) -> list[tuple[Token | None, StreamRow]]:
     """Takes a list of updates of form [(token, row)] and sets the token to
     None for all rows where the next row has the same token. This is used to

synapse/replication/tcp/streams/_base.py

@@ -26,7 +26,8 @@ from typing import (
     Any,
     Awaitable,
     Callable,
-    TypeVar,
+    Sequence,
+    Union,
 )
 
 import attr
@@ -50,11 +51,36 @@ _STREAM_UPDATE_TARGET_ROW_COUNT = 100
 # A stream position token
 Token = int
 
+RdataSafeValue = Union[
+    str,
+    int,
+    bool,
+    float,
+    None,
+    # We could probably expand to support immutable types of these, but
+    # when they come off the wire they will deserialise to the mutable
+    # types again.
+    # Since nobody is using it right now, stick to `list` and `dict`
+    list["RdataSafeValue"],
+    dict[str, "RdataSafeValue"],
+]
+"""
+Safe types that can be used in RDATA commands and thus used as
+the wire format of stream rows.
+
+Prevents you from thinking you can push e.g. a `frozenset` over
+the wire and get it back on the other end.
+
+At the moment, to be safe, a type has to roundtrip correctly with our JSON codec.
+Consult the RdataCommand `from_line` and `to_line` for information.
+"""
+
 # The type of a stream update row, after JSON deserialisation, but before
 # parsing with Stream.parse_row (which turns it into a `ROW_TYPE`). Normally it's
 # just a row from a database query, though this is dependent on the stream in question.
 #
-StreamRow = TypeVar("StreamRow", bound=tuple)
+# NOTE: Prefer to use tuples, but since we have some streams still using list, support those for now.
+StreamRow = Union[tuple[RdataSafeValue, ...], list[RdataSafeValue]]
 
 # The type returned by the update_function of a stream, as well as get_updates(),
 # get_updates_since, etc.
@@ -64,7 +90,7 @@ StreamRow = TypeVar("StreamRow", bound=tuple)
 #   * `new_last_token` is the new position in stream.
 #   * `limited` is whether there are more updates to fetch.
 #
-StreamUpdateResult = tuple[list[tuple[Token, StreamRow]], Token, bool]
+StreamUpdateResult = tuple[Sequence[tuple[Token, StreamRow]], Token, bool]
 
 # The type of an update_function for a stream
 #
@@ -407,9 +433,9 @@ class TypingStream(Stream):
         if hs.get_instance_name() in hs.config.worker.writers.typing:
             # On the writer, query the typing handler
             typing_writer_handler = hs.get_typing_writer_handler()
-            update_function: Callable[
-                [str, int, int, int], Awaitable[tuple[list[tuple[int, Any]], int, bool]]
-            ] = typing_writer_handler.get_all_typing_updates
+            update_function: UpdateFunction = (
+                typing_writer_handler.get_all_typing_updates
+            )
             self.current_token_function = typing_writer_handler.get_current_token
         else:
             # Query the typing writer process
@@ -819,7 +845,7 @@ class ProfileUpdatesStream(_StreamFromIdGen):
         updates = await self.store.get_updated_profile_updates(
             from_id=from_token, to_id=to_token, limit=limit
         )
-        rows = [
+        rows: list[tuple[int, tuple[RdataSafeValue, ...]]] = [
             (
                 stream_id,
                 # These are the args to `ProfileUpdatesStreamRow`

synapse/replication/tcp/streams/federation.py

@@ -18,13 +18,14 @@
 # [This file includes modifications made by New Vector Limited]
 #
 #
-from typing import TYPE_CHECKING, Any, Awaitable, Callable
+from typing import TYPE_CHECKING
 
 import attr
 
 from synapse.replication.tcp.streams._base import (
     Stream,
     Token,
+    UpdateFunction,
     current_token_without_instance,
     make_http_update_function,
 )
@@ -57,9 +58,7 @@ class FederationStream(Stream):
             self.current_token_func = current_token_without_instance(
                 federation_sender.get_current_token
             )
-            update_function: Callable[
-                [str, int, int, int], Awaitable[tuple[list[tuple[int, Any]], int, bool]]
-            ] = federation_sender.get_replication_rows
+            update_function: UpdateFunction = federation_sender.get_replication_rows
 
         elif hs.should_send_federation():
             # federation sender: Query master process