1 files changed, 13 insertions, 0 deletions
diff --git a/synapse/storage/databases/main/room_batch.py b/synapse/storage/databases/main/room_batch.py
index 300a563c9e..dcbce8fdcf 100644
--- a/synapse/storage/databases/main/room_batch.py
+++ b/synapse/storage/databases/main/room_batch.py
@@ -36,3 +36,16 @@ class RoomBatchStore(SQLBaseStore):
             retcol="event_id",
             allow_none=True,
         )
+
+    async def store_state_group_id_for_event_id(
+        self, event_id: str, state_group_id: int
+    ) -> Optional[str]:
+        {
+            await self.db_pool.simple_upsert(
+                table="event_to_state_groups",
+                keyvalues={"event_id": event_id},
+                values={"state_group": state_group_id, "event_id": event_id},
+                # Unique constraint on event_id so we don't have to lock
+                lock=False,
+            )
+        }