Make _get_state_group_for_events raise if state is unknown

It seems like calling `_get_state_group_for_events` for an event where the
state is unknown is an error. Accordingly, let's raise an exception rather than
silently returning an empty result.

Author: richvdh

synapse/storage/databases/main/state.py

@@ -14,10 +14,10 @@
 # limitations under the License.
 import collections.abc
 import logging
-from typing import TYPE_CHECKING, Iterable, Optional, Set
+from typing import TYPE_CHECKING, Collection, Dict, Iterable, Optional, Set
 
 from synapse.api.constants import EventTypes, Membership
-from synapse.api.errors import NotFoundError, UnsupportedRoomVersionError
+from synapse.api.errors import NotFoundError, StoreError, UnsupportedRoomVersionError
 from synapse.api.room_versions import KNOWN_ROOM_VERSIONS, RoomVersion
 from synapse.events import EventBase
 from synapse.storage._base import SQLBaseStore
@@ -304,8 +304,13 @@ async def _get_state_group_for_event(self, event_id: str) -> Optional[int]:
         list_name="event_ids",
         num_args=1,
     )
-    async def _get_state_group_for_events(self, event_ids):
-        """Returns mapping event_id -> state_group"""
+    async def _get_state_group_for_events(
+        self, event_ids: Collection[str]
+    ) -> Dict[str, int]:
+        """Returns mapping event_id -> state_group.
+
+        Raises a StoreError if the state is unknown at any of the given events
+        """
         rows = await self.db_pool.simple_select_many_batch(
             table="event_to_state_groups",
             column="event_id",
@@ -315,7 +320,13 @@ async def _get_state_group_for_events(self, event_ids):
             desc="_get_state_group_for_events",
         )
 
-        return {row["event_id"]: row["state_group"] for row in rows}
+        res = {row["event_id"]: row["state_group"] for row in rows}
+        for e in event_ids:
+            if e not in res:
+                raise StoreError(
+                    404, "No state group for unknown or outlier event %s" % e
+                )
+        return res
 
     async def get_referenced_state_groups(
         self, state_groups: Iterable[int]

synapse/storage/state.py

@@ -571,6 +571,10 @@ async def get_state_groups_ids(
 
         Returns:
             dict of state_group_id -> (dict of (type, state_key) -> event id)
+
+        Raises:
+            StoreError if we don't have a state group for any of the events (ie they
+               are outliers or unknown)
         """
         if not event_ids:
             return {}
@@ -659,6 +663,10 @@ async def get_state_for_events(
 
         Returns:
             A dict of (event_id) -> (type, state_key) -> [state_events]
+
+        Raises:
+            StoreError if we don't have a state group for any of the events (ie they
+               are outliers or unknown)
         """
         event_to_groups = await self.stores.main._get_state_group_for_events(event_ids)
 
@@ -696,6 +704,10 @@ async def get_state_ids_for_events(
 
         Returns:
             A dict from event_id -> (type, state_key) -> event_id
+
+        Raises:
+            StoreError if we don't have a state group for any of the events (ie they
+               are outliers or unknown)
         """
         event_to_groups = await self.stores.main._get_state_group_for_events(event_ids)
 
@@ -723,6 +735,10 @@ async def get_state_for_event(
 
         Returns:
             A dict from (type, state_key) -> state_event
+
+        Raises:
+            StoreError if we don't have a state group for the event (ie it is an
+               outlier or is unknown)
         """
         state_map = await self.get_state_for_events(
             [event_id], state_filter or StateFilter.all()