From 088a4c7cf0390a6d2a4359807ec80b4ad357c2e6 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <eric.eastwood@beta.gouv.fr>
Date: Thu, 22 Aug 2024 18:26:37 -0500
Subject: [PATCH] Use `simple_upsert_txn` to update `sliding_sync_joined_rooms`

See https://github.com/element-hq/synapse/pull/17512#discussion_r1726817206
---
 synapse/storage/database.py              | 18 ++++-----
 synapse/storage/databases/main/events.py | 50 ++++++++----------------
 2 files changed, 26 insertions(+), 42 deletions(-)

diff --git a/synapse/storage/database.py b/synapse/storage/database.py
index 66a7238deb..ba2616b479 100644
--- a/synapse/storage/database.py
+++ b/synapse/storage/database.py
@@ -1255,9 +1255,9 @@ class DatabasePool:
         self,
         txn: LoggingTransaction,
         table: str,
-        keyvalues: Dict[str, Any],
-        values: Dict[str, Any],
-        insertion_values: Optional[Dict[str, Any]] = None,
+        keyvalues: Mapping[str, Any],
+        values: Mapping[str, Any],
+        insertion_values: Optional[Mapping[str, Any]] = None,
         where_clause: Optional[str] = None,
     ) -> bool:
         """
@@ -1300,9 +1300,9 @@ class DatabasePool:
         self,
         txn: LoggingTransaction,
         table: str,
-        keyvalues: Dict[str, Any],
-        values: Dict[str, Any],
-        insertion_values: Optional[Dict[str, Any]] = None,
+        keyvalues: Mapping[str, Any],
+        values: Mapping[str, Any],
+        insertion_values: Optional[Mapping[str, Any]] = None,
         where_clause: Optional[str] = None,
         lock: bool = True,
     ) -> bool:
@@ -1381,9 +1381,9 @@ class DatabasePool:
         self,
         txn: LoggingTransaction,
         table: str,
-        keyvalues: Dict[str, Any],
-        values: Dict[str, Any],
-        insertion_values: Optional[Dict[str, Any]] = None,
+        keyvalues: Mapping[str, Any],
+        values: Mapping[str, Any],
+        insertion_values: Optional[Mapping[str, Any]] = None,
         where_clause: Optional[str] = None,
     ) -> bool:
         """
diff --git a/synapse/storage/databases/main/events.py b/synapse/storage/databases/main/events.py
index b8ad60194c..bc19003544 100644
--- a/synapse/storage/databases/main/events.py
+++ b/synapse/storage/databases/main/events.py
@@ -1708,15 +1708,8 @@ class PersistEventsStore:
             # persisting stack (see
             # `_update_sliding_sync_tables_with_new_persisted_events_txn()`)
             #
-            # Pulling keys/values separately is safe and will produce congruent lists
-            sliding_sync_updates_keys = (
-                sliding_sync_table_changes.joined_room_updates.keys()
-            )
-            sliding_sync_updates_values = (
-                sliding_sync_table_changes.joined_room_updates.values()
-            )
             # We only need to update when one of the relevant state values has changed
-            if sliding_sync_updates_keys:
+            if sliding_sync_table_changes.joined_room_updates:
                 # This should be *some* value that points to a real event in the room if
                 # we are still joined to the room.
                 assert (
@@ -1724,31 +1717,22 @@ class PersistEventsStore:
                     is not None
                 )
 
-                args: List[Any] = [
-                    room_id,
-                    sliding_sync_table_changes.joined_room_best_effort_most_recent_stream_ordering,
-                ]
-                args.extend(iter(sliding_sync_updates_values))
-
-                # We don't update `event_stream_ordering` `ON CONFLICT` because it's
-                # simpler and we can just rely on
-                # `_update_sliding_sync_tables_with_new_persisted_events_txn()` to
-                # do the right thing (same for `bump_stamp`). The only reason we're
-                # inserting `event_stream_ordering` here is because the column has a
-                # `NON NULL` constraint and we need some answer.
-                txn.execute(
-                    f"""
-                    INSERT INTO sliding_sync_joined_rooms
-                        (room_id, event_stream_ordering, {", ".join(sliding_sync_updates_keys)})
-                    VALUES (
-                        ?, ?,
-                        {", ".join("?" for _ in sliding_sync_updates_values)}
-                    )
-                    ON CONFLICT (room_id)
-                    DO UPDATE SET
-                        {", ".join(f"{key} = EXCLUDED.{key}" for key in sliding_sync_updates_keys)}
-                    """,
-                    args,
+                self.db_pool.simple_upsert_txn(
+                    txn,
+                    table="sliding_sync_joined_rooms",
+                    keyvalues={"room_id": room_id},
+                    values=sliding_sync_table_changes.joined_room_updates,
+                    insertion_values={
+                        # The reason we're only *inserting* `event_stream_ordering` here
+                        # is because the column has a `NON NULL` constraint and we need
+                        # *some* answer. If the row already exists, we are trying to
+                        # avoid doing an `UPDATE` and accidentally overwriting the value
+                        # with some stale data since this is just a "best effort" value.
+                        # It's better to just rely on
+                        # `_update_sliding_sync_tables_with_new_persisted_events_txn()`
+                        # to do the right thing (same for `bump_stamp`).
+                        "event_stream_ordering": sliding_sync_table_changes.joined_room_best_effort_most_recent_stream_ordering
+                    },
                 )
 
         # We now update `local_current_membership`. We do this regardless