Add back the foreign key constraint

synapse/handlers/federation.py

@@ -584,10 +584,9 @@ async def do_invite_join(
                 # The background process is responsible for unmarking this flag,
                 # even if the join fails.
                 await self.store.store_partial_state_room(
-                    room_id,
-                    ret.servers_in_room,
-                    ret.event.event_id,
-                    self.store.get_device_stream_token(),
+                    room_id=room_id,
+                    servers=ret.servers_in_room,
+                    device_lists_stream_id=self.store.get_device_stream_token(),
                 )
 
             try:
@@ -613,6 +612,14 @@ async def do_invite_join(
                     room_id,
                 )
                 raise LimitExceededError(msg=e.msg, errcode=e.errcode, retry_after_ms=0)
+            else:
+                # Record the join event id for future use (when we finish the full
+                # join). We have to do this after persisting the event to keep foreign
+                # key constraints intact.
+                if ret.partial_state:
+                    await self.store.write_partial_state_rooms_join_event_id(
+                        room_id, event.event_id
+                    )
             finally:
                 # Always kick off the background process that asynchronously fetches
                 # state for the room.

synapse/storage/databases/main/room.py

@@ -1777,7 +1777,6 @@ async def store_partial_state_room(
         self,
         room_id: str,
         servers: Collection[str],
-        join_event_id: str,
         device_lists_stream_id: int,
     ) -> None:
         """Mark the given room as containing events with partial state.
@@ -1786,11 +1785,12 @@ async def store_partial_state_room(
         room, which helps us to keep other homeservers in sync when we finally fully
         join this room.
 
+        We do not include a `join_event_id` here---we need to wait for the join event
+        to be persisted first.
+
         Args:
             room_id: the ID of the room
             servers: other servers known to be in the room
-            join_event_id: the event ID of the join membership event returned in the
-                (partial) /send_join response.
             device_lists_stream_id: the device_lists stream ID at the time when we first
                 joined the room.
         """
@@ -1799,7 +1799,6 @@ async def store_partial_state_room(
             self._store_partial_state_room_txn,
             room_id,
             servers,
-            join_event_id,
             device_lists_stream_id,
         )
 
@@ -1808,7 +1807,6 @@ def _store_partial_state_room_txn(
         txn: LoggingTransaction,
         room_id: str,
         servers: Collection[str],
-        join_event_id: str,
         device_lists_stream_id: int,
     ) -> None:
         DatabasePool.simple_insert_txn(
@@ -1817,7 +1815,8 @@ def _store_partial_state_room_txn(
             values={
                 "room_id": room_id,
                 "device_lists_stream_id": device_lists_stream_id,
-                "join_event_id": join_event_id,
+                # To be updated later once the join event is persisted.
+                "join_event_id": None,
             },
         )
         DatabasePool.simple_insert_many_txn(
@@ -1828,6 +1827,36 @@ def _store_partial_state_room_txn(
         )
         self._invalidate_cache_and_stream(txn, self.is_partial_state_room, (room_id,))
 
+    async def write_partial_state_rooms_join_event_id(
+        self,
+        room_id: str,
+        join_event_id: str,
+    ) -> None:
+        """Record the join event which resulted from a partial join.
+
+        We do this separately to `store_partial_state_room` because we need to wait for
+        the join event to be persisted. Otherwise we violate a foreign key constraint.
+        """
+        await self.db_pool.runInteraction(
+            "write_partial_state_rooms_join_event_id",
+            self._store_partial_state_room_txn,
+            room_id,
+            join_event_id,
+        )
+
+    async def _write_partial_state_rooms_join_event_id(
+        self,
+        txn: LoggingTransaction,
+        room_id: str,
+        join_event_id: str,
+    ) -> None:
+        DatabasePool.simple_update_txn(
+            txn,
+            table="partial_state_rooms",
+            keyvalues={"room_id": room_id},
+            updatevalues={"join_event_id": join_event_id},
+        )
+
     async def maybe_store_room_on_outlier_membership(
         self, room_id: str, room_version: RoomVersion
     ) -> None:

synapse/storage/schema/main/delta/73/04partial_join_details.sql

@@ -20,4 +20,4 @@
 --
 -- Both columns are backwards compatible.
 ALTER TABLE partial_state_rooms ADD COLUMN device_lists_stream_id BIGINT NOT NULL DEFAULT 0;
-ALTER TABLE partial_state_rooms ADD COLUMN join_event_id TEXT;
+ALTER TABLE partial_state_rooms ADD COLUMN join_event_id TEXT REFERENCES events(event_id);