diff --git a/fastapi_users/db/base.py b/fastapi_users/db/base.py index 9c31844e..6ac78730 100644 --- a/fastapi_users/db/base.py +++ b/fastapi_users/db/base.py @@ -2,7 +2,7 @@ from typing import Any, Dict, Generic, Optional 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 @@ class BaseUserDatabase(Generic[UP]): """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]] diff --git a/fastapi_users/manager.py b/fastapi_users/manager.py index df8186f0..77847561 100644 --- a/fastapi_users/manager.py +++ b/fastapi_users/manager.py @@ -171,7 +171,12 @@ class BaseUserManager(Generic[models.UP]): 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 @@ class BaseUserManager(Generic[models.UP]): 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 diff --git a/fastapi_users/router/oauth.py b/fastapi_users/router/oauth.py index 74bb3fb9..00499d2b 100644 --- a/fastapi_users/router/oauth.py +++ b/fastapi_users/router/oauth.py @@ -6,7 +6,7 @@ from httpx_oauth.integrations.fastapi import OAuth2AuthorizeCallback 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 @@ def get_oauth_router( 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, diff --git a/tests/conftest.py b/tests/conftest.py index a756a3bb..c995b13b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -430,6 +430,32 @@ def mock_user_db_oauth( 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() diff --git a/tests/test_db_base.py b/tests/test_db_base.py index 5cfea0d2..82501d2f 100644 --- a/tests/test_db_base.py +++ b/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("lancelot@camelot.bt") @@ -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, {}) diff --git a/tests/test_manager.py b/tests/test_manager.py index fac99a0a..cc1360d9 100644 --- a/tests/test_manager.py +++ b/tests/test_manager.py @@ -173,14 +173,18 @@ class TestOAuthCallback: 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 @@ class TestOAuthCallback: 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="galahad@camelot.bt", + user = await user_manager_oauth.oauth_callback( + "service1", "TOKEN", "new_user_oauth1", "galahad@camelot.bt", 1579000751 ) - user = await user_manager_oauth.oauth_callback(oauth_account) - assert user.email == "galahad@camelot.bt" 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