diff --git a/synapse/storage/roommember.py b/synapse/storage/roommember.py
index a422ddf633..3ffad672a7 100644
--- a/synapse/storage/roommember.py
+++ b/synapse/storage/roommember.py
@@ -20,7 +20,7 @@ from collections import namedtuple
from ._base import SQLBaseStore
from synapse.util.caches.descriptors import cached, cachedInlineCallbacks
-from synapse.api.constants import Membership
+from synapse.api.constants import Membership, EventTypes
from synapse.types import get_domain_from_id
import logging
@@ -325,7 +325,8 @@ class RoomMemberStore(SQLBaseStore):
@cachedInlineCallbacks(num_args=3)
def was_forgotten_at(self, user_id, room_id, event_id):
- """Returns whether user_id has elected to discard history for room_id at event_id.
+ """Returns whether user_id has elected to discard history for room_id at
+ event_id.
event_id must be a membership event."""
def f(txn):
@@ -358,3 +359,43 @@ class RoomMemberStore(SQLBaseStore):
},
desc="who_forgot"
)
+
+ def get_joined_users_from_context(self, room_id, context):
+ state_group = context.state_group
+ if not state_group:
+ # If state_group is None it means it has yet to be assigned a
+ # state group, i.e. we need to make sure that calls with a state_group
+ # of None don't hit previous cached calls with a None state_group.
+ # To do this we set the state_group to a new object as object() != object()
+ state_group = object()
+
+ return self._get_joined_users_from_context(
+ room_id, state_group, context.current_state_ids
+ )
+
+ @cachedInlineCallbacks(num_args=2, cache_context=True)
+ def _get_joined_users_from_context(self, room_id, state_group, current_state_ids,
+ cache_context):
+ # We don't use `state_group`, its there so that we can cache based
+ # on it. However, its important that its never None, since two current_state's
+ # with a state_group of None are likely to be different.
+ # See bulk_get_push_rules_for_room for how we work around this.
+ assert state_group is not None
+
+ member_event_ids = [
+ e_id
+ for key, e_id in current_state_ids.iteritems()
+ if key[0] == EventTypes.Member
+ ]
+
+ rows = yield self._simple_select_many_batch(
+ table="room_memberships",
+ column="event_id",
+ iterable=member_event_ids,
+ retcols=['user_id'],
+ keyvalues={
+ "membership": Membership.JOIN,
+ }
+ )
+
+ defer.returnValue(set(row["user_id"] for row in rows))
|
