From 6c2fc1d20f830adca8c7ff1d80941283d9054ed1 Mon Sep 17 00:00:00 2001 From: Eric Eastwood Date: Thu, 15 Aug 2024 20:51:05 -0500 Subject: [PATCH] Move background updates to `StateBackgroundUpdateStore` So we can access `_get_state_groups_from_groups_txn(...)` --- synapse/storage/databases/main/events.py | 7 +- .../databases/main/events_bg_updates.py | 363 +---------------- synapse/storage/databases/state/bg_updates.py | 382 +++++++++++++++++- tests/storage/test_events.py | 2 +- 4 files changed, 386 insertions(+), 368 deletions(-) diff --git a/synapse/storage/databases/main/events.py b/synapse/storage/databases/main/events.py index 025576eb6e..23dd9f2cce 100644 --- a/synapse/storage/databases/main/events.py +++ b/synapse/storage/databases/main/events.py @@ -1194,7 +1194,7 @@ class PersistEventsStore: if ev_type == EventTypes.Member } - # We now update `sliding_sync_membership_snapshots`. + # Handle updating the `sliding_sync_membership_snapshots` table # # This would only happen if someone was state reset out of the room if to_delete: @@ -1674,7 +1674,7 @@ class PersistEventsStore: events_and_contexts: The events being persisted """ - # Handle updating `sliding_sync_joined_rooms` + # Handle updating the `sliding_sync_joined_rooms` table. room_id_to_stream_ordering_map: Dict[str, int] = {} room_id_to_bump_stamp_map: Dict[str, int] = {} for event, _ in events_and_contexts: @@ -2390,7 +2390,8 @@ class PersistEventsStore: }, ) - # Update the `sliding_sync_membership_snapshots` table + # Handle updating the `sliding_sync_membership_snapshots` table + # (out-of-band membership events only) # raw_stripped_state_events = None if event.membership == Membership.INVITE: diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py index cef354dd5b..64d303e330 100644 --- a/synapse/storage/databases/main/events_bg_updates.py +++ b/synapse/storage/databases/main/events_bg_updates.py @@ -20,12 +20,11 @@ # import logging -from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union, cast +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, cast import attr -from typing_extensions import assert_never -from synapse.api.constants import EventContentFields, Membership, RelationTypes +from synapse.api.constants import EventContentFields, RelationTypes from synapse.api.room_versions import KNOWN_ROOM_VERSIONS from synapse.events import make_event_from_dict from synapse.storage._base import SQLBaseStore, db_to_json, make_in_list_sql_clause @@ -36,10 +35,8 @@ from synapse.storage.database import ( make_tuple_comparison_clause, ) from synapse.storage.databases.main.events import PersistEventsStore -from synapse.storage.engines import BaseDatabaseEngine from synapse.storage.types import Cursor from synapse.types import JsonDict, StrCollection -from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES if TYPE_CHECKING: from synapse.server import HomeServer @@ -81,11 +78,6 @@ class _BackgroundUpdates: EVENTS_JUMP_TO_DATE_INDEX = "events_jump_to_date_index" - SLIDING_SYNC_JOINED_ROOMS_BACKFILL = "sliding_sync_joined_rooms_backfill" - SLIDING_SYNC_MEMBERSHIP_SNAPSHOTS_BACKFILL = ( - "sliding_sync_membership_snapshots_backfill" - ) - @attr.s(slots=True, frozen=True, auto_attribs=True) class _CalculateChainCover: @@ -287,16 +279,6 @@ class EventsBackgroundUpdatesStore(SQLBaseStore): where_clause="NOT outlier", ) - # Backfill the sliding sync tables - self.db_pool.updates.register_background_update_handler( - _BackgroundUpdates.SLIDING_SYNC_JOINED_ROOMS_BACKFILL, - self._sliding_sync_joined_rooms_backfill, - ) - self.db_pool.updates.register_background_update_handler( - _BackgroundUpdates.SLIDING_SYNC_MEMBERSHIP_SNAPSHOTS_BACKFILL, - self._sliding_sync_membership_snapshots_backfill, - ) - async def _background_reindex_fields_sender( self, progress: JsonDict, batch_size: int ) -> int: @@ -1534,344 +1516,3 @@ class EventsBackgroundUpdatesStore(SQLBaseStore): ) return batch_size - - async def _sliding_sync_joined_rooms_backfill( - self, progress: JsonDict, batch_size: int - ) -> int: - """ - Handles backfilling the `sliding_sync_joined_rooms` table. - """ - last_room_id = progress.get("last_room_id", "") - - def make_sql_clause_for_get_last_event_pos_in_room( - database_engine: BaseDatabaseEngine, - event_types: Optional[StrCollection] = None, - ) -> Tuple[str, list]: - """ - Returns the ID and event position of the last event in a room at or before a - stream ordering. - - Based on `get_last_event_pos_in_room_before_stream_ordering(...)` - - Args: - database_engine - event_types: Optional allowlist of event types to filter by - - Returns: - A tuple of SQL query and the args - """ - event_type_clause = "" - event_type_args: List[str] = [] - if event_types is not None and len(event_types) > 0: - event_type_clause, event_type_args = make_in_list_sql_clause( - database_engine, "type", event_types - ) - event_type_clause = f"AND {event_type_clause}" - - sql = f""" - SELECT stream_ordering - FROM events - LEFT JOIN rejections USING (event_id) - WHERE room_id = ? - {event_type_clause} - AND NOT outlier - AND rejections.event_id IS NULL - ORDER BY stream_ordering DESC - LIMIT 1 - """ - - return sql, event_type_args - - def _txn(txn: LoggingTransaction) -> int: - # Fetch the set of room IDs that we want to update - txn.execute( - """ - SELECT DISTINCT room_id FROM current_state_events - WHERE room_id > ? - ORDER BY room_id ASC - LIMIT ? - """, - (last_room_id, batch_size), - ) - - rooms_to_update_rows = txn.fetchall() - if not rooms_to_update_rows: - return 0 - - for (room_id,) in rooms_to_update_rows: - current_state_map = PersistEventsStore._get_relevant_sliding_sync_current_state_event_ids_txn( - txn, room_id - ) - # We're iterating over rooms pulled from the current_state_events table - # 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( - txn, current_state_map - ) - # We should have some insert values for each room, even if they are `None` - assert sliding_sync_joined_rooms_insert_map - - ( - most_recent_event_stream_ordering_clause, - most_recent_event_stream_ordering_args, - ) = make_sql_clause_for_get_last_event_pos_in_room( - txn.database_engine, event_types=None - ) - bump_stamp_clause, bump_stamp_args = ( - make_sql_clause_for_get_last_event_pos_in_room( - txn.database_engine, - event_types=SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES, - ) - ) - - # Pulling keys/values separately is safe and will produce congruent - # lists - insert_keys = sliding_sync_joined_rooms_insert_map.keys() - insert_values = sliding_sync_joined_rooms_insert_map.values() - - sql = f""" - INSERT INTO sliding_sync_joined_rooms - (room_id, event_stream_ordering, bump_stamp, {", ".join(insert_keys)}) - VALUES ( - ?, - ({most_recent_event_stream_ordering_clause}), - ({bump_stamp_clause}), - {", ".join("?" for _ in insert_values)} - ) - ON CONFLICT (room_id) - DO UPDATE SET - event_stream_ordering = EXCLUDED.event_stream_ordering, - bump_stamp = EXCLUDED.bump_stamp, - {", ".join(f"{key} = EXCLUDED.{key}" for key in insert_keys)} - """ - args = ( - [room_id, room_id] - + most_recent_event_stream_ordering_args - + [room_id] - + bump_stamp_args - + list(insert_values) - ) - txn.execute(sql, args) - - self.db_pool.updates._background_update_progress_txn( - txn, - _BackgroundUpdates.SLIDING_SYNC_JOINED_ROOMS_BACKFILL, - {"last_room_id": rooms_to_update_rows[-1][0]}, - ) - - return len(rooms_to_update_rows) - - count = await self.db_pool.runInteraction( - "sliding_sync_joined_rooms_backfill", _txn - ) - - if not count: - await self.db_pool.updates._end_background_update( - _BackgroundUpdates.SLIDING_SYNC_JOINED_ROOMS_BACKFILL - ) - - return count - - async def _sliding_sync_membership_snapshots_backfill( - self, progress: JsonDict, batch_size: int - ) -> int: - """ - Handles backfilling the `sliding_sync_membership_snapshots` table. - """ - last_event_stream_ordering = progress.get( - "last_event_stream_ordering", -(1 << 31) - ) - - def _txn(txn: LoggingTransaction) -> int: - # Fetch the set of event IDs that we want to update - txn.execute( - """ - SELECT - c.room_id, - c.user_id, - c.event_id, - c.membership, - c.event_stream_ordering, - e.outlier - FROM local_current_membership as c - INNER JOIN events AS e USING (event_id) - WHERE event_stream_ordering > ? - ORDER BY event_stream_ordering ASC - LIMIT ? - """, - (last_event_stream_ordering, batch_size), - ) - - memberships_to_update_rows = txn.fetchall() - if not memberships_to_update_rows: - return 0 - - for ( - room_id, - user_id, - membership_event_id, - membership, - _membership_event_stream_ordering, - is_outlier, - ) in memberships_to_update_rows: - # We don't know how to handle `membership` values other than these. The - # code below would need to be updated. - assert membership in ( - Membership.JOIN, - Membership.INVITE, - Membership.KNOCK, - Membership.LEAVE, - Membership.BAN, - ) - - # Map of values to insert/update in the `sliding_sync_membership_snapshots` table - sliding_sync_membership_snapshots_insert_map: Dict[ - str, Optional[Union[str, bool]] - ] = {} - if membership == Membership.JOIN: - # If we're still joined, we can pull from current state - current_state_map = PersistEventsStore._get_relevant_sliding_sync_current_state_event_ids_txn( - txn, room_id - ) - # We're iterating over rooms that we are joined to so they should - # have `current_state_events` and we should have some current state - # for each room - assert current_state_map - - sliding_sync_membership_snapshots_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_current_state_map_txn( - txn, current_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 current state to work from - sliding_sync_membership_snapshots_insert_map["has_known_state"] = ( - True - ) - elif membership in (Membership.INVITE, Membership.KNOCK) or ( - membership == Membership.LEAVE and is_outlier - ): - invite_or_knock_event_id = membership_event_id - invite_or_knock_membership = membership - - # If the event is an `out_of_band_membership` (special case of - # `outlier`), we never had historical state so we have to pull from - # the stripped state on the previous invite/knock event. This gives - # us a consistent view of the room state regardless of your - # membership (i.e. the room shouldn't disappear if your using the - # `is_encrypted` filter and you leave). - if membership == Membership.LEAVE and is_outlier: - # Find the previous invite/knock event before the leave event - txn.execute( - """ - SELECT event_id, membership - FROM room_memberships - WHERE - room_id = ? - AND user_id = ? - AND event_stream_ordering < ? - ORDER BY event_stream_ordering DESC - LIMIT 1 - """ - ) - row = txn.fetchone() - # We should see a corresponding previous invite/knock event - assert row is not None - invite_or_knock_event_id, invite_or_knock_membership = row - - # Pull from the stripped state on the invite/knock event - txn.execute( - """ - SELECT json FROM event_json - WHERE event_id = ? - """, - (invite_or_knock_event_id), - ) - row = txn.fetchone() - # We should find a corresponding event - assert row is not None - json = row[0] - event_json = db_to_json(json) - - raw_stripped_state_events = None - if invite_or_knock_membership == Membership.INVITE: - invite_room_state = event_json.get("unsigned").get( - "invite_room_state" - ) - raw_stripped_state_events = invite_room_state - elif invite_or_knock_membership == Membership.KNOCK: - knock_room_state = event_json.get("unsigned").get( - "knock_room_state" - ) - raw_stripped_state_events = knock_room_state - - sliding_sync_membership_snapshots_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_stripped_state_txn( - txn, raw_stripped_state_events - ) - # We should have some insert values for each room, even if no - # stripped state is on the event because we still want to record - # that we have no known state - assert sliding_sync_membership_snapshots_insert_map - elif membership == Membership.BAN: - # Pull from historical state - # TODO - pass - else: - assert_never(membership) - - # Pulling keys/values separately is safe and will produce congruent - # lists - insert_keys = sliding_sync_membership_snapshots_insert_map.keys() - insert_values = sliding_sync_membership_snapshots_insert_map.values() - # We don't need to do anything `ON CONFLICT` because we never partially - # insert/update the snapshots - txn.execute( - f""" - INSERT INTO sliding_sync_membership_snapshots - (room_id, user_id, membership_event_id, membership, event_stream_ordering - {("," + ", ".join(insert_keys)) if insert_keys else ""}) - VALUES ( - ?, ?, ?, ?, - (SELECT stream_ordering FROM events WHERE event_id = ?) - {("," + ", ".join("?" for _ in insert_values)) if insert_values else ""} - ) - ON CONFLICT (room_id, user_id) - DO NOTHING - """, - [ - room_id, - user_id, - membership_event_id, - membership, - membership_event_id, - ] - + list(insert_values), - ) - - ( - _room_id, - _user_id, - _membership_event_id, - _membership, - membership_event_stream_ordering, - _is_outlier, - ) = memberships_to_update_rows[-1] - self.db_pool.updates._background_update_progress_txn( - txn, - _BackgroundUpdates.SLIDING_SYNC_MEMBERSHIP_SNAPSHOTS_BACKFILL, - {"last_event_stream_ordering": membership_event_stream_ordering}, - ) - - return len(memberships_to_update_rows) - - count = await self.db_pool.runInteraction( - "sliding_sync_membership_snapshots_backfill", _txn - ) - - if not count: - await self.db_pool.updates._end_background_update( - _BackgroundUpdates.SLIDING_SYNC_MEMBERSHIP_SNAPSHOTS_BACKFILL - ) - - return count diff --git a/synapse/storage/databases/state/bg_updates.py b/synapse/storage/databases/state/bg_updates.py index ea7d8199a7..39c82807d2 100644 --- a/synapse/storage/databases/state/bg_updates.py +++ b/synapse/storage/databases/state/bg_updates.py @@ -22,15 +22,20 @@ import logging from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Tuple, Union +from typing_extensions import assert_never + +from synapse.api.constants import Membership from synapse.logging.opentracing import tag_args, trace -from synapse.storage._base import SQLBaseStore +from synapse.storage._base import SQLBaseStore, db_to_json, make_in_list_sql_clause from synapse.storage.database import ( DatabasePool, LoggingDatabaseConnection, LoggingTransaction, ) -from synapse.storage.engines import PostgresEngine -from synapse.types import MutableStateMap, StateMap +from synapse.storage.databases.main.events import 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 from synapse.types.state import StateFilter from synapse.util.caches import intern_string @@ -43,6 +48,13 @@ logger = logging.getLogger(__name__) MAX_STATE_DELTA_HOPS = 100 +class _BackgroundUpdates: + SLIDING_SYNC_JOINED_ROOMS_BACKFILL = "sliding_sync_joined_rooms_backfill" + SLIDING_SYNC_MEMBERSHIP_SNAPSHOTS_BACKFILL = ( + "sliding_sync_membership_snapshots_backfill" + ) + + class StateGroupBackgroundUpdateStore(SQLBaseStore): """Defines functions related to state groups needed to run the state background updates. @@ -349,6 +361,16 @@ class StateBackgroundUpdateStore(StateGroupBackgroundUpdateStore): columns=["event_stream_ordering"], ) + # Backfill the sliding sync tables + self.db_pool.updates.register_background_update_handler( + _BackgroundUpdates.SLIDING_SYNC_JOINED_ROOMS_BACKFILL, + self._sliding_sync_joined_rooms_backfill, + ) + self.db_pool.updates.register_background_update_handler( + _BackgroundUpdates.SLIDING_SYNC_MEMBERSHIP_SNAPSHOTS_BACKFILL, + self._sliding_sync_membership_snapshots_backfill, + ) + async def _background_deduplicate_state( self, progress: dict, batch_size: int ) -> int: @@ -524,3 +546,357 @@ class StateBackgroundUpdateStore(StateGroupBackgroundUpdateStore): ) return 1 + + async def _sliding_sync_joined_rooms_backfill( + self, progress: JsonDict, batch_size: int + ) -> int: + """ + Handles backfilling the `sliding_sync_joined_rooms` table. + """ + last_room_id = progress.get("last_room_id", "") + + def make_sql_clause_for_get_last_event_pos_in_room( + database_engine: BaseDatabaseEngine, + event_types: Optional[StrCollection] = None, + ) -> Tuple[str, list]: + """ + Returns the ID and event position of the last event in a room at or before a + stream ordering. + + Based on `get_last_event_pos_in_room_before_stream_ordering(...)` + + Args: + database_engine + event_types: Optional allowlist of event types to filter by + + Returns: + A tuple of SQL query and the args + """ + event_type_clause = "" + event_type_args: List[str] = [] + if event_types is not None and len(event_types) > 0: + event_type_clause, event_type_args = make_in_list_sql_clause( + database_engine, "type", event_types + ) + event_type_clause = f"AND {event_type_clause}" + + sql = f""" + SELECT stream_ordering + FROM events + LEFT JOIN rejections USING (event_id) + WHERE room_id = ? + {event_type_clause} + AND NOT outlier + AND rejections.event_id IS NULL + ORDER BY stream_ordering DESC + LIMIT 1 + """ + + return sql, event_type_args + + def _txn(txn: LoggingTransaction) -> int: + # Fetch the set of room IDs that we want to update + txn.execute( + """ + SELECT DISTINCT room_id FROM current_state_events + WHERE room_id > ? + ORDER BY room_id ASC + LIMIT ? + """, + (last_room_id, batch_size), + ) + + rooms_to_update_rows = txn.fetchall() + if not rooms_to_update_rows: + return 0 + + for (room_id,) in rooms_to_update_rows: + current_state_map = PersistEventsStore._get_relevant_sliding_sync_current_state_event_ids_txn( + txn, room_id + ) + # We're iterating over rooms pulled from the current_state_events table + # 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( + txn, current_state_map + ) + # We should have some insert values for each room, even if they are `None` + assert sliding_sync_joined_rooms_insert_map + + ( + most_recent_event_stream_ordering_clause, + most_recent_event_stream_ordering_args, + ) = make_sql_clause_for_get_last_event_pos_in_room( + txn.database_engine, event_types=None + ) + bump_stamp_clause, bump_stamp_args = ( + make_sql_clause_for_get_last_event_pos_in_room( + txn.database_engine, + event_types=SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES, + ) + ) + + # Pulling keys/values separately is safe and will produce congruent + # lists + insert_keys = sliding_sync_joined_rooms_insert_map.keys() + insert_values = sliding_sync_joined_rooms_insert_map.values() + + sql = f""" + INSERT INTO sliding_sync_joined_rooms + (room_id, event_stream_ordering, bump_stamp, {", ".join(insert_keys)}) + VALUES ( + ?, + ({most_recent_event_stream_ordering_clause}), + ({bump_stamp_clause}), + {", ".join("?" for _ in insert_values)} + ) + ON CONFLICT (room_id) + DO UPDATE SET + event_stream_ordering = EXCLUDED.event_stream_ordering, + bump_stamp = EXCLUDED.bump_stamp, + {", ".join(f"{key} = EXCLUDED.{key}" for key in insert_keys)} + """ + args = ( + [room_id, room_id] + + most_recent_event_stream_ordering_args + + [room_id] + + bump_stamp_args + + list(insert_values) + ) + txn.execute(sql, args) + + self.db_pool.updates._background_update_progress_txn( + txn, + _BackgroundUpdates.SLIDING_SYNC_JOINED_ROOMS_BACKFILL, + {"last_room_id": rooms_to_update_rows[-1][0]}, + ) + + return len(rooms_to_update_rows) + + count = await self.db_pool.runInteraction( + "sliding_sync_joined_rooms_backfill", _txn + ) + + if not count: + await self.db_pool.updates._end_background_update( + _BackgroundUpdates.SLIDING_SYNC_JOINED_ROOMS_BACKFILL + ) + + return count + + async def _sliding_sync_membership_snapshots_backfill( + self, progress: JsonDict, batch_size: int + ) -> int: + """ + Handles backfilling the `sliding_sync_membership_snapshots` table. + """ + last_event_stream_ordering = progress.get( + "last_event_stream_ordering", -(1 << 31) + ) + + def _txn(txn: LoggingTransaction) -> int: + # Fetch the set of event IDs that we want to update + txn.execute( + """ + SELECT + c.room_id, + c.user_id, + c.event_id, + c.membership, + c.event_stream_ordering, + e.outlier + FROM local_current_membership as c + INNER JOIN events AS e USING (event_id) + WHERE event_stream_ordering > ? + ORDER BY event_stream_ordering ASC + LIMIT ? + """, + (last_event_stream_ordering, batch_size), + ) + + memberships_to_update_rows = txn.fetchall() + if not memberships_to_update_rows: + return 0 + + for ( + room_id, + user_id, + membership_event_id, + membership, + _membership_event_stream_ordering, + is_outlier, + ) in memberships_to_update_rows: + # We don't know how to handle `membership` values other than these. The + # code below would need to be updated. + assert membership in ( + Membership.JOIN, + Membership.INVITE, + Membership.KNOCK, + Membership.LEAVE, + Membership.BAN, + ) + + # Map of values to insert/update in the `sliding_sync_membership_snapshots` table + sliding_sync_membership_snapshots_insert_map: Dict[ + str, Optional[Union[str, bool]] + ] = {} + if membership == Membership.JOIN: + # If we're still joined, we can pull from current state + current_state_map = PersistEventsStore._get_relevant_sliding_sync_current_state_event_ids_txn( + txn, room_id + ) + # We're iterating over rooms that we are joined to so they should + # have `current_state_events` and we should have some current state + # for each room + assert current_state_map + + sliding_sync_membership_snapshots_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_current_state_map_txn( + txn, current_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 current state to work from + sliding_sync_membership_snapshots_insert_map["has_known_state"] = ( + True + ) + elif membership in (Membership.INVITE, Membership.KNOCK) or ( + membership == Membership.LEAVE and is_outlier + ): + invite_or_knock_event_id = membership_event_id + invite_or_knock_membership = membership + + # If the event is an `out_of_band_membership` (special case of + # `outlier`), we never had historical state so we have to pull from + # the stripped state on the previous invite/knock event. This gives + # us a consistent view of the room state regardless of your + # membership (i.e. the room shouldn't disappear if your using the + # `is_encrypted` filter and you leave). + if membership == Membership.LEAVE and is_outlier: + # Find the previous invite/knock event before the leave event + txn.execute( + """ + SELECT event_id, membership + FROM room_memberships + WHERE + room_id = ? + AND user_id = ? + AND event_stream_ordering < ? + ORDER BY event_stream_ordering DESC + LIMIT 1 + """ + ) + row = txn.fetchone() + # We should see a corresponding previous invite/knock event + assert row is not None + invite_or_knock_event_id, invite_or_knock_membership = row + + # Pull from the stripped state on the invite/knock event + txn.execute( + """ + SELECT json FROM event_json + WHERE event_id = ? + """, + (invite_or_knock_event_id), + ) + row = txn.fetchone() + # We should find a corresponding event + assert row is not None + json = row[0] + event_json = db_to_json(json) + + raw_stripped_state_events = None + if invite_or_knock_membership == Membership.INVITE: + invite_room_state = event_json.get("unsigned").get( + "invite_room_state" + ) + raw_stripped_state_events = invite_room_state + elif invite_or_knock_membership == Membership.KNOCK: + knock_room_state = event_json.get("unsigned").get( + "knock_room_state" + ) + raw_stripped_state_events = knock_room_state + + sliding_sync_membership_snapshots_insert_map = PersistEventsStore._get_sliding_sync_insert_values_from_stripped_state_txn( + txn, raw_stripped_state_events + ) + # We should have some insert values for each room, even if no + # stripped state is on the event because we still want to record + # that we have no known state + 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_by_group = self._get_state_groups_from_groups_txn( + # txn, [state_group] + # ) + # state_map = state_by_group[state_group] + pass + else: + assert_never(membership) + + # Pulling keys/values separately is safe and will produce congruent + # lists + insert_keys = sliding_sync_membership_snapshots_insert_map.keys() + insert_values = sliding_sync_membership_snapshots_insert_map.values() + # We don't need to do anything `ON CONFLICT` because we never partially + # insert/update the snapshots + txn.execute( + f""" + INSERT INTO sliding_sync_membership_snapshots + (room_id, user_id, membership_event_id, membership, event_stream_ordering + {("," + ", ".join(insert_keys)) if insert_keys else ""}) + VALUES ( + ?, ?, ?, ?, + (SELECT stream_ordering FROM events WHERE event_id = ?) + {("," + ", ".join("?" for _ in insert_values)) if insert_values else ""} + ) + ON CONFLICT (room_id, user_id) + DO NOTHING + """, + [ + room_id, + user_id, + membership_event_id, + membership, + membership_event_id, + ] + + list(insert_values), + ) + + ( + _room_id, + _user_id, + _membership_event_id, + _membership, + membership_event_stream_ordering, + _is_outlier, + ) = memberships_to_update_rows[-1] + self.db_pool.updates._background_update_progress_txn( + txn, + _BackgroundUpdates.SLIDING_SYNC_MEMBERSHIP_SNAPSHOTS_BACKFILL, + {"last_event_stream_ordering": membership_event_stream_ordering}, + ) + + return len(memberships_to_update_rows) + + count = await self.db_pool.runInteraction( + "sliding_sync_membership_snapshots_backfill", _txn + ) + + if not count: + await self.db_pool.updates._end_background_update( + _BackgroundUpdates.SLIDING_SYNC_MEMBERSHIP_SNAPSHOTS_BACKFILL + ) + + return count diff --git a/tests/storage/test_events.py b/tests/storage/test_events.py index 8875a9364e..e6a7d4dba1 100644 --- a/tests/storage/test_events.py +++ b/tests/storage/test_events.py @@ -35,7 +35,7 @@ from synapse.federation.federation_base import event_from_pdu_json from synapse.rest import admin from synapse.rest.client import login, room from synapse.server import HomeServer -from synapse.storage.databases.main.events_bg_updates import _BackgroundUpdates +from synapse.storage.databases.state.bg_updates import _BackgroundUpdates from synapse.types import StateMap from synapse.util import Clock