From 1f5dec35ffbd5af0995674aa3d5179fb4fec53e2 Mon Sep 17 00:00:00 2001 From: Johannes Marbach Date: Thu, 15 Oct 2026 15:52:38 +0200 Subject: [PATCH] Add a query for which rooms have sticky events in a stream range --- .../storage/databases/main/sticky_events.py | 49 +++++++++++++ tests/storage/test_sticky_events.py | 68 +++++++++++++++++++ 2 files changed, 117 insertions(+) diff --git a/synapse/storage/databases/main/sticky_events.py b/synapse/storage/databases/main/sticky_events.py index 4f95736709..5830bf2b33 100644 --- a/synapse/storage/databases/main/sticky_events.py +++ b/synapse/storage/databases/main/sticky_events.py @@ -273,6 +273,55 @@ class StickyEventsWorkerStore(StateGroupWorkerStore, CacheInvalidationWorkerStor ) return cast(list[tuple[int, str, str]], txn.fetchall()) + async def get_rooms_with_sticky_events_between( + self, + room_ids: Collection[str], + *, + from_token: MultiWriterStreamToken, + to_token: MultiWriterStreamToken, + ) -> set[str]: + """Given a set of room IDs, find out which of them have sticky events + (unexpired or not) with sticky stream positions after `from_token` and + at or before `to_token`. + + Args: + room_ids: The room IDs to check. + from_token: The sticky stream position to check from (exclusive). + to_token: The sticky stream position to check to (inclusive). + Returns: + The subset of `room_ids` that have sticky events in the range. + """ + if len(room_ids) == 0 or to_token.is_before_or_eq(from_token): + return set() + + def _get_rooms_with_sticky_events_between_txn( + txn: LoggingTransaction, + ) -> set[str]: + room_id_in_list_clause, room_id_in_list_values = make_in_list_sql_clause( + txn.database_engine, "room_id", room_ids + ) + stream_clause, stream_values = make_multiwriter_sharded_token_bounds_sql( + self.database_engine, + stream_id_column="stream_id", + instance_name_column="instance_name", + from_token_exclusive=from_token, + to_token_inclusive=to_token, + ) + txn.execute( + f""" + SELECT DISTINCT room_id + FROM sticky_events + WHERE {stream_clause} AND {room_id_in_list_clause} + """, + (*stream_values, *room_id_in_list_values), + ) + return {room_id for (room_id,) in txn} + + return await self.db_pool.runInteraction( + "get_rooms_with_sticky_events_between", + _get_rooms_with_sticky_events_between_txn, + ) + async def get_updated_sticky_events( self, *, from_id: int, to_id: int, limit: int ) -> list[StickyEventUpdate]: diff --git a/tests/storage/test_sticky_events.py b/tests/storage/test_sticky_events.py index e81ce129d4..9c8a7937b4 100644 --- a/tests/storage/test_sticky_events.py +++ b/tests/storage/test_sticky_events.py @@ -299,6 +299,74 @@ class StickyEventsTestCase(unittest.HomeserverTestCase): ), ) + def test_get_rooms_with_sticky_events_between(self) -> None: + """ + Tests `get_rooms_with_sticky_events_between` returns only the rooms with + sticky events in the given stream range. + """ + other_room_id = self.helper.create_room_as(self.user_id, tok=self.token) + room_ids = [self.room_id, other_room_id] + + start_token = self.store.get_sticky_events_stream_token() + + self.helper.send_sticky_event( + self.room_id, + EventTypes.Message, + duration=Duration(minutes=1), + content={"body": "message 1", "msgtype": "m.text"}, + tok=self.token, + ) + mid_token = self.store.get_sticky_events_stream_token() + + self.helper.send_sticky_event( + other_room_id, + EventTypes.Message, + duration=Duration(minutes=1), + content={"body": "message 2", "msgtype": "m.text"}, + tok=self.token, + ) + end_token = self.store.get_sticky_events_stream_token() + + # Both rooms changed over the whole range + self.assertEqual( + self.get_success( + self.store.get_rooms_with_sticky_events_between( + room_ids, from_token=start_token, to_token=end_token + ) + ), + {self.room_id, other_room_id}, + ) + + # Only the other room changed in the second half of the range + self.assertEqual( + self.get_success( + self.store.get_rooms_with_sticky_events_between( + room_ids, from_token=mid_token, to_token=end_token + ) + ), + {other_room_id}, + ) + + # Nothing changed in an empty range + self.assertEqual( + self.get_success( + self.store.get_rooms_with_sticky_events_between( + room_ids, from_token=end_token, to_token=end_token + ) + ), + set(), + ) + + # Rooms we didn't ask about are not returned + self.assertEqual( + self.get_success( + self.store.get_rooms_with_sticky_events_between( + [self.room_id], from_token=start_token, to_token=end_token + ) + ), + {self.room_id}, + ) + def test_outlier_events_not_in_table(self) -> None: """ Tests the behaviour of outliered and then de-outliered events in the