Add a query for which rooms have sticky events in a stream range

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]:

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