Upgrade for FastAPI Users V10

Author: frankie567

fastapi_users_db_sqlmodel/__init__.py

@@ -1,9 +1,9 @@
 """FastAPI Users database adapter for SQLModel."""
 import uuid
-from typing import Generic, Optional, Type, TypeVar
+from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type
 
 from fastapi_users.db.base import BaseUserDatabase
-from fastapi_users.models import BaseOAuthAccount, BaseUserDB
+from fastapi_users.models import ID, OAP, UP
 from pydantic import UUID4, EmailStr
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
@@ -12,13 +12,17 @@
 __version__ = "0.1.2"
 
 
-class SQLModelBaseUserDB(BaseUserDB, SQLModel):
+class SQLModelBaseUserDB(SQLModel):
     __tablename__ = "user"
 
     id: UUID4 = Field(default_factory=uuid.uuid4, primary_key=True, nullable=False)
-    email: EmailStr = Field(
-        sa_column_kwargs={"unique": True, "index": True}, nullable=False
-    )
+    if TYPE_CHECKING:  # pragma: no cover
+        email: str
+    else:
+        email: EmailStr = Field(
+            sa_column_kwargs={"unique": True, "index": True}, nullable=False
+        )
+    hashed_password: str
 
     is_active: bool = Field(True, nullable=False)
     is_superuser: bool = Field(False, nullable=False)
@@ -28,68 +32,59 @@ class Config:
         orm_mode = True
 
 
-class SQLModelBaseOAuthAccount(BaseOAuthAccount, SQLModel):
+class SQLModelBaseOAuthAccount(SQLModel):
     __tablename__ = "oauthaccount"
 
     id: UUID4 = Field(default_factory=uuid.uuid4, primary_key=True)
     user_id: UUID4 = Field(foreign_key="user.id", nullable=False)
+    oauth_name: str = Field(index=True, nullable=False)
+    access_token: str = Field(nullable=False)
+    expires_at: Optional[int] = Field(nullable=True)
+    refresh_token: Optional[str] = Field(nullable=True)
+    account_id: str = Field(index=True, nullable=False)
+    account_email: str = Field(nullable=False)
 
     class Config:
         orm_mode = True
 
 
-UD = TypeVar("UD", bound=SQLModelBaseUserDB)
-OA = TypeVar("OA", bound=SQLModelBaseOAuthAccount)
-
-
-class NotSetOAuthAccountTableError(Exception):
-    """
-    OAuth table was not set in DB adapter but was needed.
-
-    Raised when trying to create/update a user with OAuth accounts set
-    but no table were specified in the DB adapter.
-    """
-
-    pass
-
-
-class SQLModelUserDatabase(Generic[UD, OA], BaseUserDatabase[UD]):
+class SQLModelUserDatabase(Generic[UP, ID], BaseUserDatabase[UP, ID]):
     """
     Database adapter for SQLModel.
 
-    :param user_db_model: SQLModel model of a DB representation of a user.
     :param session: SQLAlchemy session.
     """
 
     session: Session
-    oauth_account_model: Optional[Type[OA]]
+    user_model: Type[UP]
+    oauth_account_model: Optional[Type[SQLModelBaseOAuthAccount]]
 
     def __init__(
         self,
-        user_db_model: Type[UD],
         session: Session,
-        oauth_account_model: Optional[Type[OA]] = None,
+        user_model: Type[UP],
+        oauth_account_model: Optional[Type[SQLModelBaseOAuthAccount]] = None,
     ):
-        super().__init__(user_db_model)
         self.session = session
+        self.user_model = user_model
         self.oauth_account_model = oauth_account_model
 
-    async def get(self, id: UUID4) -> Optional[UD]:
+    async def get(self, id: ID) -> Optional[UP]:
         """Get a single user by id."""
-        return self.session.get(self.user_db_model, id)
+        return self.session.get(self.user_model, id)
 
-    async def get_by_email(self, email: str) -> Optional[UD]:
+    async def get_by_email(self, email: str) -> Optional[UP]:
         """Get a single user by email."""
-        statement = select(self.user_db_model).where(
-            func.lower(self.user_db_model.email) == func.lower(email)
+        statement = select(self.user_model).where(
+            func.lower(self.user_model.email) == func.lower(email)
         )
         results = self.session.exec(statement)
         return results.first()
 
-    async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UD]:
+    async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UP]:
         """Get a single user by OAuth account id."""
