diff --git a/synapse/handlers/profile.py b/synapse/handlers/profile.py index 4ee0396b40..d10e6ece02 100644 --- a/synapse/handlers/profile.py +++ b/synapse/handlers/profile.py @@ -125,9 +125,9 @@ class ProfileHandler: self._profile_updates_writer_instance = self.hs.config.worker.writers.events[0] async def get_profiles( - self, user_ids: set[str] + self, user_ids: set[str], field_names: set[str] | None = None ) -> dict[str, dict[str, JsonValue | dict[str, JsonValue]]]: - return await self.store.get_profile_data_for_users(user_ids) + return await self.store.get_profile_data_for_users(user_ids, field_names) async def get_profile(self, user_id: str, ignore_backoff: bool = True) -> JsonDict: """ diff --git a/synapse/rest/client/profile.py b/synapse/rest/client/profile.py index 830963a854..4b0cfb4abf 100644 --- a/synapse/rest/client/profile.py +++ b/synapse/rest/client/profile.py @@ -370,6 +370,19 @@ class ProfileQueryRestServlet(RestServlet): Codes.INVALID_PARAM, ) target_users = set(users) + field_names = content.get("fields") + if field_names is not None and ( + not isinstance(field_names, list) + or not field_names + or not all(isinstance(field_name, str) for field_name in field_names) + ): + raise SynapseError( + HTTPStatus.BAD_REQUEST, + "`fields`, if given, must be a non-empty list of strings.", + Codes.INVALID_PARAM, + ) + if field_names is not None: + field_names = set(field_names) # Limit the amount of users pulled in one go target_users = set(list(target_users)[:PROFILE_QUERY_MAX_TARGET_USERS]) @@ -389,7 +402,9 @@ class ProfileQueryRestServlet(RestServlet): ) ) - ret = await self.profile_handler.get_profiles(filtered_profiles) + ret = await self.profile_handler.get_profiles( + filtered_profiles, field_names=field_names + ) return 200, ret diff --git a/synapse/storage/databases/main/profile.py b/synapse/storage/databases/main/profile.py index 8ed0d645bc..fbf3104930 100644 --- a/synapse/storage/databases/main/profile.py +++ b/synapse/storage/databases/main/profile.py @@ -20,7 +20,7 @@ # import json from collections.abc import Set -from typing import TYPE_CHECKING, Collection, Iterable, cast +from typing import TYPE_CHECKING, Collection, cast import attr from canonicaljson import encode_canonical_json @@ -43,6 +43,7 @@ from synapse.storage.databases.main.roommember import ProfileInfo from synapse.storage.engines import PostgresEngine, Sqlite3Engine from synapse.storage.util.id_generators import MultiWriterIdGenerator from synapse.types import JsonDict, JsonValue, UserID +from synapse.util.iterutils import batch_iter from synapse.util.json import json_encoder if TYPE_CHECKING: @@ -633,7 +634,7 @@ class ProfileWorkerStore(SQLBaseStore): ) async def get_profile_data_for_users( - self, user_ids: Collection[str] + self, user_ids: Collection[str], field_names: Collection[str] | None = None ) -> dict[str, dict[str, JsonValue | dict[str, JsonValue]]]: """Fetch displayname/avatar_url/custom fields for a list of users. @@ -642,32 +643,106 @@ class ProfileWorkerStore(SQLBaseStore): Args: user_ids: List of user IDs to filter against. + field_names: If given, only these fields are returned. Applies to + `displayname` and `avatar_url` as well as custom fields; names + that aren't set for a user are left out. `None` returns all + fields. Returns: Dictionary from user_id -> field name -> field value for the requested users. - This includes `displayname`, `avatar_url` and all custom fields. + Every requested user that has a profile row gets an entry, even when + none of the requested fields are set (it then maps to `{}`). Users + without a profile row are omitted. + For `displayname` and `avatar_url`, when they are stored as NULL in the database column, the dictionary entry will be omitted. """ if not user_ids: return {} - rows: Iterable[ - tuple[str, str | None, str | None, str | JsonDict | None] - ] = await self.db_pool.simple_select_many_batch( - table="profiles", - column="full_user_id", - iterable=user_ids, - retcols=("full_user_id", "displayname", "avatar_url", "fields"), - desc="get_profile_data_for_users", + include_displayname = ( + field_names is None or ProfileFields.DISPLAYNAME in field_names + ) + include_avatar_url = ( + field_names is None or ProfileFields.AVATAR_URL in field_names + ) + custom_field_names: set[str] | None = ( + None + if field_names is None + else set(field_names) + - {ProfileFields.DISPLAYNAME, ProfileFields.AVATAR_URL} + ) + + # Pick the custom fields in SQL rather than in Python. + fields_args: list[list[str]] = [] + if custom_field_names is None: + fields_column = "fields" + elif not custom_field_names: + fields_column = "NULL" + elif isinstance(self.database_engine, PostgresEngine): + # `JSONB_OBJECT_AGG` over no rows gives NULL, which we treat as `{}`. + key_clause, key_args = make_in_list_sql_clause( + self.database_engine, "f.key", custom_field_names + ) + fields_column = f"""( + SELECT JSONB_OBJECT_AGG(f.key, f.value) + FROM JSONB_EACH(fields) AS f + WHERE {key_clause} + )""" + fields_args.extend(key_args) + else: + # SQLite runs in-process, so there is no transfer to save. Rebuilding + # the object with `json_each`/`json_group_object` would also change + # value types (`true` becomes `1`, nested objects become strings), so + # filter in Python below instead. + fields_column = "fields" + + def _get_profile_data_for_users_txn( + txn: LoggingTransaction, + ) -> list[tuple[str, str | None, str | None, str | JsonDict | None]]: + rows: list[tuple[str, str | None, str | None, str | JsonDict | None]] = [] + for batch_user_ids in batch_iter(user_ids, 100): + user_id_clause, user_id_args = make_in_list_sql_clause( + self.database_engine, "full_user_id", batch_user_ids + ) + txn.execute( + f""" + SELECT + full_user_id, + {"displayname" if include_displayname else "NULL"}, + {"avatar_url" if include_avatar_url else "NULL"}, + {fields_column} + FROM profiles + WHERE {user_id_clause} + """, + (*fields_args, *user_id_args), + ) + rows.extend( + cast( + list[tuple[str, str | None, str | None, str | JsonDict | None]], + txn.fetchall(), + ) + ) + return rows + + rows = await self.db_pool.runInteraction( + "get_profile_data_for_users", _get_profile_data_for_users_txn ) results: dict[str, dict[str, JsonValue | dict[str, JsonValue]]] = {} for full_user_id, displayname, avatar_url, fields in rows: user_fields = db_to_json(fields or {}) + # On SQLite the custom fields haven't been filtered yet + if custom_field_names is not None and isinstance( + self.database_engine, Sqlite3Engine + ): + user_fields = { + k: v for k, v in user_fields.items() if k in custom_field_names + } + # When the displayname and avatar URL aren't set, # they are stored as NULL in the database. # To make them behave the same as custom fields, diff --git a/tests/storage/test_profile.py b/tests/storage/test_profile.py index 0d9abaff47..64b6a2d410 100644 --- a/tests/storage/test_profile.py +++ b/tests/storage/test_profile.py @@ -28,7 +28,7 @@ from synapse.api.errors import StoreError from synapse.server import HomeServer from synapse.storage.database import LoggingTransaction from synapse.storage.engines import PostgresEngine -from synapse.types import UserID +from synapse.types import JsonValue, UserID from synapse.util.clock import Clock from tests import unittest @@ -197,3 +197,129 @@ class ProfileStoreTestCase(unittest.HomeserverTestCase): ) self.assertEqual(len(res), len(expected_values)) self.assertEqual(res, expected_values) + + def _set_up_profiles_for_field_filtering(self) -> tuple[UserID, UserID, UserID]: + """Create three local profiles for the `get_profile_data_for_users` tests: + + - frank: displayname, avatar_url and custom fields with string, nested + object and boolean values (so JSON types surviving the filter are + checked, not just strings). + - bob: displayname only. + - carol: a profile row with nothing set. + """ + u_bob = UserID.from_string("@bob:test") + u_carol = UserID.from_string("@carol:test") + for user_id in (self.u_frank, u_bob, u_carol): + self.get_success(self.store.create_profile(user_id)) + + frank_fields: list[tuple[str, JsonValue | dict[str, JsonValue]]] = [ + (ProfileFields.DISPLAYNAME, "Frank"), + (ProfileFields.AVATAR_URL, "mxc://test/frank"), + ("m.status", {"emoji": "💬", "text": "In a meeting"}), + ("org.example.pronouns", "he/him"), + ("org.example.verified", True), + ] + for field_name, value in frank_fields: + self.get_success( + self.store.set_profile_field( + user_id=self.u_frank, field_name=field_name, new_value=value + ) + ) + self.get_success( + self.store.set_profile_field( + user_id=u_bob, field_name=ProfileFields.DISPLAYNAME, new_value="Bob" + ) + ) + + return self.u_frank, u_bob, u_carol + + def test_get_profile_data_for_users_all_fields(self) -> None: + """Without `field_names`, every set field is returned, users with an + empty profile map to `{}`, and users without a profile row are omitted. + """ + u_frank, u_bob, u_carol = self._set_up_profiles_for_field_filtering() + + result = self.get_success( + self.store.get_profile_data_for_users( + [ + u_frank.to_string(), + u_bob.to_string(), + u_carol.to_string(), + "@nobody:test", + ] + ) + ) + + self.assertEqual( + result, + { + "@frank:test": { + ProfileFields.DISPLAYNAME: "Frank", + ProfileFields.AVATAR_URL: "mxc://test/frank", + "m.status": {"emoji": "💬", "text": "In a meeting"}, + "org.example.pronouns": "he/him", + "org.example.verified": True, + }, + "@bob:test": {ProfileFields.DISPLAYNAME: "Bob"}, + "@carol:test": {}, + }, + ) + + def test_get_profile_data_for_users_filters_fields(self) -> None: + """`field_names` restricts the result to the named fields, covering the + column-backed fields (`displayname`, `avatar_url`), custom fields, and + names that aren't set for some or all users. + """ + u_frank, u_bob, u_carol = self._set_up_profiles_for_field_filtering() + user_ids = [u_frank.to_string(), u_bob.to_string(), u_carol.to_string()] + + # (field_names, expected result) + cases: list[tuple[set[str], dict]] = [ + # Only a column-backed field. + ( + {ProfileFields.DISPLAYNAME}, + { + "@frank:test": {ProfileFields.DISPLAYNAME: "Frank"}, + "@bob:test": {ProfileFields.DISPLAYNAME: "Bob"}, + "@carol:test": {}, + }, + ), + # Only custom fields, including a nested object and a boolean. + ( + {"m.status", "org.example.verified"}, + { + "@frank:test": { + "m.status": {"emoji": "💬", "text": "In a meeting"}, + "org.example.verified": True, + }, + "@bob:test": {}, + "@carol:test": {}, + }, + ), + # A mix of column-backed, custom and unknown fields. + ( + {ProfileFields.AVATAR_URL, "org.example.pronouns", "org.example.unset"}, + { + "@frank:test": { + ProfileFields.AVATAR_URL: "mxc://test/frank", + "org.example.pronouns": "he/him", + }, + "@bob:test": {}, + "@carol:test": {}, + }, + ), + # No fields at all: every user with a profile still appears, empty. + ( + set(), + {"@frank:test": {}, "@bob:test": {}, "@carol:test": {}}, + ), + ] + + for field_names, expected in cases: + with self.subTest(field_names=field_names): + result = self.get_success( + self.store.get_profile_data_for_users( + user_ids, field_names=field_names + ) + ) + self.assertEqual(result, expected)