Allow filtering by fields

Getting the right `get_profile_data_for_users` SQL changes involved some Claude support.

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

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
 

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,

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)