From dff23cf704c4f7db0b2f7e1c643cd069bdcbcfee Mon Sep 17 00:00:00 2001 From: Ryan Skonnord Date: Fri, 12 Apr 2024 15:50:37 -0700 Subject: [PATCH 1/2] ref(hc): Introduce RpcUserProfile and UserService.get_many_profiles Offer a more performant alternative to UserService.get_many that provides all attributes that can be provided without fetching from any related tables. --- src/sentry/services/hybrid_cloud/user/impl.py | 16 +++++++-- .../services/hybrid_cloud/user/model.py | 8 +++-- .../services/hybrid_cloud/user/serial.py | 34 ++++++++++++++----- .../services/hybrid_cloud/user/service.py | 12 ++++++- tests/sentry/hybridcloud/test_user.py | 7 ++++ 5 files changed, 63 insertions(+), 14 deletions(-) diff --git a/src/sentry/services/hybrid_cloud/user/impl.py b/src/sentry/services/hybrid_cloud/user/impl.py index 1944f33a409805..1a42e53f03555d 100644 --- a/src/sentry/services/hybrid_cloud/user/impl.py +++ b/src/sentry/services/hybrid_cloud/user/impl.py @@ -37,8 +37,16 @@ UserSerializeType, UserUpdateArgs, ) -from sentry.services.hybrid_cloud.user.model import RpcVerifyUserEmail, UserIdEmailArgs -from sentry.services.hybrid_cloud.user.serial import serialize_rpc_user, serialize_user_avatar +from sentry.services.hybrid_cloud.user.model import ( + RpcUserProfile, + RpcVerifyUserEmail, + UserIdEmailArgs, +) +from sentry.services.hybrid_cloud.user.serial import ( + serialize_rpc_user, + serialize_rpc_user_profile, + serialize_user_avatar, +) from sentry.services.hybrid_cloud.user.service import UserService from sentry.signals import user_signup @@ -62,6 +70,10 @@ def get_many(self, *, filter: UserFilterArgs) -> list[RpcUser]: def get_many_ids(self, *, filter: UserFilterArgs) -> list[int]: return self._FQ.get_many_ids(filter) + def get_many_profiles(self, *, filter: UserFilterArgs) -> list[RpcUserProfile]: + users = self._FQ._query_many(filter, ids_only=True) + return [serialize_rpc_user_profile(user) for user in users] + def get_many_by_email( self, emails: list[str], diff --git a/src/sentry/services/hybrid_cloud/user/model.py b/src/sentry/services/hybrid_cloud/user/model.py index df9038b68cc66b..92d88342f1ce73 100644 --- a/src/sentry/services/hybrid_cloud/user/model.py +++ b/src/sentry/services/hybrid_cloud/user/model.py @@ -34,12 +34,13 @@ class RpcAuthenticator(RpcModel): config: Any = None -class RpcUser(RpcModel): +class RpcUserProfile(RpcModel): + """Minimal set of user attributes that can be fetched efficiently.""" + id: int = -1 pk: int = -1 name: str = "" email: str = "" - emails: frozenset[str] = frozenset() username: str = "" actor_id: int | None = None display_name: str = "" @@ -55,9 +56,12 @@ class RpcUser(RpcModel): is_password_expired: bool = False session_nonce: str | None = None + +class RpcUser(RpcUserProfile): roles: frozenset[str] = frozenset() permissions: frozenset[str] = frozenset() avatar: RpcAvatar | None = None + emails: frozenset[str] = frozenset() useremails: list[RpcUserEmail] = Field(default_factory=list) authenticators: list[RpcAuthenticator] = Field(default_factory=list) diff --git a/src/sentry/services/hybrid_cloud/user/serial.py b/src/sentry/services/hybrid_cloud/user/serial.py index 7e8adb3768c985..a53e3faadd9fb4 100644 --- a/src/sentry/services/hybrid_cloud/user/serial.py +++ b/src/sentry/services/hybrid_cloud/user/serial.py @@ -8,7 +8,13 @@ from sentry.db.models import BaseQuerySet from sentry.models.avatars.user_avatar import UserAvatar from sentry.models.user import User -from sentry.services.hybrid_cloud.user import RpcAuthenticator, RpcAvatar, RpcUser, RpcUserEmail +from sentry.services.hybrid_cloud.user import ( + RpcAuthenticator, + RpcAvatar, + RpcUser, + RpcUserEmail, + RpcUserProfile, +) def serialize_generic_user(user: Any) -> RpcUser | None: @@ -27,10 +33,10 @@ def serialize_generic_user(user: Any) -> RpcUser | None: raise TypeError(f"Can't serialize {type(user)} to RpcUser") -def serialize_rpc_user(user: User) -> RpcUser: +def _serialize_from_user_fields(user: User) -> dict[str, Any]: args = { field_name: getattr(user, field_name) - for field_name in RpcUser.__fields__ + for field_name in RpcUserProfile.__fields__ if hasattr(user, field_name) } args["pk"] = user.pk @@ -39,22 +45,32 @@ def serialize_rpc_user(user: User) -> RpcUser: args["is_superuser"] = user.is_superuser args["is_sentry_app"] = bool(user.is_sentry_app) args["password_usable"] = user.has_usable_password() + args["session_nonce"] = user.session_nonce + + if args["name"] is None: + # This field is non-nullable according to the Django schema, but may be null + # on some servers due to migration history + args["name"] = "" + + return args + + +def serialize_rpc_user_profile(user: User) -> RpcUserProfile: + return RpcUserProfile(**_serialize_from_user_fields(user)) + + +def serialize_rpc_user(user: User) -> RpcUser: + args = _serialize_from_user_fields(user) # Prefer eagerloaded attributes from _base_query if hasattr(user, "useremails") and user.useremails is not None: args["emails"] = frozenset([e["email"] for e in user.useremails if e["is_verified"]]) else: args["emails"] = frozenset([email.email for email in user.get_verified_emails()]) - args["session_nonce"] = user.session_nonce # And process the _base_query special data additions args["permissions"] = frozenset(getattr(user, "permissions", None) or ()) - if args["name"] is None: - # This field is non-nullable according to the Django schema, but may be null - # on some servers due to migration history - args["name"] = "" - roles: frozenset[str] = frozenset() if hasattr(user, "roles") and user.roles is not None: roles = frozenset(_flatten(user.roles)) diff --git a/src/sentry/services/hybrid_cloud/user/service.py b/src/sentry/services/hybrid_cloud/user/service.py index ac9de3881f4c2a..39f0706716590e 100644 --- a/src/sentry/services/hybrid_cloud/user/service.py +++ b/src/sentry/services/hybrid_cloud/user/service.py @@ -17,7 +17,12 @@ UserSerializeType, UserUpdateArgs, ) -from sentry.services.hybrid_cloud.user.model import RpcAvatar, RpcVerifyUserEmail, UserIdEmailArgs +from sentry.services.hybrid_cloud.user.model import ( + RpcAvatar, + RpcUserProfile, + RpcVerifyUserEmail, + UserIdEmailArgs, +) from sentry.silo import SiloMode @@ -53,6 +58,11 @@ def get_many(self, *, filter: UserFilterArgs) -> list[RpcUser]: def get_many_ids(self, *, filter: UserFilterArgs) -> list[int]: pass + @rpc_method + @abstractmethod + def get_many_profiles(self, *, filter: UserFilterArgs) -> list[RpcUserProfile]: + pass + @rpc_method @abstractmethod def get_many_by_email( diff --git a/tests/sentry/hybridcloud/test_user.py b/tests/sentry/hybridcloud/test_user.py index 727234de4748c3..5ef1ba737a4151 100644 --- a/tests/sentry/hybridcloud/test_user.py +++ b/tests/sentry/hybridcloud/test_user.py @@ -37,3 +37,10 @@ def test_user_serialize_multiple_emails(self): assert len(rpc_user.useremails) == 3 expected = {self.user.email, email.email, unverified_email.email} assert expected == {e.email for e in rpc_user.useremails} + + def test_get_many_profiles(self): + users = [self.create_user() for _ in range(2)] + target_ids = [users[0].id] + profiles = user_service.get_many_profiles(filter=dict(user_ids=target_ids)) + assert len(profiles) == 1 + assert profiles[0].id == users[0].id From b22d1d19aaeab54144296754554c48857610a4be Mon Sep 17 00:00:00 2001 From: Ryan Skonnord Date: Fri, 12 Apr 2024 15:57:19 -0700 Subject: [PATCH 2/2] ref(hc): Change FilterQueryDatabaseImpl signatures Rename `ids_only` to more accurately reflect what it does, since `get_many_profiles` is now using it to fetch more than just IDs but still wants to avoid selecting any related data. Make `query_many` no longer private by convention. We now want to call it directly, so that we can then do a different kind of serialization than _UserFilterQuery's default. (DatabaseBackedUserSocialAuthService also was previously calling it directly.) --- src/sentry/services/hybrid_cloud/app/impl.py | 4 ++-- src/sentry/services/hybrid_cloud/filter_query.py | 12 ++++++------ src/sentry/services/hybrid_cloud/identity/impl.py | 2 +- src/sentry/services/hybrid_cloud/user/impl.py | 6 +++--- src/sentry/services/hybrid_cloud/user_option/impl.py | 2 +- .../services/hybrid_cloud/usersocialauth/impl.py | 6 +++--- 6 files changed, 16 insertions(+), 16 deletions(-) diff --git a/src/sentry/services/hybrid_cloud/app/impl.py b/src/sentry/services/hybrid_cloud/app/impl.py index 7f382542eee11c..a73e4b5a7ded5f 100644 --- a/src/sentry/services/hybrid_cloud/app/impl.py +++ b/src/sentry/services/hybrid_cloud/app/impl.py @@ -168,8 +168,8 @@ class _AppServiceFilterQuery( SentryAppInstallation, SentryAppInstallationFilterArgs, RpcSentryAppInstallation, None ] ): - def base_query(self, ids_only: bool = False) -> QuerySet[SentryAppInstallation]: - if ids_only: + def base_query(self, select_related: bool = True) -> QuerySet[SentryAppInstallation]: + if not select_related: return SentryAppInstallation.objects return SentryAppInstallation.objects.select_related("sentry_app") diff --git a/src/sentry/services/hybrid_cloud/filter_query.py b/src/sentry/services/hybrid_cloud/filter_query.py index 0818cabaea5efb..72ab32d79048d2 100644 --- a/src/sentry/services/hybrid_cloud/filter_query.py +++ b/src/sentry/services/hybrid_cloud/filter_query.py @@ -50,7 +50,7 @@ class FilterQueryDatabaseImpl( # Required Overrides @abc.abstractmethod - def base_query(self, ids_only: bool = False) -> QuerySet[BASE_MODEL]: + def base_query(self, select_related: bool = True) -> QuerySet[BASE_MODEL]: # This should return a QuerySet for the model in question along with any other required data # that is not a filter pass @@ -91,14 +91,14 @@ def validator(d: FILTER_ARGS) -> str | None: # Helpers - def _query_many(self, filter: FILTER_ARGS, ids_only: bool = False) -> QuerySet: + def query_many(self, filter: FILTER_ARGS, select_related: bool = True) -> QuerySet: validation_error = self.filter_arg_validator()(filter) if validation_error is not None: raise TypeError( f"Failed to validate filter arguments passed to {self.__class__.__name__}: {validation_error}" ) - query = self.base_query(ids_only=ids_only) + query = self.base_query(select_related=select_related) return self.apply_filters(query, filter) # Public Interface @@ -121,7 +121,7 @@ def serialize_many( if as_user is None and auth_context: as_user = auth_context.user - result = self._query_many(filter=filter) + result = self.query_many(filter=filter) return serialize( list(result), user=as_user, @@ -129,7 +129,7 @@ def serialize_many( ) def get_many(self, filter: FILTER_ARGS) -> list[RPC_RESPONSE]: - return [self.serialize_rpc(o) for o in self._query_many(filter=filter)] + return [self.serialize_rpc(o) for o in self.query_many(filter=filter)] def get_many_ids(self, filter: FILTER_ARGS) -> list[int]: - return [o.id for o in self._query_many(filter=filter, ids_only=True)] + return [o.id for o in self.query_many(filter=filter, select_related=False)] diff --git a/src/sentry/services/hybrid_cloud/identity/impl.py b/src/sentry/services/hybrid_cloud/identity/impl.py index 8d528c20ae7508..f632fc0edfd889 100644 --- a/src/sentry/services/hybrid_cloud/identity/impl.py +++ b/src/sentry/services/hybrid_cloud/identity/impl.py @@ -108,7 +108,7 @@ def apply_filters( query = query.filter(idp__type=filters["provider_type"]) return query - def base_query(self, ids_only: bool = False) -> QuerySet[Identity]: + def base_query(self, select_related: bool = True) -> QuerySet[Identity]: return Identity.objects def filter_arg_validator(self) -> Callable[[IdentityFilterArgs], str | None]: diff --git a/src/sentry/services/hybrid_cloud/user/impl.py b/src/sentry/services/hybrid_cloud/user/impl.py index 1a42e53f03555d..3fe3705514cd7e 100644 --- a/src/sentry/services/hybrid_cloud/user/impl.py +++ b/src/sentry/services/hybrid_cloud/user/impl.py @@ -71,7 +71,7 @@ def get_many_ids(self, *, filter: UserFilterArgs) -> list[int]: return self._FQ.get_many_ids(filter) def get_many_profiles(self, *, filter: UserFilterArgs) -> list[RpcUserProfile]: - users = self._FQ._query_many(filter, ids_only=True) + users = self._FQ.query_many(filter, select_related=False) return [serialize_rpc_user_profile(user) for user in users] def get_many_by_email( @@ -325,8 +325,8 @@ def apply_filters(self, query: QuerySet[User], filters: UserFilterArgs) -> Query return query - def base_query(self, ids_only: bool = False) -> QuerySet[User]: - if ids_only: + def base_query(self, select_related: bool = True) -> QuerySet[User]: + if not select_related: return User.objects return User.objects.extra( diff --git a/src/sentry/services/hybrid_cloud/user_option/impl.py b/src/sentry/services/hybrid_cloud/user_option/impl.py index ccc6f67381ed79..9c25b6a9a82798 100644 --- a/src/sentry/services/hybrid_cloud/user_option/impl.py +++ b/src/sentry/services/hybrid_cloud/user_option/impl.py @@ -56,7 +56,7 @@ def set_option( class _UserOptionFilterQuery( FilterQueryDatabaseImpl[UserOption, UserOptionFilterArgs, RpcUserOption, None] ): - def base_query(self, ids_only: bool = False) -> QuerySet[UserOption]: + def base_query(self, select_related: bool = True) -> QuerySet[UserOption]: return UserOption.objects def filter_arg_validator(self) -> Callable[[UserOptionFilterArgs], str | None]: diff --git a/src/sentry/services/hybrid_cloud/usersocialauth/impl.py b/src/sentry/services/hybrid_cloud/usersocialauth/impl.py index d08facddfdf498..05c783ecf48314 100644 --- a/src/sentry/services/hybrid_cloud/usersocialauth/impl.py +++ b/src/sentry/services/hybrid_cloud/usersocialauth/impl.py @@ -37,7 +37,7 @@ def revoke_token( """ Calls UserSocialAuth.revoke_token() on all matching results, returning the modified RpcUserSocialAuths. """ - db_auths = self._FQ._query_many(filter=filter) + db_auths = self._FQ.query_many(filter=filter) for db_auth in db_auths: db_auth.revoke_token(drop_token=drop_token) return self.get_many(filter=filter) @@ -46,7 +46,7 @@ def refresh_token(self, *, filter: UserSocialAuthFilterArgs) -> list[RpcUserSoci """ Calls UserSocialAuth.refresh_token() on all matching results, returning the modified RpcUserSocialAuths. """ - db_auths = self._FQ._query_many(filter=filter) + db_auths = self._FQ.query_many(filter=filter) for db_auth in db_auths: db_auth.refresh_token() return self.get_many(filter=filter) @@ -78,7 +78,7 @@ def apply_filters( query = query.filter(uid=filters["uid"]) return query - def base_query(self, ids_only: bool = False) -> QuerySet[UserSocialAuth]: + def base_query(self, select_related: bool = True) -> QuerySet[UserSocialAuth]: return UserSocialAuth.objects.filter() def filter_arg_validator(self) -> Callable[[UserSocialAuthFilterArgs], str | None]: