You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 

1646 lines
60 KiB

  1. # Copyright 2014-2016 OpenMarket Ltd
  2. # Copyright 2018 New Vector Ltd
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. import logging
  16. from itertools import chain
  17. from typing import (
  18. TYPE_CHECKING,
  19. AbstractSet,
  20. Collection,
  21. Dict,
  22. FrozenSet,
  23. Iterable,
  24. List,
  25. Mapping,
  26. Optional,
  27. Sequence,
  28. Set,
  29. Tuple,
  30. Union,
  31. )
  32. import attr
  33. from synapse.api.constants import EventTypes, Membership
  34. from synapse.metrics import LaterGauge
  35. from synapse.metrics.background_process_metrics import wrap_as_background_process
  36. from synapse.storage._base import SQLBaseStore, db_to_json, make_in_list_sql_clause
  37. from synapse.storage.database import (
  38. DatabasePool,
  39. LoggingDatabaseConnection,
  40. LoggingTransaction,
  41. )
  42. from synapse.storage.databases.main.cache import CacheInvalidationWorkerStore
  43. from synapse.storage.databases.main.events_worker import EventsWorkerStore
  44. from synapse.storage.engines import Sqlite3Engine
  45. from synapse.storage.roommember import (
  46. GetRoomsForUserWithStreamOrdering,
  47. MemberSummary,
  48. ProfileInfo,
  49. RoomsForUser,
  50. )
  51. from synapse.types import (
  52. JsonDict,
  53. PersistedEventPosition,
  54. StateMap,
  55. StrCollection,
  56. get_domain_from_id,
  57. )
  58. from synapse.util.async_helpers import Linearizer
  59. from synapse.util.caches import intern_string
  60. from synapse.util.caches.descriptors import _CacheContext, cached, cachedList
  61. from synapse.util.iterutils import batch_iter
  62. from synapse.util.metrics import Measure
  63. if TYPE_CHECKING:
  64. from synapse.server import HomeServer
  65. from synapse.state import _StateCacheEntry
  66. logger = logging.getLogger(__name__)
  67. _MEMBERSHIP_PROFILE_UPDATE_NAME = "room_membership_profile_update"
  68. _CURRENT_STATE_MEMBERSHIP_UPDATE_NAME = "current_state_events_membership"
  69. @attr.s(frozen=True, slots=True, auto_attribs=True)
  70. class EventIdMembership:
  71. """Returned by `get_membership_from_event_ids`"""
  72. user_id: str
  73. membership: str
  74. class RoomMemberWorkerStore(EventsWorkerStore, CacheInvalidationWorkerStore):
  75. def __init__(
  76. self,
  77. database: DatabasePool,
  78. db_conn: LoggingDatabaseConnection,
  79. hs: "HomeServer",
  80. ):
  81. super().__init__(database, db_conn, hs)
  82. # Used by `_get_joined_hosts` to ensure only one thing mutates the cache
  83. # at a time. Keyed by room_id.
  84. self._joined_host_linearizer = Linearizer("_JoinedHostsCache")
  85. self._server_notices_mxid = hs.config.servernotices.server_notices_mxid
  86. if (
  87. self.hs.config.worker.run_background_tasks
  88. and self.hs.config.metrics.metrics_flags.known_servers
  89. ):
  90. self._known_servers_count = 1
  91. self.hs.get_clock().looping_call(
  92. self._count_known_servers,
  93. 60 * 1000,
  94. )
  95. self.hs.get_clock().call_later(
  96. 1,
  97. self._count_known_servers,
  98. )
  99. LaterGauge(
  100. "synapse_federation_known_servers",
  101. "",
  102. [],
  103. lambda: self._known_servers_count,
  104. )
  105. @wrap_as_background_process("_count_known_servers")
  106. async def _count_known_servers(self) -> int:
  107. """
  108. Count the servers that this server knows about.
  109. The statistic is stored on the class for the
  110. `synapse_federation_known_servers` LaterGauge to collect.
  111. """
  112. def _transact(txn: LoggingTransaction) -> int:
  113. if isinstance(self.database_engine, Sqlite3Engine):
  114. query = """
  115. SELECT COUNT(DISTINCT substr(out.user_id, pos+1))
  116. FROM (
  117. SELECT rm.user_id as user_id, instr(rm.user_id, ':')
  118. AS pos FROM room_memberships as rm
  119. INNER JOIN current_state_events as c ON rm.event_id = c.event_id
  120. WHERE c.type = 'm.room.member'
  121. ) as out
  122. """
  123. else:
  124. query = """
  125. SELECT COUNT(DISTINCT split_part(state_key, ':', 2))
  126. FROM current_state_events
  127. WHERE type = 'm.room.member' AND membership = 'join';
  128. """
  129. txn.execute(query)
  130. return list(txn)[0][0]
  131. count = await self.db_pool.runInteraction("get_known_servers", _transact)
  132. # We always know about ourselves, even if we have nothing in
  133. # room_memberships (for example, the server is new).
  134. self._known_servers_count = max([count, 1])
  135. return self._known_servers_count
  136. @cached(max_entries=100000, iterable=True)
  137. async def get_users_in_room(self, room_id: str) -> Sequence[str]:
  138. """Returns a list of users in the room.
  139. Will return inaccurate results for rooms with partial state, since the state for
  140. the forward extremities of those rooms will exclude most members. We may also
  141. calculate room state incorrectly for such rooms and believe that a member is or
  142. is not in the room when the opposite is true.
  143. Note: If you only care about users in the room local to the homeserver, use
  144. `get_local_users_in_room(...)` instead which will be more performant.
  145. """
  146. return await self.db_pool.simple_select_onecol(
  147. table="current_state_events",
  148. keyvalues={
  149. "type": EventTypes.Member,
  150. "room_id": room_id,
  151. "membership": Membership.JOIN,
  152. },
  153. retcol="state_key",
  154. desc="get_users_in_room",
  155. )
  156. def get_users_in_room_txn(self, txn: LoggingTransaction, room_id: str) -> List[str]:
  157. """Returns a list of users in the room."""
  158. return self.db_pool.simple_select_onecol_txn(
  159. txn,
  160. table="current_state_events",
  161. keyvalues={
  162. "type": EventTypes.Member,
  163. "room_id": room_id,
  164. "membership": Membership.JOIN,
  165. },
  166. retcol="state_key",
  167. )
  168. @cached()
  169. def get_user_in_room_with_profile(self, room_id: str, user_id: str) -> ProfileInfo:
  170. raise NotImplementedError()
  171. @cachedList(
  172. cached_method_name="get_user_in_room_with_profile", list_name="user_ids"
  173. )
  174. async def get_subset_users_in_room_with_profiles(
  175. self, room_id: str, user_ids: Collection[str]
  176. ) -> Dict[str, ProfileInfo]:
  177. """Get a mapping from user ID to profile information for a list of users
  178. in a given room.
  179. The profile information comes directly from this room's `m.room.member`
  180. events, and so may be specific to this room rather than part of a user's
  181. global profile. To avoid privacy leaks, the profile data should only be
  182. revealed to users who are already in this room.
  183. Args:
  184. room_id: The ID of the room to retrieve the users of.
  185. user_ids: a list of users in the room to run the query for
  186. Returns:
  187. A mapping from user ID to ProfileInfo.
  188. """
  189. def _get_subset_users_in_room_with_profiles(
  190. txn: LoggingTransaction,
  191. ) -> Dict[str, ProfileInfo]:
  192. clause, ids = make_in_list_sql_clause(
  193. self.database_engine, "c.state_key", user_ids
  194. )
  195. sql = """
  196. SELECT state_key, display_name, avatar_url FROM room_memberships as m
  197. INNER JOIN current_state_events as c
  198. ON m.event_id = c.event_id
  199. AND m.room_id = c.room_id
  200. AND m.user_id = c.state_key
  201. WHERE c.type = 'm.room.member' AND c.room_id = ? AND m.membership = ? AND %s
  202. """ % (
  203. clause,
  204. )
  205. txn.execute(sql, (room_id, Membership.JOIN, *ids))
  206. return {r[0]: ProfileInfo(display_name=r[1], avatar_url=r[2]) for r in txn}
  207. return await self.db_pool.runInteraction(
  208. "get_subset_users_in_room_with_profiles",
  209. _get_subset_users_in_room_with_profiles,
  210. )
  211. @cached(max_entries=100000, iterable=True)
  212. async def get_users_in_room_with_profiles(
  213. self, room_id: str
  214. ) -> Mapping[str, ProfileInfo]:
  215. """Get a mapping from user ID to profile information for all users in a given room.
  216. The profile information comes directly from this room's `m.room.member`
  217. events, and so may be specific to this room rather than part of a user's
  218. global profile. To avoid privacy leaks, the profile data should only be
  219. revealed to users who are already in this room.
  220. Args:
  221. room_id: The ID of the room to retrieve the users of.
  222. Returns:
  223. A mapping from user ID to ProfileInfo.
  224. Preconditions:
  225. - There is full state available for the room (it is not partial-stated).
  226. """
  227. def _get_users_in_room_with_profiles(
  228. txn: LoggingTransaction,
  229. ) -> Dict[str, ProfileInfo]:
  230. sql = """
  231. SELECT state_key, display_name, avatar_url FROM room_memberships as m
  232. INNER JOIN current_state_events as c
  233. ON m.event_id = c.event_id
  234. AND m.room_id = c.room_id
  235. AND m.user_id = c.state_key
  236. WHERE c.type = 'm.room.member' AND c.room_id = ? AND m.membership = ?
  237. """
  238. txn.execute(sql, (room_id, Membership.JOIN))
  239. return {r[0]: ProfileInfo(display_name=r[1], avatar_url=r[2]) for r in txn}
  240. return await self.db_pool.runInteraction(
  241. "get_users_in_room_with_profiles",
  242. _get_users_in_room_with_profiles,
  243. )
  244. @cached(max_entries=100000)
  245. async def get_room_summary(self, room_id: str) -> Mapping[str, MemberSummary]:
  246. """Get the details of a room roughly suitable for use by the room
  247. summary extension to /sync. Useful when lazy loading room members.
  248. Args:
  249. room_id: The room ID to query
  250. Returns:
  251. dict of membership states, pointing to a MemberSummary named tuple.
  252. """
  253. def _get_room_summary_txn(
  254. txn: LoggingTransaction,
  255. ) -> Dict[str, MemberSummary]:
  256. # first get counts.
  257. # We do this all in one transaction to keep the cache small.
  258. # FIXME: get rid of this when we have room_stats
  259. # Note, rejected events will have a null membership field, so
  260. # we we manually filter them out.
  261. sql = """
  262. SELECT count(*), membership FROM current_state_events
  263. WHERE type = 'm.room.member' AND room_id = ?
  264. AND membership IS NOT NULL
  265. GROUP BY membership
  266. """
  267. txn.execute(sql, (room_id,))
  268. res: Dict[str, MemberSummary] = {}
  269. for count, membership in txn:
  270. res.setdefault(membership, MemberSummary([], count))
  271. # we order by membership and then fairly arbitrarily by event_id so
  272. # heroes are consistent
  273. # Note, rejected events will have a null membership field, so
  274. # we we manually filter them out.
  275. sql = """
  276. SELECT state_key, membership, event_id
  277. FROM current_state_events
  278. WHERE type = 'm.room.member' AND room_id = ?
  279. AND membership IS NOT NULL
  280. ORDER BY
  281. CASE membership WHEN ? THEN 1 WHEN ? THEN 2 ELSE 3 END ASC,
  282. event_id ASC
  283. LIMIT ?
  284. """
  285. # 6 is 5 (number of heroes) plus 1, in case one of them is the calling user.
  286. txn.execute(sql, (room_id, Membership.JOIN, Membership.INVITE, 6))
  287. for user_id, membership, event_id in txn:
  288. summary = res[membership]
  289. # we will always have a summary for this membership type at this
  290. # point given the summary currently contains the counts.
  291. members = summary.members
  292. members.append((user_id, event_id))
  293. return res
  294. return await self.db_pool.runInteraction(
  295. "get_room_summary", _get_room_summary_txn
  296. )
  297. @cached()
  298. async def get_number_joined_users_in_room(self, room_id: str) -> int:
  299. return await self.db_pool.simple_select_one_onecol(
  300. table="current_state_events",
  301. keyvalues={"room_id": room_id, "membership": Membership.JOIN},
  302. retcol="COUNT(*)",
  303. desc="get_number_joined_users_in_room",
  304. )
  305. @cached()
  306. async def get_invited_rooms_for_local_user(
  307. self, user_id: str
  308. ) -> Sequence[RoomsForUser]:
  309. """Get all the rooms the *local* user is invited to.
  310. Args:
  311. user_id: The user ID.
  312. Returns:
  313. A list of RoomsForUser.
  314. """
  315. return await self.get_rooms_for_local_user_where_membership_is(
  316. user_id, [Membership.INVITE]
  317. )
  318. async def get_invite_for_local_user_in_room(
  319. self, user_id: str, room_id: str
  320. ) -> Optional[RoomsForUser]:
  321. """Gets the invite for the given *local* user and room.
  322. Args:
  323. user_id: The user ID to find the invite of.
  324. room_id: The room to user was invited to.
  325. Returns:
  326. Either a RoomsForUser or None if no invite was found.
  327. """
  328. invites = await self.get_invited_rooms_for_local_user(user_id)
  329. for invite in invites:
  330. if invite.room_id == room_id:
  331. return invite
  332. return None
  333. async def get_rooms_for_local_user_where_membership_is(
  334. self,
  335. user_id: str,
  336. membership_list: Collection[str],
  337. excluded_rooms: StrCollection = (),
  338. ) -> List[RoomsForUser]:
  339. """Get all the rooms for this *local* user where the membership for this user
  340. matches one in the membership list.
  341. Filters out forgotten rooms.
  342. Args:
  343. user_id: The user ID.
  344. membership_list: A list of synapse.api.constants.Membership
  345. values which the user must be in.
  346. excluded_rooms: A list of rooms to ignore.
  347. Returns:
  348. The RoomsForUser that the user matches the membership types.
  349. """
  350. if not membership_list:
  351. return []
  352. rooms = await self.db_pool.runInteraction(
  353. "get_rooms_for_local_user_where_membership_is",
  354. self._get_rooms_for_local_user_where_membership_is_txn,
  355. user_id,
  356. membership_list,
  357. )
  358. # Now we filter out forgotten and excluded rooms
  359. rooms_to_exclude: AbstractSet[str] = set()
  360. # Users can't forget joined/invited rooms, so we skip the check for such look ups.
  361. if any(m not in (Membership.JOIN, Membership.INVITE) for m in membership_list):
  362. rooms_to_exclude = await self.get_forgotten_rooms_for_user(user_id)
  363. if excluded_rooms is not None:
  364. # Take a copy to avoid mutating the in-cache set
  365. rooms_to_exclude = set(rooms_to_exclude)
  366. rooms_to_exclude.update(excluded_rooms)
  367. return [room for room in rooms if room.room_id not in rooms_to_exclude]
  368. def _get_rooms_for_local_user_where_membership_is_txn(
  369. self,
  370. txn: LoggingTransaction,
  371. user_id: str,
  372. membership_list: List[str],
  373. ) -> List[RoomsForUser]:
  374. """Get all the rooms for this *local* user where the membership for this user
  375. matches one in the membership list.
  376. Args:
  377. user_id: The user ID.
  378. membership_list: A list of synapse.api.constants.Membership
  379. values which the user must be in.
  380. Returns:
  381. The RoomsForUser that the user matches the membership types.
  382. """
  383. # Paranoia check.
  384. if not self.hs.is_mine_id(user_id):
  385. raise Exception(
  386. "Cannot call 'get_rooms_for_local_user_where_membership_is' on non-local user %r"
  387. % (user_id,),
  388. )
  389. clause, args = make_in_list_sql_clause(
  390. self.database_engine, "c.membership", membership_list
  391. )
  392. sql = """
  393. SELECT room_id, e.sender, c.membership, event_id, e.stream_ordering, r.room_version
  394. FROM local_current_membership AS c
  395. INNER JOIN events AS e USING (room_id, event_id)
  396. INNER JOIN rooms AS r USING (room_id)
  397. WHERE
  398. user_id = ?
  399. AND %s
  400. """ % (
  401. clause,
  402. )
  403. txn.execute(sql, (user_id, *args))
  404. results = [RoomsForUser(*r) for r in txn]
  405. return results
  406. @cached(iterable=True)
  407. async def get_local_users_in_room(self, room_id: str) -> Sequence[str]:
  408. """
  409. Retrieves a list of the current roommembers who are local to the server.
  410. """
  411. return await self.db_pool.simple_select_onecol(
  412. table="local_current_membership",
  413. keyvalues={"room_id": room_id, "membership": Membership.JOIN},
  414. retcol="user_id",
  415. desc="get_local_users_in_room",
  416. )
  417. async def check_local_user_in_room(self, user_id: str, room_id: str) -> bool:
  418. """
  419. Check whether a given local user is currently joined to the given room.
  420. Returns:
  421. A boolean indicating whether the user is currently joined to the room
  422. Raises:
  423. Exeption when called with a non-local user to this homeserver
  424. """
  425. if not self.hs.is_mine_id(user_id):
  426. raise Exception(
  427. "Cannot call 'check_local_user_in_room' on "
  428. "non-local user %s" % (user_id,),
  429. )
  430. (
  431. membership,
  432. member_event_id,
  433. ) = await self.get_local_current_membership_for_user_in_room(
  434. user_id=user_id,
  435. room_id=room_id,
  436. )
  437. return membership == Membership.JOIN
  438. async def is_server_notice_room(self, room_id: str) -> bool:
  439. """
  440. Determines whether the given room is a 'Server Notices' room, used for
  441. sending server notices to a user.
  442. This is determined by seeing whether the server notices user is present
  443. in the room.
  444. """
  445. if self._server_notices_mxid is None:
  446. return False
  447. is_server_notices_room = await self.check_local_user_in_room(
  448. user_id=self._server_notices_mxid, room_id=room_id
  449. )
  450. return is_server_notices_room
  451. async def get_local_current_membership_for_user_in_room(
  452. self, user_id: str, room_id: str
  453. ) -> Tuple[Optional[str], Optional[str]]:
  454. """Retrieve the current local membership state and event ID for a user in a room.
  455. Args:
  456. user_id: The ID of the user.
  457. room_id: The ID of the room.
  458. Returns:
  459. A tuple of (membership_type, event_id). Both will be None if a
  460. room_id/user_id pair is not found.
  461. """
  462. # Paranoia check.
  463. if not self.hs.is_mine_id(user_id):
  464. raise Exception(
  465. "Cannot call 'get_local_current_membership_for_user_in_room' on "
  466. "non-local user %s" % (user_id,),
  467. )
  468. results_dict = await self.db_pool.simple_select_one(
  469. "local_current_membership",
  470. {"room_id": room_id, "user_id": user_id},
  471. ("membership", "event_id"),
  472. allow_none=True,
  473. desc="get_local_current_membership_for_user_in_room",
  474. )
  475. if not results_dict:
  476. return None, None
  477. return results_dict.get("membership"), results_dict.get("event_id")
  478. @cached(max_entries=500000, iterable=True)
  479. async def get_rooms_for_user_with_stream_ordering(
  480. self, user_id: str
  481. ) -> FrozenSet[GetRoomsForUserWithStreamOrdering]:
  482. """Returns a set of room_ids the user is currently joined to.
  483. If a remote user only returns rooms this server is currently
  484. participating in.
  485. Args:
  486. user_id
  487. Returns:
  488. Returns the rooms the user is in currently, along with the stream
  489. ordering of the most recent join for that user and room, along with
  490. the room version of the room.
  491. """
  492. return await self.db_pool.runInteraction(
  493. "get_rooms_for_user_with_stream_ordering",
  494. self._get_rooms_for_user_with_stream_ordering_txn,
  495. user_id,
  496. )
  497. def _get_rooms_for_user_with_stream_ordering_txn(
  498. self, txn: LoggingTransaction, user_id: str
  499. ) -> FrozenSet[GetRoomsForUserWithStreamOrdering]:
  500. # We use `current_state_events` here and not `local_current_membership`
  501. # as a) this gets called with remote users and b) this only gets called
  502. # for rooms the server is participating in.
  503. sql = """
  504. SELECT room_id, e.instance_name, e.stream_ordering
  505. FROM current_state_events AS c
  506. INNER JOIN events AS e USING (room_id, event_id)
  507. WHERE
  508. c.type = 'm.room.member'
  509. AND c.state_key = ?
  510. AND c.membership = ?
  511. """
  512. txn.execute(sql, (user_id, Membership.JOIN))
  513. return frozenset(
  514. GetRoomsForUserWithStreamOrdering(
  515. room_id, PersistedEventPosition(instance, stream_id)
  516. )
  517. for room_id, instance, stream_id in txn
  518. )
  519. async def get_users_server_still_shares_room_with(
  520. self, user_ids: Collection[str]
  521. ) -> Set[str]:
  522. """Given a list of users return the set that the server still share a
  523. room with.
  524. """
  525. if not user_ids:
  526. return set()
  527. return await self.db_pool.runInteraction(
  528. "get_users_server_still_shares_room_with",
  529. self.get_users_server_still_shares_room_with_txn,
  530. user_ids,
  531. )
  532. def get_users_server_still_shares_room_with_txn(
  533. self,
  534. txn: LoggingTransaction,
  535. user_ids: Collection[str],
  536. ) -> Set[str]:
  537. if not user_ids:
  538. return set()
  539. sql = """
  540. SELECT state_key FROM current_state_events
  541. WHERE
  542. type = 'm.room.member'
  543. AND membership = 'join'
  544. AND %s
  545. GROUP BY state_key
  546. """
  547. clause, args = make_in_list_sql_clause(
  548. self.database_engine, "state_key", user_ids
  549. )
  550. txn.execute(sql % (clause,), args)
  551. return {row[0] for row in txn}
  552. @cached(max_entries=500000, iterable=True)
  553. async def get_rooms_for_user(self, user_id: str) -> FrozenSet[str]:
  554. """Returns a set of room_ids the user is currently joined to.
  555. If a remote user only returns rooms this server is currently
  556. participating in.
  557. """
  558. rooms = self.get_rooms_for_user_with_stream_ordering.cache.get_immediate(
  559. (user_id,),
  560. None,
  561. update_metrics=False,
  562. )
  563. if rooms:
  564. return frozenset(r.room_id for r in rooms)
  565. room_ids = await self.db_pool.simple_select_onecol(
  566. table="current_state_events",
  567. keyvalues={
  568. "type": EventTypes.Member,
  569. "membership": Membership.JOIN,
  570. "state_key": user_id,
  571. },
  572. retcol="room_id",
  573. desc="get_rooms_for_user",
  574. )
  575. return frozenset(room_ids)
  576. @cachedList(
  577. cached_method_name="get_rooms_for_user",
  578. list_name="user_ids",
  579. )
  580. async def _get_rooms_for_users(
  581. self, user_ids: Collection[str]
  582. ) -> Dict[str, FrozenSet[str]]:
  583. """A batched version of `get_rooms_for_user`.
  584. Returns:
  585. Map from user_id to set of rooms that is currently in.
  586. """
  587. rows = await self.db_pool.simple_select_many_batch(
  588. table="current_state_events",
  589. column="state_key",
  590. iterable=user_ids,
  591. retcols=(
  592. "state_key",
  593. "room_id",
  594. ),
  595. keyvalues={
  596. "type": EventTypes.Member,
  597. "membership": Membership.JOIN,
  598. },
  599. desc="get_rooms_for_users",
  600. )
  601. user_rooms: Dict[str, Set[str]] = {user_id: set() for user_id in user_ids}
  602. for row in rows:
  603. user_rooms[row["state_key"]].add(row["room_id"])
  604. return {key: frozenset(rooms) for key, rooms in user_rooms.items()}
  605. async def get_rooms_for_users(
  606. self, user_ids: Collection[str]
  607. ) -> Dict[str, FrozenSet[str]]:
  608. """A batched wrapper around `_get_rooms_for_users`, to prevent locking
  609. other calls to `get_rooms_for_user` for large user lists.
  610. """
  611. all_user_rooms: Dict[str, FrozenSet[str]] = {}
  612. # 250 users is pretty arbitrary but the data can be quite large if users
  613. # are in many rooms.
  614. for batch_user_ids in batch_iter(user_ids, 250):
  615. all_user_rooms.update(await self._get_rooms_for_users(batch_user_ids))
  616. return all_user_rooms
  617. @cached(max_entries=10000)
  618. async def does_pair_of_users_share_a_room(
  619. self, user_id: str, other_user_id: str
  620. ) -> bool:
  621. raise NotImplementedError()
  622. @cachedList(
  623. cached_method_name="does_pair_of_users_share_a_room", list_name="other_user_ids"
  624. )
  625. async def _do_users_share_a_room(
  626. self, user_id: str, other_user_ids: Collection[str]
  627. ) -> Mapping[str, Optional[bool]]:
  628. """Return mapping from user ID to whether they share a room with the
  629. given user.
  630. Note: `None` and `False` are equivalent and mean they don't share a
  631. room.
  632. """
  633. def do_users_share_a_room_txn(
  634. txn: LoggingTransaction, user_ids: Collection[str]
  635. ) -> Dict[str, bool]:
  636. clause, args = make_in_list_sql_clause(
  637. self.database_engine, "state_key", user_ids
  638. )
  639. # This query works by fetching both the list of rooms for the target
  640. # user and the set of other users, and then checking if there is any
  641. # overlap.
  642. sql = f"""
  643. SELECT DISTINCT b.state_key
  644. FROM (
  645. SELECT room_id FROM current_state_events
  646. WHERE type = 'm.room.member' AND membership = 'join' AND state_key = ?
  647. ) AS a
  648. INNER JOIN (
  649. SELECT room_id, state_key FROM current_state_events
  650. WHERE type = 'm.room.member' AND membership = 'join' AND {clause}
  651. ) AS b using (room_id)
  652. """
  653. txn.execute(sql, (user_id, *args))
  654. return {u: True for u, in txn}
  655. to_return = {}
  656. for batch_user_ids in batch_iter(other_user_ids, 1000):
  657. res = await self.db_pool.runInteraction(
  658. "do_users_share_a_room", do_users_share_a_room_txn, batch_user_ids
  659. )
  660. to_return.update(res)
  661. return to_return
  662. async def do_users_share_a_room(
  663. self, user_id: str, other_user_ids: Collection[str]
  664. ) -> Set[str]:
  665. """Return the set of users who share a room with the first users"""
  666. user_dict = await self._do_users_share_a_room(user_id, other_user_ids)
  667. return {u for u, share_room in user_dict.items() if share_room}
  668. async def get_users_who_share_room_with_user(self, user_id: str) -> Set[str]:
  669. """Returns the set of users who share a room with `user_id`"""
  670. room_ids = await self.get_rooms_for_user(user_id)
  671. user_who_share_room: Set[str] = set()
  672. for room_id in room_ids:
  673. user_ids = await self.get_users_in_room(room_id)
  674. user_who_share_room.update(user_ids)
  675. return user_who_share_room
  676. @cached(cache_context=True, iterable=True)
  677. async def get_mutual_rooms_between_users(
  678. self, user_ids: FrozenSet[str], cache_context: _CacheContext
  679. ) -> FrozenSet[str]:
  680. """
  681. Returns the set of rooms that all users in `user_ids` share.
  682. Args:
  683. user_ids: A frozen set of all users to investigate and return
  684. overlapping joined rooms for.
  685. cache_context
  686. """
  687. shared_room_ids: Optional[FrozenSet[str]] = None
  688. for user_id in user_ids:
  689. room_ids = await self.get_rooms_for_user(
  690. user_id, on_invalidate=cache_context.invalidate
  691. )
  692. if shared_room_ids is not None:
  693. shared_room_ids &= room_ids
  694. else:
  695. shared_room_ids = room_ids
  696. return shared_room_ids or frozenset()
  697. async def get_joined_user_ids_from_state(
  698. self, room_id: str, state: StateMap[str]
  699. ) -> Set[str]:
  700. """
  701. For a given set of state IDs, get a set of user IDs in the room.
  702. This method checks the local event cache, before calling
  703. `_get_user_ids_from_membership_event_ids` for any uncached events.
  704. """
  705. with Measure(self._clock, "get_joined_user_ids_from_state"):
  706. users_in_room = set()
  707. member_event_ids = [
  708. e_id for key, e_id in state.items() if key[0] == EventTypes.Member
  709. ]
  710. # We check if we have any of the member event ids in the event cache
  711. # before we ask the DB
  712. # We don't update the event cache hit ratio as it completely throws off
  713. # the hit ratio counts. After all, we don't populate the cache if we
  714. # miss it here
  715. event_map = self._get_events_from_local_cache(
  716. member_event_ids, update_metrics=False
  717. )
  718. missing_member_event_ids = []
  719. for event_id in member_event_ids:
  720. ev_entry = event_map.get(event_id)
  721. if ev_entry and not ev_entry.event.rejected_reason:
  722. if ev_entry.event.membership == Membership.JOIN:
  723. users_in_room.add(ev_entry.event.state_key)
  724. else:
  725. missing_member_event_ids.append(event_id)
  726. if missing_member_event_ids:
  727. event_to_memberships = (
  728. await self._get_user_ids_from_membership_event_ids(
  729. missing_member_event_ids
  730. )
  731. )
  732. users_in_room.update(
  733. user_id for user_id in event_to_memberships.values() if user_id
  734. )
  735. return users_in_room
  736. @cached(
  737. max_entries=10000,
  738. # This name matches the old function that has been replaced - the cache name
  739. # is kept here to maintain backwards compatibility.
  740. name="_get_joined_profile_from_event_id",
  741. )
  742. def _get_user_id_from_membership_event_id(
  743. self, event_id: str
  744. ) -> Optional[Tuple[str, ProfileInfo]]:
  745. raise NotImplementedError()
  746. @cachedList(
  747. cached_method_name="_get_user_id_from_membership_event_id",
  748. list_name="event_ids",
  749. )
  750. async def _get_user_ids_from_membership_event_ids(
  751. self, event_ids: Iterable[str]
  752. ) -> Dict[str, Optional[str]]:
  753. """For given set of member event_ids check if they point to a join
  754. event.
  755. Args:
  756. event_ids: The member event IDs to lookup
  757. Returns:
  758. Map from event ID to `user_id`, or None if event is not a join.
  759. """
  760. rows = await self.db_pool.simple_select_many_batch(
  761. table="room_memberships",
  762. column="event_id",
  763. iterable=event_ids,
  764. retcols=("user_id", "event_id"),
  765. keyvalues={"membership": Membership.JOIN},
  766. batch_size=1000,
  767. desc="_get_user_ids_from_membership_event_ids",
  768. )
  769. return {row["event_id"]: row["user_id"] for row in rows}
  770. @cached(max_entries=10000)
  771. async def is_host_joined(self, room_id: str, host: str) -> bool:
  772. return await self._check_host_room_membership(room_id, host, Membership.JOIN)
  773. @cached(max_entries=10000)
  774. async def is_host_invited(self, room_id: str, host: str) -> bool:
  775. return await self._check_host_room_membership(room_id, host, Membership.INVITE)
  776. async def _check_host_room_membership(
  777. self, room_id: str, host: str, membership: str
  778. ) -> bool:
  779. if "%" in host or "_" in host:
  780. raise Exception("Invalid host name")
  781. sql = """
  782. SELECT state_key FROM current_state_events AS c
  783. INNER JOIN room_memberships AS m USING (event_id)
  784. WHERE m.membership = ?
  785. AND type = 'm.room.member'
  786. AND c.room_id = ?
  787. AND state_key LIKE ?
  788. LIMIT 1
  789. """
  790. # We do need to be careful to ensure that host doesn't have any wild cards
  791. # in it, but we checked above for known ones and we'll check below that
  792. # the returned user actually has the correct domain.
  793. like_clause = "%:" + host
  794. rows = await self.db_pool.execute(
  795. "is_host_joined", None, sql, membership, room_id, like_clause
  796. )
  797. if not rows:
  798. return False
  799. user_id = rows[0][0]
  800. if get_domain_from_id(user_id) != host:
  801. # This can only happen if the host name has something funky in it
  802. raise Exception("Invalid host name")
  803. return True
  804. @cached(iterable=True, max_entries=10000)
  805. async def get_current_hosts_in_room(self, room_id: str) -> AbstractSet[str]:
  806. """Get current hosts in room based on current state."""
  807. # First we check if we already have `get_users_in_room` in the cache, as
  808. # we can just calculate result from that
  809. users = self.get_users_in_room.cache.get_immediate(
  810. (room_id,), None, update_metrics=False
  811. )
  812. if users is not None:
  813. return {get_domain_from_id(u) for u in users}
  814. if isinstance(self.database_engine, Sqlite3Engine):
  815. # If we're using SQLite then let's just always use
  816. # `get_users_in_room` rather than funky SQL.
  817. users = await self.get_users_in_room(room_id)
  818. return {get_domain_from_id(u) for u in users}
  819. # For PostgreSQL we can use a regex to pull out the domains from the
  820. # joined users in `current_state_events` via regex.
  821. def get_current_hosts_in_room_txn(txn: LoggingTransaction) -> Set[str]:
  822. sql = """
  823. SELECT DISTINCT substring(state_key FROM '@[^:]*:(.*)$')
  824. FROM current_state_events
  825. WHERE
  826. type = 'm.room.member'
  827. AND membership = 'join'
  828. AND room_id = ?
  829. """
  830. txn.execute(sql, (room_id,))
  831. return {d for d, in txn}
  832. return await self.db_pool.runInteraction(
  833. "get_current_hosts_in_room", get_current_hosts_in_room_txn
  834. )
  835. @cached(iterable=True, max_entries=10000)
  836. async def get_current_hosts_in_room_ordered(self, room_id: str) -> List[str]:
  837. """
  838. Get current hosts in room based on current state.
  839. The heuristic of sorting by servers who have been in the room the
  840. longest is good because they're most likely to have anything we ask
  841. about.
  842. For SQLite the returned list is not ordered, as SQLite doesn't support
  843. the appropriate SQL.
  844. Uses `m.room.member`s in the room state at the current forward
  845. extremities to determine which hosts are in the room.
  846. Will return inaccurate results for rooms with partial state, since the
  847. state for the forward extremities of those rooms will exclude most
  848. members. We may also calculate room state incorrectly for such rooms and
  849. believe that a host is or is not in the room when the opposite is true.
  850. Returns:
  851. Returns a list of servers sorted by longest in the room first. (aka.
  852. sorted by join with the lowest depth first).
  853. """
  854. if isinstance(self.database_engine, Sqlite3Engine):
  855. # If we're using SQLite then let's just always use
  856. # `get_users_in_room` rather than funky SQL.
  857. domains = await self.get_current_hosts_in_room(room_id)
  858. return list(domains)
  859. # For PostgreSQL we can use a regex to pull out the domains from the
  860. # joined users in `current_state_events` via regex.
  861. def get_current_hosts_in_room_ordered_txn(txn: LoggingTransaction) -> List[str]:
  862. # Returns a list of servers currently joined in the room sorted by
  863. # longest in the room first (aka. with the lowest depth). The
  864. # heuristic of sorting by servers who have been in the room the
  865. # longest is good because they're most likely to have anything we
  866. # ask about.
  867. sql = """
  868. SELECT
  869. /* Match the domain part of the MXID */
  870. substring(c.state_key FROM '@[^:]*:(.*)$') as server_domain
  871. FROM current_state_events c
  872. /* Get the depth of the event from the events table */
  873. INNER JOIN events AS e USING (event_id)
  874. WHERE
  875. /* Find any join state events in the room */
  876. c.type = 'm.room.member'
  877. AND c.membership = 'join'
  878. AND c.room_id = ?
  879. /* Group all state events from the same domain into their own buckets (groups) */
  880. GROUP BY server_domain
  881. /* Sorted by lowest depth first */
  882. ORDER BY min(e.depth) ASC;
  883. """
  884. txn.execute(sql, (room_id,))
  885. # `server_domain` will be `NULL` for malformed MXIDs with no colons.
  886. return [d for d, in txn if d is not None]
  887. return await self.db_pool.runInteraction(
  888. "get_current_hosts_in_room_ordered", get_current_hosts_in_room_ordered_txn
  889. )
  890. async def get_joined_hosts(
  891. self, room_id: str, state: StateMap[str], state_entry: "_StateCacheEntry"
  892. ) -> FrozenSet[str]:
  893. state_group: Union[object, int] = state_entry.state_group
  894. if not state_group:
  895. # If state_group is None it means it has yet to be assigned a
  896. # state group, i.e. we need to make sure that calls with a state_group
  897. # of None don't hit previous cached calls with a None state_group.
  898. # To do this we set the state_group to a new object as object() != object()
  899. state_group = object()
  900. assert state_group is not None
  901. with Measure(self._clock, "get_joined_hosts"):
  902. return await self._get_joined_hosts(
  903. room_id, state_group, state, state_entry=state_entry
  904. )
  905. @cached(num_args=2, max_entries=10000, iterable=True)
  906. async def _get_joined_hosts(
  907. self,
  908. room_id: str,
  909. state_group: Union[object, int],
  910. state: StateMap[str],
  911. state_entry: "_StateCacheEntry",
  912. ) -> FrozenSet[str]:
  913. # We don't use `state_group`, it's there so that we can cache based on
  914. # it. However, its important that its never None, since two
  915. # current_state's with a state_group of None are likely to be different.
  916. #
  917. # The `state_group` must match the `state_entry.state_group` (if not None).
  918. assert state_group is not None
  919. assert state_entry.state_group is None or state_entry.state_group == state_group
  920. # We use a secondary cache of previous work to allow us to build up the
  921. # joined hosts for the given state group based on previous state groups.
  922. #
  923. # We cache one object per room containing the results of the last state
  924. # group we got joined hosts for. The idea is that generally
  925. # `get_joined_hosts` is called with the "current" state group for the
  926. # room, and so consecutive calls will be for consecutive state groups
  927. # which point to the previous state group.
  928. cache = await self._get_joined_hosts_cache(room_id)
  929. # If the state group in the cache matches, we already have the data we need.
  930. if state_entry.state_group == cache.state_group:
  931. return frozenset(cache.hosts_to_joined_users)
  932. # Since we'll mutate the cache we need to lock.
  933. async with self._joined_host_linearizer.queue(room_id):
  934. if state_entry.state_group == cache.state_group:
  935. # Same state group, so nothing to do. We've already checked for
  936. # this above, but the cache may have changed while waiting on
  937. # the lock.
  938. pass
  939. elif state_entry.prev_group == cache.state_group:
  940. # The cached work is for the previous state group, so we work out
  941. # the delta.
  942. assert state_entry.delta_ids is not None
  943. for (typ, state_key), event_id in state_entry.delta_ids.items():
  944. if typ != EventTypes.Member:
  945. continue
  946. host = intern_string(get_domain_from_id(state_key))
  947. user_id = state_key
  948. known_joins = cache.hosts_to_joined_users.setdefault(host, set())
  949. event = await self.get_event(event_id)
  950. if event.membership == Membership.JOIN:
  951. known_joins.add(user_id)
  952. else:
  953. known_joins.discard(user_id)
  954. if not known_joins:
  955. cache.hosts_to_joined_users.pop(host, None)
  956. else:
  957. # The cache doesn't match the state group or prev state group,
  958. # so we calculate the result from first principles.
  959. #
  960. # We need to fetch all hosts joined to the room according to `state` by
  961. # inspecting all join memberships in `state`. However, if the `state` is
  962. # relatively recent then many of its events are likely to be held in
  963. # the current state of the room, which is easily available and likely
  964. # cached.
  965. #
  966. # We therefore compute the set of `state` events not in the
  967. # current state and only fetch those.
  968. current_memberships = (
  969. await self._get_approximate_current_memberships_in_room(room_id)
  970. )
  971. unknown_state_events = {}
  972. joined_users_in_current_state = []
  973. for (type, state_key), event_id in state.items():
  974. if event_id not in current_memberships:
  975. unknown_state_events[type, state_key] = event_id
  976. elif current_memberships[event_id] == Membership.JOIN:
  977. joined_users_in_current_state.append(state_key)
  978. joined_user_ids = await self.get_joined_user_ids_from_state(
  979. room_id, unknown_state_events
  980. )
  981. cache.hosts_to_joined_users = {}
  982. for user_id in chain(joined_user_ids, joined_users_in_current_state):
  983. host = intern_string(get_domain_from_id(user_id))
  984. cache.hosts_to_joined_users.setdefault(host, set()).add(user_id)
  985. if state_entry.state_group:
  986. cache.state_group = state_entry.state_group
  987. else:
  988. cache.state_group = object()
  989. return frozenset(cache.hosts_to_joined_users)
  990. async def _get_approximate_current_memberships_in_room(
  991. self, room_id: str
  992. ) -> Mapping[str, Optional[str]]:
  993. """Build a map from event id to membership, for all events in the current state.
  994. The event ids of non-memberships events (e.g. `m.room.power_levels`) are present
  995. in the result, mapped to values of `None`.
  996. The result is approximate for partially-joined rooms. It is fully accurate
  997. for fully-joined rooms.
  998. """
  999. rows = await self.db_pool.simple_select_list(
  1000. "current_state_events",
  1001. keyvalues={"room_id": room_id},
  1002. retcols=("event_id", "membership"),
  1003. desc="has_completed_background_updates",
  1004. )
  1005. return {row["event_id"]: row["membership"] for row in rows}
  1006. @cached(max_entries=10000)
  1007. def _get_joined_hosts_cache(self, room_id: str) -> "_JoinedHostsCache":
  1008. return _JoinedHostsCache()
  1009. @cached(num_args=2)
  1010. async def did_forget(self, user_id: str, room_id: str) -> bool:
  1011. """Returns whether user_id has elected to discard history for room_id.
  1012. Returns False if they have since re-joined."""
  1013. def f(txn: LoggingTransaction) -> int:
  1014. sql = (
  1015. "SELECT"
  1016. " COUNT(*)"
  1017. " FROM"
  1018. " room_memberships"
  1019. " WHERE"
  1020. " user_id = ?"
  1021. " AND"
  1022. " room_id = ?"
  1023. " AND"
  1024. " forgotten = 0"
  1025. )
  1026. txn.execute(sql, (user_id, room_id))
  1027. rows = txn.fetchall()
  1028. return rows[0][0]
  1029. count = await self.db_pool.runInteraction("did_forget_membership", f)
  1030. return count == 0
  1031. @cached()
  1032. async def get_forgotten_rooms_for_user(self, user_id: str) -> AbstractSet[str]:
  1033. """Gets all rooms the user has forgotten.
  1034. Args:
  1035. user_id: The user ID to query the rooms of.
  1036. Returns:
  1037. The forgotten rooms.
  1038. """
  1039. def _get_forgotten_rooms_for_user_txn(txn: LoggingTransaction) -> Set[str]:
  1040. # This is a slightly convoluted query that first looks up all rooms
  1041. # that the user has forgotten in the past, then rechecks that list
  1042. # to see if any have subsequently been updated. This is done so that
  1043. # we can use a partial index on `forgotten = 1` on the assumption
  1044. # that few users will actually forget many rooms.
  1045. #
  1046. # Note that a room is considered "forgotten" if *all* membership
  1047. # events for that user and room have the forgotten field set (as
  1048. # when a user forgets a room we update all rows for that user and
  1049. # room, not just the current one).
  1050. sql = """
  1051. SELECT room_id, (
  1052. SELECT count(*) FROM room_memberships
  1053. WHERE room_id = m.room_id AND user_id = m.user_id AND forgotten = 0
  1054. ) AS count
  1055. FROM room_memberships AS m
  1056. WHERE user_id = ? AND forgotten = 1
  1057. GROUP BY room_id, user_id;
  1058. """
  1059. txn.execute(sql, (user_id,))
  1060. return {row[0] for row in txn if row[1] == 0}
  1061. return await self.db_pool.runInteraction(
  1062. "get_forgotten_rooms_for_user", _get_forgotten_rooms_for_user_txn
  1063. )
  1064. async def is_locally_forgotten_room(self, room_id: str) -> bool:
  1065. """Returns whether all local users have forgotten this room_id.
  1066. Args:
  1067. room_id: The room ID to query.
  1068. Returns:
  1069. Whether the room is forgotten.
  1070. """
  1071. sql = """
  1072. SELECT count(*) > 0 FROM local_current_membership
  1073. INNER JOIN room_memberships USING (room_id, event_id)
  1074. WHERE
  1075. room_id = ?
  1076. AND forgotten = 0;
  1077. """
  1078. rows = await self.db_pool.execute("is_forgotten_room", None, sql, room_id)
  1079. # `count(*)` returns always an integer
  1080. # If any rows still exist it means someone has not forgotten this room yet
  1081. return not rows[0][0]
  1082. async def get_rooms_user_has_been_in(self, user_id: str) -> Set[str]:
  1083. """Get all rooms that the user has ever been in.
  1084. Args:
  1085. user_id: The user ID to get the rooms of.
  1086. Returns:
  1087. Set of room IDs.
  1088. """
  1089. room_ids = await self.db_pool.simple_select_onecol(
  1090. table="room_memberships",
  1091. keyvalues={"membership": Membership.JOIN, "user_id": user_id},
  1092. retcol="room_id",
  1093. desc="get_rooms_user_has_been_in",
  1094. )
  1095. return set(room_ids)
  1096. @cached(max_entries=5000)
  1097. async def _get_membership_from_event_id(
  1098. self, member_event_id: str
  1099. ) -> Optional[EventIdMembership]:
  1100. raise NotImplementedError()
  1101. @cachedList(
  1102. cached_method_name="_get_membership_from_event_id", list_name="member_event_ids"
  1103. )
  1104. async def get_membership_from_event_ids(
  1105. self, member_event_ids: Iterable[str]
  1106. ) -> Dict[str, Optional[EventIdMembership]]:
  1107. """Get user_id and membership of a set of event IDs.
  1108. Returns:
  1109. Mapping from event ID to `EventIdMembership` if the event is a
  1110. membership event, otherwise the value is None.
  1111. """
  1112. rows = await self.db_pool.simple_select_many_batch(
  1113. table="room_memberships",
  1114. column="event_id",
  1115. iterable=member_event_ids,
  1116. retcols=("user_id", "membership", "event_id"),
  1117. keyvalues={},
  1118. batch_size=500,
  1119. desc="get_membership_from_event_ids",
  1120. )
  1121. return {
  1122. row["event_id"]: EventIdMembership(
  1123. membership=row["membership"], user_id=row["user_id"]
  1124. )
  1125. for row in rows
  1126. }
  1127. async def is_local_host_in_room_ignoring_users(
  1128. self, room_id: str, ignore_users: Collection[str]
  1129. ) -> bool:
  1130. """Check if there are any local users, excluding those in the given
  1131. list, in the room.
  1132. """
  1133. clause, args = make_in_list_sql_clause(
  1134. self.database_engine, "user_id", ignore_users
  1135. )
  1136. sql = """
  1137. SELECT 1 FROM local_current_membership
  1138. WHERE
  1139. room_id = ? AND membership = ?
  1140. AND NOT (%s)
  1141. LIMIT 1
  1142. """ % (
  1143. clause,
  1144. )
  1145. def _is_local_host_in_room_ignoring_users_txn(
  1146. txn: LoggingTransaction,
  1147. ) -> bool:
  1148. txn.execute(sql, (room_id, Membership.JOIN, *args))
  1149. return bool(txn.fetchone())
  1150. return await self.db_pool.runInteraction(
  1151. "is_local_host_in_room_ignoring_users",
  1152. _is_local_host_in_room_ignoring_users_txn,
  1153. )
  1154. async def forget(self, user_id: str, room_id: str) -> None:
  1155. """Indicate that user_id wishes to discard history for room_id."""
  1156. def f(txn: LoggingTransaction) -> None:
  1157. self.db_pool.simple_update_txn(
  1158. txn,
  1159. table="room_memberships",
  1160. keyvalues={"user_id": user_id, "room_id": room_id},
  1161. updatevalues={"forgotten": 1},
  1162. )
  1163. self._invalidate_cache_and_stream(txn, self.did_forget, (user_id, room_id))
  1164. self._invalidate_cache_and_stream(
  1165. txn, self.get_forgotten_rooms_for_user, (user_id,)
  1166. )
  1167. await self.db_pool.runInteraction("forget_membership", f)
  1168. async def get_room_forgetter_stream_pos(self) -> int:
  1169. """Get the stream position of the background process to forget rooms when left
  1170. by users.
  1171. """
  1172. return await self.db_pool.simple_select_one_onecol(
  1173. table="room_forgetter_stream_pos",
  1174. keyvalues={},
  1175. retcol="stream_id",
  1176. desc="room_forgetter_stream_pos",
  1177. )
  1178. async def update_room_forgetter_stream_pos(self, stream_id: int) -> None:
  1179. """Update the stream position of the background process to forget rooms when
  1180. left by users.
  1181. Must only be used by the worker running the background process.
  1182. """
  1183. assert self.hs.config.worker.run_background_tasks
  1184. await self.db_pool.simple_update_one(
  1185. table="room_forgetter_stream_pos",
  1186. keyvalues={},
  1187. updatevalues={"stream_id": stream_id},
  1188. desc="room_forgetter_stream_pos",
  1189. )
  1190. class RoomMemberBackgroundUpdateStore(SQLBaseStore):
  1191. def __init__(
  1192. self,
  1193. database: DatabasePool,
  1194. db_conn: LoggingDatabaseConnection,
  1195. hs: "HomeServer",
  1196. ):
  1197. super().__init__(database, db_conn, hs)
  1198. self.db_pool.updates.register_background_update_handler(
  1199. _MEMBERSHIP_PROFILE_UPDATE_NAME, self._background_add_membership_profile
  1200. )
  1201. self.db_pool.updates.register_background_update_handler(
  1202. _CURRENT_STATE_MEMBERSHIP_UPDATE_NAME,
  1203. self._background_current_state_membership,
  1204. )
  1205. self.db_pool.updates.register_background_index_update(
  1206. "room_membership_forgotten_idx",
  1207. index_name="room_memberships_user_room_forgotten",
  1208. table="room_memberships",
  1209. columns=["user_id", "room_id"],
  1210. where_clause="forgotten = 1",
  1211. )
  1212. self.db_pool.updates.register_background_index_update(
  1213. "room_membership_user_room_index",
  1214. index_name="room_membership_user_room_idx",
  1215. table="room_memberships",
  1216. columns=["user_id", "room_id"],
  1217. )
  1218. async def _background_add_membership_profile(
  1219. self, progress: JsonDict, batch_size: int
  1220. ) -> int:
  1221. target_min_stream_id = progress.get(
  1222. "target_min_stream_id_inclusive", self._min_stream_order_on_start # type: ignore[attr-defined]
  1223. )
  1224. max_stream_id = progress.get(
  1225. "max_stream_id_exclusive", self._stream_order_on_start + 1 # type: ignore[attr-defined]
  1226. )
  1227. def add_membership_profile_txn(txn: LoggingTransaction) -> int:
  1228. sql = """
  1229. SELECT stream_ordering, event_id, events.room_id, event_json.json
  1230. FROM events
  1231. INNER JOIN event_json USING (event_id)
  1232. INNER JOIN room_memberships USING (event_id)
  1233. WHERE ? <= stream_ordering AND stream_ordering < ?
  1234. AND type = 'm.room.member'
  1235. ORDER BY stream_ordering DESC
  1236. LIMIT ?
  1237. """
  1238. txn.execute(sql, (target_min_stream_id, max_stream_id, batch_size))
  1239. rows = self.db_pool.cursor_to_dict(txn)
  1240. if not rows:
  1241. return 0
  1242. min_stream_id = rows[-1]["stream_ordering"]
  1243. to_update = []
  1244. for row in rows:
  1245. event_id = row["event_id"]
  1246. room_id = row["room_id"]
  1247. try:
  1248. event_json = db_to_json(row["json"])
  1249. content = event_json["content"]
  1250. except Exception:
  1251. continue
  1252. display_name = content.get("displayname", None)
  1253. avatar_url = content.get("avatar_url", None)
  1254. if display_name or avatar_url:
  1255. to_update.append((display_name, avatar_url, event_id, room_id))
  1256. to_update_sql = """
  1257. UPDATE room_memberships SET display_name = ?, avatar_url = ?
  1258. WHERE event_id = ? AND room_id = ?
  1259. """
  1260. txn.execute_batch(to_update_sql, to_update)
  1261. progress = {
  1262. "target_min_stream_id_inclusive": target_min_stream_id,
  1263. "max_stream_id_exclusive": min_stream_id,
  1264. }
  1265. self.db_pool.updates._background_update_progress_txn(
  1266. txn, _MEMBERSHIP_PROFILE_UPDATE_NAME, progress
  1267. )
  1268. return len(rows)
  1269. result = await self.db_pool.runInteraction(
  1270. _MEMBERSHIP_PROFILE_UPDATE_NAME, add_membership_profile_txn
  1271. )
  1272. if not result:
  1273. await self.db_pool.updates._end_background_update(
  1274. _MEMBERSHIP_PROFILE_UPDATE_NAME
  1275. )
  1276. return result
  1277. async def _background_current_state_membership(
  1278. self, progress: JsonDict, batch_size: int
  1279. ) -> int:
  1280. """Update the new membership column on current_state_events.
  1281. This works by iterating over all rooms in alphebetical order.
  1282. """
  1283. def _background_current_state_membership_txn(
  1284. txn: LoggingTransaction, last_processed_room: str
  1285. ) -> Tuple[int, bool]:
  1286. processed = 0
  1287. while processed < batch_size:
  1288. txn.execute(
  1289. """
  1290. SELECT MIN(room_id) FROM current_state_events WHERE room_id > ?
  1291. """,
  1292. (last_processed_room,),
  1293. )
  1294. row = txn.fetchone()
  1295. if not row or not row[0]:
  1296. return processed, True
  1297. (next_room,) = row
  1298. sql = """
  1299. UPDATE current_state_events
  1300. SET membership = (
  1301. SELECT membership FROM room_memberships
  1302. WHERE event_id = current_state_events.event_id
  1303. )
  1304. WHERE room_id = ?
  1305. """
  1306. txn.execute(sql, (next_room,))
  1307. processed += txn.rowcount
  1308. last_processed_room = next_room
  1309. self.db_pool.updates._background_update_progress_txn(
  1310. txn,
  1311. _CURRENT_STATE_MEMBERSHIP_UPDATE_NAME,
  1312. {"last_processed_room": last_processed_room},
  1313. )
  1314. return processed, False
  1315. # If we haven't got a last processed room then just use the empty
  1316. # string, which will compare before all room IDs correctly.
  1317. last_processed_room = progress.get("last_processed_room", "")
  1318. row_count, finished = await self.db_pool.runInteraction(
  1319. "_background_current_state_membership_update",
  1320. _background_current_state_membership_txn,
  1321. last_processed_room,
  1322. )
  1323. if finished:
  1324. await self.db_pool.updates._end_background_update(
  1325. _CURRENT_STATE_MEMBERSHIP_UPDATE_NAME
  1326. )
  1327. return row_count
  1328. class RoomMemberStore(
  1329. RoomMemberWorkerStore,
  1330. RoomMemberBackgroundUpdateStore,
  1331. CacheInvalidationWorkerStore,
  1332. ):
  1333. def __init__(
  1334. self,
  1335. database: DatabasePool,
  1336. db_conn: LoggingDatabaseConnection,
  1337. hs: "HomeServer",
  1338. ):
  1339. super().__init__(database, db_conn, hs)
  1340. def extract_heroes_from_room_summary(
  1341. details: Mapping[str, MemberSummary], me: str
  1342. ) -> List[str]:
  1343. """Determine the users that represent a room, from the perspective of the `me` user.
  1344. The rules which say which users we select are specified in the "Room Summary"
  1345. section of
  1346. https://spec.matrix.org/v1.4/client-server-api/#get_matrixclientv3sync
  1347. Returns a list (possibly empty) of heroes' mxids.
  1348. """
  1349. empty_ms = MemberSummary([], 0)
  1350. joined_user_ids = [
  1351. r[0] for r in details.get(Membership.JOIN, empty_ms).members if r[0] != me
  1352. ]
  1353. invited_user_ids = [
  1354. r[0] for r in details.get(Membership.INVITE, empty_ms).members if r[0] != me
  1355. ]
  1356. gone_user_ids = [
  1357. r[0] for r in details.get(Membership.LEAVE, empty_ms).members if r[0] != me
  1358. ] + [r[0] for r in details.get(Membership.BAN, empty_ms).members if r[0] != me]
  1359. # FIXME: order by stream ordering rather than as returned by SQL
  1360. if joined_user_ids or invited_user_ids:
  1361. return sorted(joined_user_ids + invited_user_ids)[0:5]
  1362. else:
  1363. return sorted(gone_user_ids)[0:5]
  1364. @attr.s(slots=True, auto_attribs=True)
  1365. class _JoinedHostsCache:
  1366. """The cached data used by the `_get_joined_hosts_cache`."""
  1367. # Dict of host to the set of their users in the room at the state group.
  1368. hosts_to_joined_users: Dict[str, Set[str]] = attr.Factory(dict)
  1369. # The state group `hosts_to_joined_users` is derived from. Will be an object
  1370. # if the instance is newly created or if the state is not based on a state
  1371. # group. (An object is used as a sentinel value to ensure that it never is
  1372. # equal to anything else).
  1373. state_group: Union[object, int] = attr.Factory(object)
  1374. def __len__(self) -> int:
  1375. return sum(len(v) for v in self.hosts_to_joined_users.values())