Revamp OAuth account model management

Author: frankie567

fastapi_users/db/base.py

@@ -2,7 +2,7 @@
 
 from pydantic import UUID4
 
-from fastapi_users.models import UP
+from fastapi_users.models import OAP, UP
 from fastapi_users.types import DependencyCallable
 
 
@@ -33,5 +33,15 @@ async def delete(self, user: UP) -> None:
         """Delete a user."""
         raise NotImplementedError()
 
+    async def add_oauth_account(self, user: UP, create_dict: Dict[str, Any]) -> UP:
+        """Create an OAuth account and add it to the user."""
+        raise NotImplementedError()
+
+    async def update_oauth_account(
+        self, user: UP, oauth_account: OAP, update_dict: Dict[str, Any]
+    ) -> UP:
+        """Update an OAuth account on a user."""
+        raise NotImplementedError()
+
 
 UserDatabaseDependency = DependencyCallable[BaseUserDatabase[UP]]

fastapi_users/manager.py

@@ -171,7 +171,12 @@ async def create(
 
     async def oauth_callback(
         self: "BaseUserManager[models.UOAP]",
-        oauth_account: models.OAP,
+        oauth_name: str,
+        access_token: str,
+        account_id: str,
+        account_email: str,
+        expires_at: Optional[int] = None,
+        refresh_token: Optional[str] = None,
         request: Optional[Request] = None,
     ) -> models.UOAP:
         """
@@ -185,44 +190,53 @@ async def oauth_callback(
         If the user does not exist, it is created and the on_after_register handler
         is triggered.
 
-        :param oauth_account: The new OAuth account to create.
+        :param oauth_name: Name of the OAuth client.
+        :param access_token: Valid access token for the service provider.
+        :param account_id: ID of the user on the service provider.
+        :param account_email: E-mail of the user on the service provider.
+        :param expires_at: Optional timestamp at which the access token expires.
+        :param refresh_token: Optional refresh token to get a
+        fresh access token from the service provider.
         :param request: Optional FastAPI request that
         triggered the operation, defaults to None
         :return: A user.
         """
+        oauth_account_dict = {
+            "oauth_name": oauth_name,
+            "access_token": access_token,
+            "account_id": account_id,
+            "account_email": account_email,
+            "expires_at": expires_at,
+            "refresh_token": refresh_token,
+        }
+
         try:
-            user = await self.get_by_oauth_account(
-                oauth_account.oauth_name, oauth_account.account_id
-            )
+            user = await self.get_by_oauth_account(oauth_name, account_id)
         except UserNotExists:
             try:
                 # Link account
-                user = await self.get_by_email(oauth_account.account_email)
-                oauth_accounts = [*user.oauth_accounts, oauth_account]
-                await self.user_db.update(user, {"oauth_accounts": oauth_accounts})
+                user = await self.get_by_email(account_email)
+                user = await self.user_db.add_oauth_account(user, oauth_account_dict)
             except UserNotExists:
                 # Create account
                 password = self.password_helper.generate()
                 user_dict = {
-                    "email": oauth_account.account_email,
+                    "email": account_email,
                     "hashed_password": self.password_helper.hash(password),
-                    "oauth_accounts": [oauth_account],
                 }
                 user = await self.user_db.create(user_dict)
+                user = await self.user_db.add_oauth_account(user, oauth_account_dict)
                 await self.on_after_register(user, request)
         else:
             # Update oauth
-            updated_oauth_accounts = []
-            for existing_oauth_account in user.oauth_accounts:  # type: ignore
+            for existing_oauth_account in user.oauth_accounts:
                 if (
-                    existing_oauth_account.account_id == oauth_account.account_id
-                    and existing_oauth_account.oauth_name == oauth_account.oauth_name
+                    existing_oauth_account.account_id == account_id
+                    and existing_oauth_account.oauth_name == oauth_name
                 ):
-                    oauth_account.id = existing_oauth_account.id
-                    updated_oauth_accounts.append(oauth_account)
-                else:
-                    updated_oauth_accounts.append(existing_oauth_account)
-            await self.user_db.update(user, {"oauth_accounts": updated_oauth_accounts})
+                    user = await self.user_db.update_oauth_account(
+                        user, existing_oauth_account, oauth_account_dict
+                    )
 
         return user
 

fastapi_users/router/oauth.py

@@ -6,7 +6,7 @@
 from httpx_oauth.oauth2 import BaseOAuth2, OAuth2Token
 from pydantic import BaseModel
 
-from fastapi_users import models, schemas
+from fastapi_users import models
 from fastapi_users.authentication import AuthenticationBackend, Strategy
 from fastapi_users.jwt import SecretType, decode_jwt, generate_jwt
 from fastapi_users.manager import BaseUserManager, UserManagerDependency
@@ -114,17 +114,16 @@ async def callback(
         except jwt.DecodeError:
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
 
-        new_oauth_account = schemas.BaseOAuthAccount(
-            oauth_name=oauth_client.name,
-            access_token=token["access_token"],
-            expires_at=token.get("expires_at"),
-            refresh_token=token.get("refresh_token"),
-            account_id=account_id,
-            account_email=account_email,
+        user = await user_manager.oauth_callback(
+            oauth_client.name,
+            token["access_token"],
+            account_id,
+            account_email,
+            token.get("expires_at"),
+            token.get("refresh_token"),
+            request,
         )
 
-        user = await user_manager.oauth_callback(new_oauth_account, request)
-
         if not user.is_active:
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,

tests/conftest.py

@@ -430,6 +430,32 @@ async def update(
         async def delete(self, user: UserOAuthModel) -> None:
             pass
 
+        async def add_oauth_account(
+            self, user: UserOAuthModel, create_dict: Dict[str, Any]
+        ) -> UserOAuthModel:
+            oauth_account = OAuthAccountModel(**create_dict)
+            user.oauth_accounts.append(oauth_account)
+            return user
+
+        async def update_oauth_account(  # type: ignore
+            self,
+            user: UserOAuthModel,
+            oauth_account: OAuthAccountModel,
+            update_dict: Dict[str, Any],
+        ) -> UserOAuthModel:
+            for field, value in update_dict.items():
+                setattr(oauth_account, field, value)
+            updated_oauth_accounts = []
+            for existing_oauth_account in user.oauth_accounts:
+                if (
+                    existing_oauth_account.account_id == oauth_account.account_id
+                    and existing_oauth_account.oauth_name == oauth_account.oauth_name
+                ):
+                    updated_oauth_accounts.append(oauth_account)
+                else:
+                    updated_oauth_accounts.append(existing_oauth_account)
+            return user
+
     return MockUserDatabase()
 
 

tests/test_db_base.py

@@ -1,15 +1,20 @@
+import uuid
+
 import pytest
 
 from fastapi_users.db import BaseUserDatabase
+from tests.conftest import OAuthAccountModel, UserModel
 
 
 @pytest.mark.asyncio
 @pytest.mark.db
-async def test_not_implemented_methods(user):
-    base_user_db = BaseUserDatabase()
+async def test_not_implemented_methods(
+    user: UserModel, oauth_account1: OAuthAccountModel
+):
+    base_user_db = BaseUserDatabase[UserModel]()
 
     with pytest.raises(NotImplementedError):
-        await base_user_db.get("aaa")
+        await base_user_db.get(uuid.uuid4())
 
     with pytest.raises(NotImplementedError):
         await base_user_db.get_by_email("[email protected]")
@@ -25,3 +30,9 @@ async def test_not_implemented_methods(user):
 
     with pytest.raises(NotImplementedError):
         await base_user_db.delete(user)
+
+    with pytest.raises(NotImplementedError):
+        await base_user_db.add_oauth_account(user, {})
+
+    with pytest.raises(NotImplementedError):
+        await base_user_db.update_oauth_account(user, oauth_account1, {})

tests/test_manager.py

@@ -173,14 +173,18 @@ async def test_existing_user_with_oauth(
         user_manager_oauth: UserManagerMock[UserOAuthModel],
         user_oauth: UserOAuthModel,
     ):
-        oauth_account = copy.deepcopy(user_oauth.oauth_accounts[0])
-        oauth_account.id = uuid.uuid4()
-        oauth_account.access_token = "UPDATED_TOKEN"
+        oauth_account = user_oauth.oauth_accounts[0]
 
-        user = await user_manager_oauth.oauth_callback(oauth_account)
+        user = await user_manager_oauth.oauth_callback(
+            oauth_account.oauth_name,
+            "UPDATED_TOKEN",
+            oauth_account.account_id,
+            oauth_account.account_email,
+        )
 
         assert user.id == user_oauth.id
         assert len(user.oauth_accounts) == 2
+        assert user.oauth_accounts[0].id == oauth_account.id
         assert user.oauth_accounts[0].oauth_name == "service1"
         assert user.oauth_accounts[0].access_token == "UPDATED_TOKEN"
         assert user.oauth_accounts[1].access_token == "TOKEN"
@@ -193,36 +197,24 @@ async def test_existing_user_without_oauth(
         user_manager_oauth: UserManagerMock[UserOAuthModel],
         superuser_oauth: UserOAuthModel,
     ):
-        oauth_account = OAuthAccountModel(
-            oauth_name="service1",
-            access_token="TOKEN",
-            expires_at=1579000751,
-            account_id="superuser_oauth1",
-            account_email=superuser_oauth.email,
+        user = await user_manager_oauth.oauth_callback(
+            "service1", "TOKEN", "superuser_oauth1", superuser_oauth.email, 1579000751
         )
 
-        user = await user_manager_oauth.oauth_callback(oauth_account)
-
         assert user.id == superuser_oauth.id
         assert len(user.oauth_accounts) == 1
-        assert user.oauth_accounts[0].id == oauth_account.id
+        assert user.oauth_accounts[0].id is not None
 
         assert user_manager_oauth.on_after_register.called is False
 
     async def test_new_user(self, user_manager_oauth: UserManagerMock[UserOAuthModel]):
-        oauth_account = OAuthAccountModel(
-            oauth_name="service1",
-            access_token="TOKEN",
-            expires_at=1579000751,
-            account_id="new_user_oauth1",
-            account_email="[email protected]",
+        user = await user_manager_oauth.oauth_callback(
+            "service1", "TOKEN", "new_user_oauth1", "[email protected]", 1579000751
         )
 
-        user = await user_manager_oauth.oauth_callback(oauth_account)
-
         assert user.email == "[email protected]"
         assert len(user.oauth_accounts) == 1
-        assert user.oauth_accounts[0].id == oauth_account.id
+        assert user.oauth_accounts[0].id is not None
 
         assert user_manager_oauth.on_after_register.called is True