feat: add refresh token classes, fix base models not inherited from protocols

Author: Ae-Mc

fastapi_users_db_sqlmodel/__init__.py

@@ -1,22 +1,32 @@
 """FastAPI Users database adapter for SQLModel."""
 
 import uuid
-from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type
+from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, _ProtocolMeta
 
 from fastapi_users.db.base import BaseUserDatabase
-from fastapi_users.models import ID, OAP, UP
+from fastapi_users.models import (
+    ID,
+    UP,
+    OAuthAccountProtocol,
+    UserProtocol,
+)
 from pydantic import UUID4, ConfigDict, EmailStr
 from pydantic.version import VERSION as PYDANTIC_VERSION
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
 from sqlmodel import AutoString, Field, Session, SQLModel, func, select
+from sqlmodel.main import SQLModelMetaclass
 
 __version__ = "0.3.0"
 PYDANTIC_V2 = PYDANTIC_VERSION.startswith("2.")
 
 
-class SQLModelBaseUserDB(SQLModel):
-    __tablename__ = "user"
+class SQLModelProtocolMetaclass(SQLModelMetaclass, _ProtocolMeta):
+    pass
+
+
+class SQLModelBaseUserDB(SQLModel, UserProtocol, metaclass=SQLModelProtocolMetaclass):
+    __tablename__ = "user"  # type: ignore
 
     id: UUID4 = Field(default_factory=uuid.uuid4, primary_key=True, nullable=False)
     if TYPE_CHECKING:  # pragma: no cover
@@ -41,8 +51,10 @@ class Config:
             orm_mode = True
 
 
-class SQLModelBaseOAuthAccount(SQLModel):
-    __tablename__ = "oauthaccount"
+class SQLModelBaseOAuthAccount(
+    SQLModel, OAuthAccountProtocol, metaclass=SQLModelProtocolMetaclass
+):
+    __tablename__ = "oauthaccount"  # type: ignore
 
     id: UUID4 = Field(default_factory=uuid.uuid4, primary_key=True)
     user_id: UUID4 = Field(foreign_key="user.id", nullable=False)
@@ -54,10 +66,11 @@ class SQLModelBaseOAuthAccount(SQLModel):
     account_email: str = Field(nullable=False)
 
     if PYDANTIC_V2:
+        # pragma: no cover
         model_config = ConfigDict(from_attributes=True)  # type: ignore
     else:
 
-        class Config:
+        class Config:  # pragma: no cover
             orm_mode = True
 
 
@@ -143,7 +156,7 @@ async def add_oauth_account(self, user: UP, create_dict: Dict[str, Any]) -> UP:
         return user
 
     async def update_oauth_account(
-        self, user: UP, oauth_account: OAP, update_dict: Dict[str, Any]
+        self, user: UP, oauth_account: OAuthAccountProtocol, update_dict: Dict[str, Any]
     ) -> UP:
         if self.oauth_account_model is None:
             raise NotImplementedError()
@@ -243,7 +256,7 @@ async def add_oauth_account(self, user: UP, create_dict: Dict[str, Any]) -> UP:
         return user
 
     async def update_oauth_account(
-        self, user: UP, oauth_account: OAP, update_dict: Dict[str, Any]
+        self, user: UP, oauth_account: OAuthAccountProtocol, update_dict: Dict[str, Any]
     ) -> UP:
         if self.oauth_account_model is None:
             raise NotImplementedError()

fastapi_users_db_sqlmodel/access_token.py

@@ -1,7 +1,17 @@
 from datetime import datetime
 from typing import Any, Dict, Generic, Optional, Type
 
-from fastapi_users.authentication.strategy.db import AP, AccessTokenDatabase
+from fastapi_users.authentication.strategy.db import (
+    AP,
+    APE,
+    AccessRefreshTokenDatabase,
+    AccessTokenDatabase,
+)
+from fastapi_users.authentication.strategy.db.adapter import BaseAccessTokenDatabase
+from fastapi_users.authentication.strategy.db.models import (
+    AccessRefreshTokenProtocol,
+    AccessTokenProtocol,
+)
 from pydantic import UUID4, ConfigDict
 from pydantic.version import VERSION as PYDANTIC_VERSION
 from sqlalchemy import types
@@ -10,14 +20,23 @@
 
 from fastapi_users_db_sqlmodel.generics import TIMESTAMPAware, now_utc
 
+from . import SQLModelProtocolMetaclass
+
 PYDANTIC_V2 = PYDANTIC_VERSION.startswith("2.")
-class SQLModelBaseAccessToken(SQLModel):
-    __tablename__ = "accesstoken"
 
