server/user: fix OAuth account disconnection when users have multiple ones for the same platform

Author: frankie567

server/polar/user/endpoints.py

@@ -6,7 +6,6 @@
 from polar.customer_portal.endpoints.license_keys import router as license_keys_router
 from polar.customer_portal.endpoints.order import router as order_router
 from polar.customer_portal.endpoints.subscription import router as subscription_router
-from polar.exceptions import PolarError
 from polar.models import User
 from polar.models.user import OAuthPlatform
 from polar.openapi import APITag
@@ -48,22 +47,6 @@ async def create_identity_verification(
     )
 
 
-class OAuthAccountNotFound(PolarError):
-    def __init__(self, platform: OAuthPlatform) -> None:
-        self.platform = platform
-        message = f"No {platform} OAuth account found for this user."
-        super().__init__(message, 404)
-
-
-class CannotDisconnectLastAuthMethod(PolarError):
-    def __init__(self) -> None:
-        message = (
-            "Cannot disconnect this OAuth account as it's your only authentication method. "
-            "Please verify your email or connect another OAuth provider before disconnecting."
-        )
-        super().__init__(message, 400)
-
-
 @router.delete(
     "/me/oauth-accounts/{platform}",
     status_code=204,
@@ -86,21 +69,4 @@ async def disconnect_oauth_account(
     Note: You cannot disconnect your last authentication method if your email is not verified.
     """
     user = auth_subject.subject
-
-    oauth_account = await oauth_account_service.get_by_platform_and_user_id(
-        session, platform, user.id
-    )
-
-    if oauth_account is None:
-        raise OAuthAccountNotFound(platform)
-
-    can_disconnect = await oauth_account_service.can_disconnect_oauth_account(
-        session, user, oauth_account.id
-    )
-
-    if not can_disconnect:
-        raise CannotDisconnectLastAuthMethod()
-
-    await oauth_account_service.disconnect_oauth_account(
-        session, user, oauth_account.id, platform
-    )
+    await oauth_account_service.disconnect_platform(session, user, platform)

server/polar/user/oauth_service.py

@@ -1,8 +1,7 @@
-from uuid import UUID
-
 import structlog
-from sqlalchemy import delete, func, select
+from sqlalchemy import func, select
 
+from polar.exceptions import PolarError
 from polar.kit.services import ResourceServiceReader
 from polar.logging import Logger
 from polar.models import OAuthAccount, User
@@ -12,6 +11,25 @@
 log: Logger = structlog.get_logger()
 
 
+class OAuthError(PolarError): ...
+
+
+class OAuthAccountNotFound(OAuthError):
+    def __init__(self, platform: OAuthPlatform) -> None:
+        self.platform = platform
+        message = f"No {platform} OAuth account found for this user."
+        super().__init__(message, 404)
+
+
+class CannotDisconnectLastAuthMethod(OAuthError):
+    def __init__(self) -> None:
+        message = (
+            "Cannot disconnect this OAuth account as it's your only authentication method. "
+            "Please verify your email or connect another OAuth provider before disconnecting."
+        )
+        super().__init__(message, 400)
+
+
 class OAuthAccountService(ResourceServiceReader[OAuthAccount]):
     async def get_by_platform_and_account_id(
         self, session: AsyncSession, platform: OAuthPlatform, account_id: str
@@ -23,44 +41,40 @@ async def get_by_platform_and_account_id(
         result = await session.execute(stmt)
         return result.scalars().one_or_none()
 
-    async def get_by_platform_and_user_id(
-        self, session: AsyncSession, platform: OAuthPlatform, user_id: UUID
-    ) -> OAuthAccount | None:
-        stmt = select(OAuthAccount).where(
+    async def disconnect_platform(
+        self, session: AsyncSession, user: User, platform: OAuthPlatform
+    ) -> None:
+        oauth_accounts_statement = select(OAuthAccount).where(
             OAuthAccount.platform == platform,
-            OAuthAccount.user_id == user_id,
+            OAuthAccount.user_id == user.id,
         )
-        result = await session.execute(stmt)
-        return result.scalars().one_or_none()
+        oauth_account_result = await session.execute(oauth_accounts_statement)
+        # Some users have a buggy state with multiple OAuth accounts for the same platform
+        oauth_accounts = oauth_account_result.scalars().all()
 
-    async def can_disconnect_oauth_account(
-        self, session: AsyncSession, user: User, oauth_account_id: UUID
-    ) -> bool:
-        stmt = select(func.count(OAuthAccount.id)).where(
+        if len(oauth_accounts) == 0:
+            raise OAuthAccountNotFound(platform)
+
+        other_accounts_count_statement = select(func.count(OAuthAccount.id)).where(
             OAuthAccount.user_id == user.id,
-            OAuthAccount.id != oauth_account_id,
+            OAuthAccount.id.not_in([oa.id for oa in oauth_accounts]),
+        )
+        other_accounts_count_result = await session.execute(
+            other_accounts_count_statement
         )
-        active_oauth_count = await session.scalar(stmt)
+        other_accounts_count = other_accounts_count_result.scalar_one()
 
-        if active_oauth_count == 0 and not user.email_verified:
-            return False
+        if other_accounts_count == 0 and not user.email_verified:
+            raise CannotDisconnectLastAuthMethod()
 
-        return True
+        for oauth_account in oauth_accounts:
+            await session.delete(oauth_account)
+            log.info(
+                "oauth_account.disconnect",
+                oauth_account_id=oauth_account.id,
+                platform=platform,
+            )
 
-    async def disconnect_oauth_account(
-        self,
-        session: AsyncSession,
-        user: User,
-        oauth_account_id: UUID,
-        platform: OAuthPlatform,
-    ) -> None:
-        log.info(
-            "oauth_account.disconnect",
-            oauth_account_id=oauth_account_id,
-            platform=platform,
-        )
-        stmt = delete(OAuthAccount).where(OAuthAccount.id == oauth_account_id)
-        await session.execute(stmt)
         await session.flush()