De-localpart `ProfileWorkerStore.get_profileinfo()`

Signed-off-by: Sean Quah <seanq@matrix.org>

8 files changed, 25 insertions, 18 deletions

diff --git a/synapse/handlers/admin.py b/synapse/handlers/admin.py
index b06f25b03c..2b318b10db 100644
--- a/synapse/handlers/admin.py
+++ b/synapse/handlers/admin.py
@@ -89,7 +89,7 @@ class AdminHandler:
         }
 
         # Add additional user metadata
-        profile = await self._store.get_profileinfo(user.localpart)
+        profile = await self._store.get_profileinfo(user.to_string())
         threepids = await self._store.user_get_threepids(user.to_string())
         external_ids = [
             ({"auth_provider": auth_provider, "external_id": external_id})
diff --git a/synapse/handlers/auth.py b/synapse/handlers/auth.py
index 1e89447044..369f51ce12 100644
--- a/synapse/handlers/auth.py
+++ b/synapse/handlers/auth.py
@@ -1756,9 +1756,7 @@ class AuthHandler:
             respond_with_html(request, 403, self._sso_account_deactivated_template)
             return
 
-        user_profile_data = await self.store.get_profileinfo(
-            UserID.from_string(registered_user_id).localpart
-        )
+        user_profile_data = await self.store.get_profileinfo(registered_user_id)
 
         # Store any extra attributes which will be passed in the login response.
         # Note that this is per-user so it may overwrite a previous value, this
diff --git a/synapse/handlers/deactivate_account.py b/synapse/handlers/deactivate_account.py
index bd5867491b..77665e53a8 100644
--- a/synapse/handlers/deactivate_account.py
+++ b/synapse/handlers/deactivate_account.py
@@ -282,8 +282,6 @@ class DeactivateAccountHandler:
         Args:
             user_id: ID of user to be re-activated
         """
-        user = UserID.from_string(user_id)
-
         # Ensure the user is not marked as erased.
         await self.store.mark_user_not_erased(user_id)
 
@@ -297,5 +295,5 @@ class DeactivateAccountHandler:
         # Add the user to the directory, if necessary. Note that
         # this must be done after the user is re-activated, because
         # deactivated users are excluded from the user directory.
-        profile = await self.store.get_profileinfo(user.localpart)
+        profile = await self.store.get_profileinfo(user_id)
         await self.user_directory_handler.handle_local_profile_change(user_id, profile)
diff --git a/synapse/handlers/profile.py b/synapse/handlers/profile.py
index 9a81a77cbd..4fa5a8611f 100644
--- a/synapse/handlers/profile.py
+++ b/synapse/handlers/profile.py
@@ -67,7 +67,7 @@ class ProfileHandler:
         target_user = UserID.from_string(user_id)
 
         if self.hs.is_mine(target_user):
-            profileinfo = await self.store.get_profileinfo(target_user.localpart)
+            profileinfo = await self.store.get_profileinfo(user_id)
             if profileinfo.display_name is None:
                 raise SynapseError(404, "Profile was not found", Codes.NOT_FOUND)
 
@@ -147,7 +147,7 @@ class ProfileHandler:
             raise AuthError(400, "Cannot set another user's displayname")
 
         if not by_admin and not self.hs.config.registration.enable_set_displayname:
-            profile = await self.store.get_profileinfo(target_user.localpart)
+            profile = await self.store.get_profileinfo(target_user.to_string())
             if profile.display_name:
                 raise SynapseError(
                     400,
@@ -182,7 +182,7 @@ class ProfileHandler:
             target_user.localpart, displayname_to_set
         )
 
-        profile = await self.store.get_profileinfo(target_user.localpart)
+        profile = await self.store.get_profileinfo(target_user.to_string())
         await self.user_directory_handler.handle_local_profile_change(
             target_user.to_string(), profile
         )
@@ -243,7 +243,7 @@ class ProfileHandler:
             raise AuthError(400, "Cannot set another user's avatar_url")
 
         if not by_admin and not self.hs.config.registration.enable_set_avatar_url:
-            profile = await self.store.get_profileinfo(target_user.localpart)
+            profile = await self.store.get_profileinfo(target_user.to_string())
             if profile.avatar_url:
                 raise SynapseError(
                     400, "Changing avatar is disabled on this server", Codes.FORBIDDEN
@@ -276,7 +276,7 @@ class ProfileHandler:
             target_user.localpart, avatar_url_to_set
         )
 
-        profile = await self.store.get_profileinfo(target_user.localpart)
+        profile = await self.store.get_profileinfo(target_user.to_string())
         await self.user_directory_handler.handle_local_profile_change(
             target_user.to_string(), profile
         )
diff --git a/synapse/handlers/register.py b/synapse/handlers/register.py
index c8bf2439af..4a27ccd062 100644
--- a/synapse/handlers/register.py
+++ b/synapse/handlers/register.py
@@ -314,7 +314,7 @@ class RegistrationHandler:
                 approved=approved,
             )
 
-            profile = await self.store.get_profileinfo(localpart)
+            profile = await self.store.get_profileinfo(user_id)
             await self.user_directory_handler.handle_local_profile_change(
                 user_id, profile
             )
diff --git a/synapse/module_api/__init__.py b/synapse/module_api/__init__.py
index 595c23e78d..f8834dfdb3 100644
--- a/synapse/module_api/__init__.py
+++ b/synapse/module_api/__init__.py
@@ -648,7 +648,8 @@ class ModuleApi:
         Returns:
             The profile information (i.e. display name and avatar URL).
         """
-        return await self._store.get_profileinfo(localpart)
+        user = UserID(localpart, self._hs.hostname)
+        return await self._store.get_profileinfo(user.to_string())
 
     async def get_threepids_for_user(self, user_id: str) -> List[Dict[str, str]]:
         """Look up the threepids (email addresses and phone numbers) associated with the
diff --git a/synapse/storage/databases/main/profile.py b/synapse/storage/databases/main/profile.py
index ee7dcdb300..23021a1f1f 100644
--- a/synapse/storage/databases/main/profile.py
+++ b/synapse/storage/databases/main/profile.py
@@ -21,21 +21,31 @@ from synapse.storage.database import (
     LoggingTransaction,
 )
 from synapse.storage.databases.main.roommember import ProfileInfo
-from synapse.types import JsonDict
+from synapse.types import JsonDict, UserID
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
 
 
 class ProfileWorkerStore(SQLBaseStore):
-    async def get_profileinfo(self, user_localpart: str) -> ProfileInfo:
+    async def get_profileinfo(self, user_id: str) -> ProfileInfo:
         try:
             profile = await self.db_pool.simple_select_one(
                 table="profiles",
-                keyvalues={"user_id": user_localpart},
+                keyvalues={"full_user_id": user_id},
                 retcols=("displayname", "avatar_url"),
+                allow_none=True,
                 desc="get_profileinfo",
             )
+            if profile is None:
+                # Fall back to the `user_id` column.
+                user_localpart = UserID.from_string(user_id).localpart
+                profile = await self.db_pool.simple_select_one(
+                    table="profiles",
+                    keyvalues={"user_id": user_localpart},
+                    retcols=("displayname", "avatar_url"),
+                    desc="get_profileinfo",
+                )
         except StoreError as e:
             if e.code == 404:
                 # no match
diff --git a/synapse/storage/databases/main/user_directory.py b/synapse/storage/databases/main/user_directory.py
index 5d65faed16..4ced693eab 100644
--- a/synapse/storage/databases/main/user_directory.py
+++ b/synapse/storage/databases/main/user_directory.py
@@ -383,7 +383,7 @@ class UserDirectoryBackgroundUpdateStore(StateDeltasStore):
 
         for user_id in users_to_work_on:
             if await self.should_include_local_user_in_dir(user_id):
-                profile = await self.get_profileinfo(get_localpart_from_id(user_id))  # type: ignore[attr-defined]
+                profile = await self.get_profileinfo(user_id)  # type: ignore[attr-defined]
                 await self.update_profile_in_user_dir(
                     user_id, profile.display_name, profile.avatar_url
                 )