diff --git a/synapse/storage/controllers/state.py b/synapse/storage/controllers/state.py
index 9d7a8a792f..7089b0a1d8 100644
--- a/synapse/storage/controllers/state.py
+++ b/synapse/storage/controllers/state.py
@@ -16,7 +16,6 @@ from typing import (
TYPE_CHECKING,
AbstractSet,
Any,
- Awaitable,
Callable,
Collection,
Dict,
@@ -67,6 +66,8 @@ class StateStorageController:
"""
self._partial_state_room_tracker.notify_un_partial_stated(room_id)
+ @trace
+ @tag_args
async def get_state_group_delta(
self, state_group: int
) -> Tuple[Optional[int], Optional[StateMap[str]]]:
@@ -84,6 +85,8 @@ class StateStorageController:
state_group_delta = await self.stores.state.get_state_group_delta(state_group)
return state_group_delta.prev_group, state_group_delta.delta_ids
+ @trace
+ @tag_args
async def get_state_groups_ids(
self, _room_id: str, event_ids: Collection[str], await_full_state: bool = True
) -> Dict[int, MutableStateMap[str]]:
@@ -114,6 +117,8 @@ class StateStorageController:
return group_to_state
+ @trace
+ @tag_args
async def get_state_ids_for_group(
self, state_group: int, state_filter: Optional[StateFilter] = None
) -> StateMap[str]:
@@ -130,6 +135,8 @@ class StateStorageController:
return group_to_state[state_group]
+ @trace
+ @tag_args
async def get_state_groups(
self, room_id: str, event_ids: Collection[str]
) -> Dict[int, List[EventBase]]:
@@ -165,9 +172,11 @@ class StateStorageController:
for group, event_id_map in group_to_ids.items()
}
- def _get_state_groups_from_groups(
+ @trace
+ @tag_args
+ async def _get_state_groups_from_groups(
self, groups: List[int], state_filter: StateFilter
- ) -> Awaitable[Dict[int, StateMap[str]]]:
+ ) -> Dict[int, StateMap[str]]:
"""Returns the state groups for a given set of groups, filtering on
types of state events.
@@ -180,9 +189,12 @@ class StateStorageController:
Dict of state group to state map.
"""
- return self.stores.state._get_state_groups_from_groups(groups, state_filter)
+ return await self.stores.state._get_state_groups_from_groups(
+ groups, state_filter
+ )
@trace
+ @tag_args
async def get_state_for_events(
self, event_ids: Collection[str], state_filter: Optional[StateFilter] = None
) -> Dict[str, StateMap[EventBase]]:
@@ -280,6 +292,8 @@ class StateStorageController:
return {event: event_to_state[event] for event in event_ids}
+ @trace
+ @tag_args
async def get_state_for_event(
self, event_id: str, state_filter: Optional[StateFilter] = None
) -> StateMap[EventBase]:
@@ -303,6 +317,7 @@ class StateStorageController:
return state_map[event_id]
@trace
+ @tag_args
async def get_state_ids_for_event(
self,
event_id: str,
@@ -333,9 +348,11 @@ class StateStorageController:
)
return state_map[event_id]
- def get_state_for_groups(
+ @trace
+ @tag_args
+ async def get_state_for_groups(
self, groups: Iterable[int], state_filter: Optional[StateFilter] = None
- ) -> Awaitable[Dict[int, MutableStateMap[str]]]:
+ ) -> Dict[int, MutableStateMap[str]]:
"""Gets the state at each of a list of state groups, optionally
filtering by type/state_key
@@ -347,7 +364,7 @@ class StateStorageController:
Returns:
Dict of state group to state map.
"""
- return self.stores.state._get_state_for_groups(
+ return await self.stores.state._get_state_for_groups(
groups, state_filter or StateFilter.all()
)
@@ -402,6 +419,8 @@ class StateStorageController:
event_id, room_id, prev_group, delta_ids, current_state_ids
)
+ @trace
+ @tag_args
@cancellable
async def get_current_state_ids(
self,
@@ -442,6 +461,8 @@ class StateStorageController:
room_id, on_invalidate=on_invalidate
)
+ @trace
+ @tag_args
async def get_canonical_alias_for_room(self, room_id: str) -> Optional[str]:
"""Get canonical alias for room, if any
@@ -466,6 +487,8 @@ class StateStorageController:
return event.content.get("canonical_alias")
+ @trace
+ @tag_args
async def get_current_state_deltas(
self, prev_stream_id: int, max_stream_id: int
) -> Tuple[int, List[Dict[str, Any]]]:
@@ -500,6 +523,7 @@ class StateStorageController:
)
@trace
+ @tag_args
async def get_current_state(
self, room_id: str, state_filter: Optional[StateFilter] = None
) -> StateMap[EventBase]:
@@ -516,6 +540,8 @@ class StateStorageController:
return state_map
+ @trace
+ @tag_args
async def get_current_state_event(
self, room_id: str, event_type: str, state_key: str
) -> Optional[EventBase]:
@@ -527,6 +553,8 @@ class StateStorageController:
)
return state_map.get(key)
+ @trace
+ @tag_args
async def get_current_hosts_in_room(self, room_id: str) -> AbstractSet[str]:
"""Get current hosts in room based on current state.
@@ -538,6 +566,8 @@ class StateStorageController:
return await self.stores.main.get_current_hosts_in_room(room_id)
+ @trace
+ @tag_args
async def get_current_hosts_in_room_ordered(self, room_id: str) -> List[str]:
"""Get current hosts in room based on current state.
@@ -553,6 +583,8 @@ class StateStorageController:
return await self.stores.main.get_current_hosts_in_room_ordered(room_id)
+ @trace
+ @tag_args
async def get_current_hosts_in_room_or_partial_state_approximation(
self, room_id: str
) -> Collection[str]:
@@ -582,6 +614,8 @@ class StateStorageController:
return hosts
+ @trace
+ @tag_args
async def get_users_in_room_with_profiles(
self, room_id: str
) -> Mapping[str, ProfileInfo]:
|