feat(hc): Use hybrid cloud services in AuthIdentityHandler (#43035)

These changes are targeted toward getting all AuthHelper and
AuthIdentityHandler code to run with control silo limits. Most
availability errors from AuthIdentityHandler are fixed, but we need more
work around the Actor model and ApiInviteHelper to make it stable.

Co-authored-by: Zach Collins <recursive.cookie.jar@gmail.com>

src/sentry/auth/helper.py

@@ -21,7 +21,7 @@
 from django.views import View
 
 from sentry import audit_log, features
-from sentry.api.invite_helper import ApiInviteHelper, remove_invite_details_from_session
+from sentry.api.invite_helper import remove_invite_details_from_session
 from sentry.api.utils import generate_organization_url
 from sentry.auth.email import AmbiguousUserFromEmail, resolve_email_to_user
 from sentry.auth.exceptions import IdentityNotValid
@@ -32,17 +32,17 @@
 from sentry.auth.provider import MigratingIdentityId, Provider
 from sentry.auth.superuser import is_active_superuser
 from sentry.locks import locks
-from sentry.models import (
-    AuditLogEntry,
-    AuthIdentity,
-    AuthProvider,
-    Organization,
-    OrganizationMember,
-    OrganizationMemberTeam,
-    User,
-)
+from sentry.models import AuditLogEntry, AuthIdentity, AuthProvider, Organization, User
 from sentry.pipeline import Pipeline, PipelineSessionStore
 from sentry.pipeline.provider import PipelineProvider
+from sentry.services.hybrid_cloud.audit import AuditLogMetadata, audit_log_service
+from sentry.services.hybrid_cloud.auth import ApiAuthIdentity, auth_service
+from sentry.services.hybrid_cloud.organization import (
+    ApiOrganization,
+    ApiOrganizationMember,
+    organization_service,
+)
+from sentry.services.hybrid_cloud.organization.impl import DatabaseBackedOrganizationService
 from sentry.signals import sso_enabled, user_signup
 from sentry.tasks.auth import email_missing_links
 from sentry.utils import auth, json, metrics
@@ -98,10 +98,15 @@ class AuthIdentityHandler:
 
     auth_provider: AuthProvider
     provider: Provider
-    organization: Organization
+    organization: ApiOrganization
     request: HttpRequest
     identity: Mapping[str, Any]
 
+    def __post_init__(self) -> None:
+        # For debugging. TODO: Remove when tests are stable
+        if not isinstance(self.organization, ApiOrganization):
+            raise TypeError
+
     @cached_property
     def user(self) -> User | AnonymousUser:
         email = self.identity.get("email")
@@ -159,11 +164,12 @@ def _login(self, user: Any) -> None:
         )
 
     @staticmethod