-    token: str = Field(sa_type=types.String(length=43), primary_key=True)
+
+class SQLModelBaseAccessToken(
+    SQLModel, AccessTokenProtocol, metaclass=SQLModelProtocolMetaclass
+):
+    __tablename__ = "accesstoken"  # type: ignore
+
+    token: str = Field(
+        sa_type=types.String(length=43),  # type: ignore
+        primary_key=True,
+    )
     created_at: datetime = Field(
         default_factory=now_utc,
-        sa_type=TIMESTAMPAware(timezone=True),
+        sa_type=TIMESTAMPAware(timezone=True),  # type: ignore
         nullable=False,
         index=True,
     )
@@ -26,11 +45,26 @@ class SQLModelBaseAccessToken(SQLModel):
     if PYDANTIC_V2:  # pragma: no cover
         model_config = ConfigDict(from_attributes=True)  # type: ignore
     else:  # pragma: no cover
+
         class Config:
             orm_mode = True
 
 
-class SQLModelAccessTokenDatabase(Generic[AP], AccessTokenDatabase[AP]):
+class SQLModelBaseAccessRefreshToken(
+    SQLModelBaseAccessToken,
+    AccessRefreshTokenProtocol,
+    metaclass=SQLModelProtocolMetaclass,
+):
+    __tablename__ = "accessrefreshtoken"
+
+    refresh_token: str = Field(
+        sa_type=types.String(length=43),  # type: ignore
+        unique=True,
+        index=True,
+    )
+
+
+class BaseSQLModelAccessTokenDatabase(Generic[AP], BaseAccessTokenDatabase[str, AP]):
     """
     Access token database adapter for SQLModel.
 
@@ -77,7 +111,47 @@ async def delete(self, access_token: AP) -> None:
         self.session.commit()
 
 
-class SQLModelAccessTokenDatabaseAsync(Generic[AP], AccessTokenDatabase[AP]):
+class SQLModelAccessTokenDatabase(
+    Generic[AP], BaseSQLModelAccessTokenDatabase[AP], AccessTokenDatabase[AP]
+):
+    """
+    Access token database adapter for SQLModel.
+
+    :param session: SQLAlchemy session.
+    :param access_token_model: SQLModel access token model.
+    """
+
+
+class SQLModelAccessRefreshTokenDatabase(
+    Generic[APE], BaseSQLModelAccessTokenDatabase[APE], AccessRefreshTokenDatabase[APE]
+):
+    """
+    Access token database adapter for SQLModel.
+
+    :param session: SQLAlchemy session.
+    :param access_token_model: SQLModel access refresh token model.
+    """
+
+    async def get_by_refresh_token(
+        self, refresh_token: str, max_age: Optional[datetime] = None
+    ) -> Optional[APE]:
+        statement = select(self.access_token_model).where(  # type: ignore
+            self.access_token_model.refresh_token == refresh_token
+        )
+        if max_age is not None:
+            statement = statement.where(self.access_token_model.created_at >= max_age)
+
+        results = self.session.exec(statement)
+        access_token = results.first()
+        if access_token is None:
+            return None
+
+        return access_token
+
+
+class BaseSQLModelAccessTokenDatabaseAsync(
+    Generic[AP], BaseAccessTokenDatabase[str, AP]
+):
     """
     Access token database adapter for SQLModel working purely asynchronously.
 
@@ -122,3 +196,31 @@ async def update(self, access_token: AP, update_dict: Dict[str, Any]) -> AP:
     async def delete(self, access_token: AP) -> None:
         await self.session.delete(access_token)
         await self.session.commit()
+
+
+class SQLModelAccessTokenDatabaseAsync(
+    BaseSQLModelAccessTokenDatabaseAsync[AP], AccessTokenDatabase[AP], Generic[AP]
+):
+    pass
+
+
+class SQLModelAccessRefreshTokenDatabaseAsync(
+    BaseSQLModelAccessTokenDatabaseAsync[APE],
+    AccessRefreshTokenDatabase[APE],
+    Generic[APE],
+):
+    async def get_by_refresh_token(
+        self, refresh_token: str, max_age: Optional[datetime] = None
+    ) -> Optional[APE]:
+        statement = select(self.access_token_model).where(  # type: ignore
+            self.access_token_model.refresh_token == refresh_token
+        )
+        if max_age is not None:
+            statement = statement.where(self.access_token_model.created_at >= max_age)
+
+        results = await self.session.execute(statement)
+        access_token = results.first()
+        if access_token is None:
+            return None
+
+        return access_token[0]