-        if not self.oauth_account_model:
-            raise NotSetOAuthAccountTableError()
+        if self.oauth_account_model is None:
+            raise NotImplementedError()
         statement = (
             select(self.oauth_account_model)
             .where(self.oauth_account_model.oauth_name == oauth)
@@ -102,72 +97,93 @@ async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UD
             return user
         return None
 
-    async def create(self, user: UD) -> UD:
+    async def create(self, create_dict: Dict[str, Any]) -> UP:
         """Create a user."""
+        user = self.user_model(**create_dict)
         self.session.add(user)
-        if self.oauth_account_model is not None:
-            for oauth_account in user.oauth_accounts:  # type: ignore
-                self.session.add(oauth_account)
         self.session.commit()
         self.session.refresh(user)
         return user
 
-    async def update(self, user: UD) -> UD:
-        """Update a user."""
+    async def update(self, user: UP, update_dict: Dict[str, Any]) -> UP:
+        for key, value in update_dict.items():
+            setattr(user, key, value)
         self.session.add(user)
-        if self.oauth_account_model is not None:
-            for oauth_account in user.oauth_accounts:  # type: ignore
-                self.session.add(oauth_account)
         self.session.commit()
         self.session.refresh(user)
         return user
 
-    async def delete(self, user: UD) -> None:
-        """Delete a user."""
+    async def delete(self, user: UP) -> None:
         self.session.delete(user)
         self.session.commit()
 
+    async def add_oauth_account(self, user: UP, create_dict: Dict[str, Any]) -> UP:
+        if self.oauth_account_model is None:
+            raise NotImplementedError()
 
-class SQLModelUserDatabaseAsync(Generic[UD, OA], BaseUserDatabase[UD]):
+        oauth_account = self.oauth_account_model(**create_dict)
+        user.oauth_accounts.append(oauth_account)  # type: ignore
+        self.session.add(user)
+
+        self.session.commit()
+
+        return user
+
+    async def update_oauth_account(
+        self, user: UP, oauth_account: OAP, update_dict: Dict[str, Any]
+    ) -> UP:
+        if self.oauth_account_model is None:
+            raise NotImplementedError()
+
+        for key, value in update_dict.items():
+            setattr(oauth_account, key, value)
+        self.session.add(oauth_account)
+        self.session.commit()
+
+        return user
+
+
+class SQLModelUserDatabaseAsync(Generic[UP, ID], BaseUserDatabase[UP, ID]):
     """
     Database adapter for SQLModel working purely asynchronously.
 
-    :param user_db_model: SQLModel model of a DB representation of a user.
+    :param user_model: SQLModel model of a DB representation of a user.
     :param session: SQLAlchemy async session.
     """
 
     session: AsyncSession
-    oauth_account_model: Optional[Type[OA]]
+    user_model: Type[UP]
+    oauth_account_model: Optional[Type[SQLModelBaseOAuthAccount]]
 
     def __init__(
         self,
-        user_db_model: Type[UD],
         session: AsyncSession,
-        oauth_account_model: Optional[Type[OA]] = None,
+        user_model: Type[UP],
+        oauth_account_model: Optional[Type[SQLModelBaseOAuthAccount]] = None,
     ):
-        super().__init__(user_db_model)
         self.session = session
+        self.user_model = user_model
         self.oauth_account_model = oauth_account_model
 
-    async def get(self, id: UUID4) -> Optional[UD]:
+    async def get(self, id: ID) -> Optional[UP]:
         """Get a single user by id."""
-        return await self.session.get(self.user_db_model, id)
+        return await self.session.get(self.user_model, id)
 
-    async def get_by_email(self, email: str) -> Optional[UD]:
+    async def get_by_email(self, email: str) -> Optional[UP]:
         """Get a single user by email."""
-        statement = select(self.user_db_model).where(
-            func.lower(self.user_db_model.email) == func.lower(email)
+        statement = select(self.user_model).where(
+            func.lower(self.user_model.email) == func.lower(email)
         )
         results = await self.session.execute(statement)
         object = results.first()
         if object is None:
             return None
         return object[0]
 
-    async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UD]:
+    async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UP]:
         """Get a single user by OAuth account id."""
-        if not self.oauth_account_model:
-            raise NotSetOAuthAccountTableError()
+        if self.oauth_account_model is None:
+            raise NotImplementedError()
         statement = (
             select(self.oauth_account_model)
             .where(self.oauth_account_model.oauth_name == oauth)
@@ -177,31 +193,51 @@ async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UD
         results = await self.session.execute(statement)
         oauth_account = results.first()
         if oauth_account:
-            user = oauth_account[0].user
+            user = oauth_account[0].user  # type: ignore
             return user
         return None
 
-    async def create(self, user: UD) -> UD:
+    async def create(self, create_dict: Dict[str, Any]) -> UP:
         """Create a user."""
+        user = self.user_model(**create_dict)
         self.session.add(user)
-        if self.oauth_account_model is not None:
-            for oauth_account in user.oauth_accounts:  # type: ignore
-                self.session.add(oauth_account)
         await self.session.commit()
         await self.session.refresh(user)
         return user
 
-    async def update(self, user: UD) -> UD:
-        """Update a user."""
+    async def update(self, user: UP, update_dict: Dict[str, Any]) -> UP:
+        for key, value in update_dict.items():
+            setattr(user, key, value)
         self.session.add(user)
-        if self.oauth_account_model is not None:
-            for oauth_account in user.oauth_accounts:  # type: ignore
-                self.session.add(oauth_account)
         await self.session.commit()
         await self.session.refresh(user)
         return user
 
-    async def delete(self, user: UD) -> None:
-        """Delete a user."""
+    async def delete(self, user: UP) -> None:
         await self.session.delete(user)
         await self.session.commit()
+
+    async def add_oauth_account(self, user: UP, create_dict: Dict[str, Any]) -> UP:
+        if self.oauth_account_model is None:
+            raise NotImplementedError()
+
+        oauth_account = self.oauth_account_model(**create_dict)
+        user.oauth_accounts.append(oauth_account)  # type: ignore
+        self.session.add(user)
+
+        await self.session.commit()
+
+        return user
+
+    async def update_oauth_account(
+        self, user: UP, oauth_account: OAP, update_dict: Dict[str, Any]
+    ) -> UP:
+        if self.oauth_account_model is None:
+            raise NotImplementedError()
+
+        for key, value in update_dict.items():
+            setattr(oauth_account, key, value)
+        self.session.add(oauth_account)
+        await self.session.commit()
+
+        return user