-    def _set_linked_flag(member: OrganizationMember) -> None:
-        if getattr(member.flags, "sso:invalid") or not getattr(member.flags, "sso:linked"):
-            setattr(member.flags, "sso:invalid", False)
-            setattr(member.flags, "sso:linked", True)
-            member.save()
+    def _set_linked_flag(member: ApiOrganizationMember) -> None:
+        if member.flags.sso__invalid or not member.flags.sso__linked:
+            member.flags.sso__invalid = False
+            member.flags.sso__linked = True
+
+            organization_service.update_membership_flags(organization_member=member)
 
     def handle_existing_identity(
         self,
@@ -180,11 +186,10 @@ def handle_existing_identity(
             last_synced=now,
         )
 
-        try:
-            member = OrganizationMember.objects.get(
-                user=auth_identity.user, organization=self.organization
-            )
-        except OrganizationMember.DoesNotExist:
+        member = organization_service.check_membership_by_id(
+            organization_id=self.organization.id, user_id=auth_identity.user.id
+        )
+        if member is None:
             # this is likely the case when someone was removed from the org
             # but still has access to rejoin
             member = self._handle_new_membership(auth_identity)
@@ -222,54 +227,23 @@ def _get_login_redirect(self, subdomain: str | None) -> str:
             login_redirect_url = absolute_uri(login_redirect_url, url_prefix=url_prefix)
         return login_redirect_url
 
-    def _handle_new_membership(self, auth_identity: AuthIdentity) -> OrganizationMember | None:
-        user = auth_identity.user
-
-        # If the user is either currently *pending* invite acceptance (as indicated
-        # from the invite token and member id in the session) OR an existing invite exists on this
-        # organization for the email provided by the identity provider.
-        invite_helper = ApiInviteHelper.from_session_or_email(
-            request=self.request, organization=self.organization, email=user.email
-        )
-
-        # If we are able to accept an existing invite for the user for this
-        # organization, do so, otherwise handle new membership
-        if invite_helper:
-            if invite_helper.invite_approved:
-                return invite_helper.accept_invite(user)
-
-            # It's possible the user has an _invite request_ that hasn't been approved yet,
-            # and is able to join the organization without an invite through the SSO flow.
-            # In that case, delete the invite request and create a new membership.
-            invite_helper.handle_invite_not_approved()
-
-        flags = OrganizationMember.flags["sso:linked"]
-        # if the org doesn't have the ability to add members then anyone who got added
-        # this way should be disabled until the org upgrades
-        if not features.has("organizations:invite-members", self.organization):
-            flags = flags | OrganizationMember.flags["member-limit:restricted"]
-
-        # Otherwise create a new membership
-        om = OrganizationMember.objects.create(
-            organization=self.organization,
-            role=self.organization.default_role,
-            user=user,
-            flags=flags,
+    def _handle_new_membership(self, auth_identity: ApiAuthIdentity) -> ApiOrganizationMember:
+        user, om = auth_service.handle_new_membership(
+            self.request, self.organization, auth_identity, self.auth_provider
         )
 
-        default_teams = self.auth_provider.default_teams.all()
-        for team in default_teams:
-            OrganizationMemberTeam.objects.create(team=team, organizationmember=om)
-
-        AuditLogEntry.objects.create(
-            organization=self.organization,
-            actor=user,
-            ip_address=self.request.META["REMOTE_ADDR"],
-            target_object=om.id,
-            target_user=om.user,
-            event=audit_log.get_event_id("MEMBER_ADD"),
-            data=om.get_audit_log_data(),
-        )
+        if om is not None:
+            audit_log_service.log_organization_membership(
+                metadata=AuditLogMetadata(
+                    organization=self.organization,
+                    actor=user,
+                    ip_address=self.request.META["REMOTE_ADDR"],
+                    target_object=om.id,
+                    target_user=user,
+                    event=audit_log.get_event_id("MEMBER_ADD"),
+                ),
+                organization_member=om,
+            )
 
         return om
 
@@ -280,7 +254,7 @@ def _get_auth_identity(self, **params: Any) -> AuthIdentity | None:
             return None
 
     @transaction.atomic  # type: ignore
-    def handle_attach_identity(self, member: OrganizationMember | None = None) -> AuthIdentity:
+    def handle_attach_identity(self, member: ApiOrganizationMember | None = None) -> AuthIdentity:
         """
         Given an already authenticated user, attach or re-attach an identity.
         """
@@ -340,13 +314,15 @@ def handle_attach_identity(self, member: OrganizationMember | None = None) -> Au
         self._set_linked_flag(member)
 
         if auth_is_new:
-            AuditLogEntry.objects.create(
-                organization=self.organization,
-                actor=self.user,
-                ip_address=self.request.META["REMOTE_ADDR"],
-                target_object=auth_identity.id,
-                event=audit_log.get_event_id("SSO_IDENTITY_LINK"),
-                data=auth_identity.get_audit_log_data(),
+            audit_log_service.log_auth_identity(
+                metadata=AuditLogMetadata(
+                    organization=self.organization,
+                    actor=self.user,
+                    ip_address=self.request.META["REMOTE_ADDR"],
+                    target_object=auth_identity.id,
+                    event=audit_log.get_event_id("SSO_IDENTITY_LINK"),
+                ),
+                auth_identity=auth_identity,
             )
 
             messages.add_message(self.request, messages.SUCCESS, OK_LINK_IDENTITY)
@@ -366,27 +342,29 @@ def _wipe_existing_identity(self, auth_identity: AuthIdentity) -> Any:
 
         # since we've identified an identity which is no longer valid
         # lets preemptively mark it as such
-        try:
-            other_member = OrganizationMember.objects.get(
-                user=auth_identity.user_id, organization=self.organization
-            )
-        except OrganizationMember.DoesNotExist:
+        other_member = organization_service.check_membership_by_id(
+            user_id=auth_identity.user_id, organization_id=self.organization.id
+        )
+        if other_member is None:
             return
-        other_member.flags["sso:invalid"] = True
-        other_member.flags["sso:linked"] = False
-        other_member.save()
+
+        other_member.flags.sso__invalid = True
+        other_member.flags.sso__linked = False
+        organization_service.update_membership_flags(organization_member=other_member)
 
         return deletion_result
 
-    def _get_organization_member(self, auth_identity: AuthIdentity) -> OrganizationMember:
+    def _get_organization_member(self, auth_identity: AuthIdentity) -> ApiOrganizationMember:
         """
         Check to see if the user has a member associated, if not, create a new membership
         based on the auth_identity email.
         """
-        try:
-            return OrganizationMember.objects.get(user=self.user, organization=self.organization)
-        except OrganizationMember.DoesNotExist:
+        member = organization_service.check_membership_by_id(
+            organization_id=self.organization.id, user_id=self.user.id
+        )
+        if member is None:
             return self._handle_new_membership(auth_identity)
+        return member
 
     def _respond(
         self,
@@ -474,14 +452,14 @@ def handle_unknown_identity(
                 is_account_verified = self.has_verified_account(verification_value)
 
         is_new_account = not self.user.is_authenticated  # stateful
-        if self._app_user and self.identity.get("email_verified") or is_account_verified:
+        if self._app_user and (self.identity.get("email_verified") or is_account_verified):
             # we only allow this flow to happen if the existing user has
             # membership, otherwise we short circuit because it might be
             # an attempt to hijack membership of another organization
-            has_membership = OrganizationMember.objects.filter(
-                user=self._app_user, organization=self.organization
-            ).exists()
-            if has_membership:
+            membership = organization_service.check_membership_by_id(
+                user_id=self._app_user.id, organization_id=self.organization.id
+            )
+            if membership is not None:
                 try:
                     self._login(self.user)
                 except self._NotCompletedSecurityChecks:
@@ -750,8 +728,14 @@ def finish_pipeline(self) -> HttpResponseBase:
         return response
 
     def auth_handler(self, identity: Mapping[str, Any]) -> AuthIdentityHandler:
+        # This is a temporary step to keep test_helper integrated
+        # TODO: Move this conversion further upstream
+        api_organization = DatabaseBackedOrganizationService.serialize_organization(
+            self.organization
+        )
+
         return AuthIdentityHandler(
-            self.provider_model, self.provider, self.organization, self.request, identity
+            self.provider_model, self.provider, api_organization, self.request, identity
         )
 
     @transaction.atomic  # type: ignore
@@ -828,9 +812,10 @@ def _finish_setup_pipeline(self, identity: Mapping[str, Any]) -> HttpResponseRed
         data = self.fetch_state()
         config = self.provider.build_config(data)
 
-        try:
-            om = OrganizationMember.objects.get(user=request.user, organization=self.organization)
-        except OrganizationMember.DoesNotExist:
+        om = organization_service.check_membership_by_id(
+            organization_id=self.organization.id, user_id=request.user.id
+        )
+        if om is None:
             return self.error(ERR_UID_MISMATCH)
 
         # disable require 2FA for the organization

src/sentry/auth/idpmigration.py

@@ -8,7 +8,8 @@
 from rb.clients import LocalClient
 
 from sentry import options
-from sentry.models import AuthProvider, Organization, OrganizationMember, User
+from sentry.models import AuthProvider, User
+from sentry.services.hybrid_cloud.organization import ApiOrganization, organization_service
 from sentry.utils import json, metrics, redis
 from sentry.utils.email import MessageBuilder
 from sentry.utils.http import absolute_uri
@@ -21,7 +22,7 @@
 
 def send_one_time_account_confirm_link(
     user: User,
-    org: Organization,
+    org: ApiOrganization,
     provider: AuthProvider,
     email: str,
     identity_id: str,
@@ -34,7 +35,7 @@ def send_one_time_account_confirm_link(
     in an email to the associated address.
 
     :param user: the user profile to link
-    :param organization: the organization whose SSO provider is being used
+    :param org: the organization whose SSO provider is being used
     :param provider: the SSO provider
     :param email: the email address associated with the SSO identity
     :param identity_id: the SSO identity id
@@ -52,7 +53,7 @@ def get_redis_cluster() -> LocalClient:
 @dataclass
 class AccountConfirmLink:
     user: User
-    organization: Organization
+    organization: ApiOrganization
     provider: AuthProvider
     email: str
     identity_id: str
@@ -88,17 +89,14 @@ def send_confirm_email(self) -> None:
     def store_in_redis(self) -> None:
         cluster = get_redis_cluster()
 
-        try:
-            member_id = OrganizationMember.objects.get(
-                organization=self.organization, user=self.user
-            ).id
-        except OrganizationMember.DoesNotExist:
-            member_id = None
+        member = organization_service.check_membership_by_id(
+            organization_id=self.organization.id, user_id=self.user.id
+        )
 
         verification_value = {
             "user_id": self.user.id,
             "email": self.email,
-            "member_id": member_id,
+            "member_id": member.id if member is not None else None,
             "organization_id": self.organization.id,
             "identity_id": self.identity_id,
             "provider": self.provider.provider,