From 8461faf3845a42c351bc1e76a56f65a271acfcbe Mon Sep 17 00:00:00 2001 From: Eric Eastwood Date: Thu, 15 Aug 2024 21:56:12 -0500 Subject: [PATCH] Add historical case to background update --- synapse/storage/databases/main/events.py | 10 ++-- synapse/storage/databases/state/bg_updates.py | 51 ++++++++++++------- 2 files changed, 39 insertions(+), 22 deletions(-) diff --git a/synapse/storage/databases/main/events.py b/synapse/storage/databases/main/events.py index 23dd9f2cce..21540d0279 100644 --- a/synapse/storage/databases/main/events.py +++ b/synapse/storage/databases/main/events.py @@ -1244,7 +1244,7 @@ class PersistEventsStore: ] = {} if current_state_map: sliding_sync_membership_snapshots_insert_map = ( - self._get_sliding_sync_insert_values_from_current_state_map_txn( + self._get_sliding_sync_insert_values_from_state_map_txn( txn, current_state_map ) ) @@ -1430,7 +1430,7 @@ class PersistEventsStore: # Map of values to insert/update in the `sliding_sync_joined_rooms` table sliding_sync_joined_rooms_insert_map = ( - self._get_sliding_sync_insert_values_from_current_state_map_txn( + self._get_sliding_sync_insert_values_from_state_map_txn( txn, current_state_map ) ) @@ -1593,8 +1593,8 @@ class PersistEventsStore: return current_state_map @classmethod - def _get_sliding_sync_insert_values_from_current_state_map_txn( - cls, txn: LoggingTransaction, current_state_map: StateMap[str] + def _get_sliding_sync_insert_values_from_state_map_txn( + cls, txn: LoggingTransaction, state_map: StateMap[str] ) -> Dict[str, Optional[Union[str, bool]]]: """ TODO @@ -1613,7 +1613,7 @@ class PersistEventsStore: ) = make_in_list_sql_clause( txn.database_engine, "event_id", - current_state_map.values(), + state_map.values(), ) txn.execute( f""" diff --git a/synapse/storage/databases/state/bg_updates.py b/synapse/storage/databases/state/bg_updates.py index 39c82807d2..9b7843aa3f 100644 --- a/synapse/storage/databases/state/bg_updates.py +++ b/synapse/storage/databases/state/bg_updates.py @@ -32,7 +32,10 @@ from synapse.storage.database import ( LoggingDatabaseConnection, LoggingTransaction, ) -from synapse.storage.databases.main.events import PersistEventsStore +from synapse.storage.databases.main.events import ( + SLIDING_SYNC_RELEVANT_STATE_SET, + PersistEventsStore, +) from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine from synapse.types import JsonDict, MutableStateMap, StateMap, StrCollection from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES @@ -618,7 +621,7 @@ class StateBackgroundUpdateStore(StateGroupBackgroundUpdateStore): # so we should have some current state for each room assert current_state_map - sliding_sync_joined_rooms_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_current_state_map_txn( + sliding_sync_joined_rooms_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_state_map_txn( txn, current_state_map ) # We should have some insert values for each room, even if they are `None` @@ -751,7 +754,7 @@ class StateBackgroundUpdateStore(StateGroupBackgroundUpdateStore): # for each room assert current_state_map - sliding_sync_membership_snapshots_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_current_state_map_txn( + sliding_sync_membership_snapshots_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_state_map_txn( txn, current_state_map ) # We should have some insert values for each room, even if they are `None` @@ -827,21 +830,35 @@ class StateBackgroundUpdateStore(StateGroupBackgroundUpdateStore): assert sliding_sync_membership_snapshots_insert_map elif membership == Membership.BAN: # Pull from historical state - # state_group = self.db_pool.simple_select_one_onecol_txn( - # table="event_to_state_groups", - # keyvalues={"event_id": membership_event_id}, - # retcol="state_group", - # allow_none=True, - # desc="_get_state_group_for_event", - # ) - # # We should know the state for the event - # assert state_group is not None + state_group = self.db_pool.simple_select_one_onecol_txn( + txn, + table="event_to_state_groups", + keyvalues={"event_id": membership_event_id}, + retcol="state_group", + allow_none=True, + ) + # We should know the state for the event + assert state_group is not None - # state_by_group = self._get_state_groups_from_groups_txn( - # txn, [state_group] - # ) - # state_map = state_by_group[state_group] - pass + state_by_group = self._get_state_groups_from_groups_txn( + txn, + groups=[state_group], + state_filter=StateFilter.from_types( + SLIDING_SYNC_RELEVANT_STATE_SET + ), + ) + state_map = state_by_group[state_group] + + sliding_sync_membership_snapshots_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_state_map_txn( + txn, state_map + ) + # We should have some insert values for each room, even if they are `None` + assert sliding_sync_membership_snapshots_insert_map + + # We have historical state to work from + sliding_sync_membership_snapshots_insert_map["has_known_state"] = ( + True + ) else: assert_never(membership)