Implement access token strategy db adapter

Author: frankie567

fastapi_users_db_sqlmodel/__init__.py

@@ -13,6 +13,8 @@
 
 
 class SQLModelBaseUserDB(BaseUserDB, 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
@@ -27,7 +29,10 @@ class Config:
 
 
 class SQLModelBaseOAuthAccount(BaseOAuthAccount, SQLModel):
+    __tablename__ = "oauthaccount"
+
     id: UUID4 = Field(default_factory=uuid.uuid4, primary_key=True)
+    user_id: UUID4 = Field(foreign_key="user.id", nullable=False)
 
     class Config:
         orm_mode = True
@@ -128,7 +133,7 @@ class SQLModelUserDatabaseAsync(Generic[UD, OA], BaseUserDatabase[UD]):
     Database adapter for SQLModel working purely asynchronously.
 
     :param user_db_model: SQLModel model of a DB representation of a user.
-    :param engine: SQLAlchemy async engine.
+    :param session: SQLAlchemy async session.
     """
 
     session: AsyncSession

fastapi_users_db_sqlmodel/access_token.py

@@ -0,0 +1,114 @@
+from datetime import datetime, timezone
+from typing import Generic, Optional, Type, TypeVar
+
+from fastapi_users.authentication.strategy.db import AccessTokenDatabase
+from fastapi_users.authentication.strategy.db.models import BaseAccessToken
+from pydantic import UUID4
+from sqlalchemy import Column, types
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlmodel import Field, Session, SQLModel, select
+
+
+def now_utc():
+    return datetime.now(timezone.utc)
+
+
+class SQLModelBaseAccessToken(BaseAccessToken, SQLModel):
+    __tablename__ = "accesstoken"
+
+    token: str = Field(
+        sa_column=Column("token", types.String(length=43), primary_key=True)
+    )
+    created_at: datetime = Field(default_factory=now_utc, nullable=False)
+    user_id: UUID4 = Field(foreign_key="user.id", nullable=False)
+
+    class Config:
+        orm_mode = True
+
+
+A = TypeVar("A", bound=SQLModelBaseAccessToken)
+
+
+class SQLModelAccessTokenDatabase(Generic[A], AccessTokenDatabase[A]):
+    """
+    Access token database adapter for SQLModel.
+
+    :param user_db_model: SQLModel model of a DB representation of an access token.
+    :param session: SQLAlchemy session.
+    """
+
+    def __init__(self, access_token_model: Type[A], session: Session):
+        self.access_token_model = access_token_model
+        self.session = session
+
+    async def get_by_token(
+        self, token: str, max_age: Optional[datetime] = None
+    ) -> Optional[A]:
+        statement = select(self.access_token_model).where(
+            self.access_token_model.token == token
+        )
+        if max_age is not None:
+            statement = statement.where(self.access_token_model.created_at >= max_age)
+
+        results = self.session.exec(statement)
+        return results.first()
+
+    async def create(self, access_token: A) -> A:
+        self.session.add(access_token)
+        self.session.commit()
+        self.session.refresh(access_token)
+        return access_token
+
+    async def update(self, access_token: A) -> A:
+        self.session.add(access_token)
+        self.session.commit()
+        self.session.refresh(access_token)
+        return access_token
+
+    async def delete(self, access_token: A) -> None:
+        self.session.delete(access_token)
+        self.session.commit()
+
+
+class SQLModelAccessTokenDatabaseAsync(Generic[A], AccessTokenDatabase[A]):
+    """
+    Access token database adapter for SQLModel working purely asynchronously.
+
+    :param user_db_model: SQLModel model of a DB representation of an access token.
+    :param session: SQLAlchemy async session.
+    """
+
+    def __init__(self, access_token_model: Type[A], session: AsyncSession):
+        self.access_token_model = access_token_model
+        self.session = session
+
+    async def get_by_token(
+        self, token: str, max_age: Optional[datetime] = None
+    ) -> Optional[A]:
+        statement = select(self.access_token_model).where(
+            self.access_token_model.token == token
+        )
+        if max_age is not None:
+            statement = statement.where(self.access_token_model.created_at >= max_age)
+
+        results = await self.session.execute(statement)
+        object = results.first()
+        if object is None:
+            return None
+        return object[0]
+
+    async def create(self, access_token: A) -> A:
+        self.session.add(access_token)
+        await self.session.commit()
+        await self.session.refresh(access_token)
+        return access_token
+
+    async def update(self, access_token: A) -> A:
+        self.session.add(access_token)
+        await self.session.commit()
+        await self.session.refresh(access_token)
+        return access_token
+
+    async def delete(self, access_token: A) -> None:
+        await self.session.delete(access_token)
+        await self.session.commit()

pyproject.toml

@@ -22,7 +22,7 @@ classifiers = [
 description-file = "README.md"
 requires-python = ">=3.7"
 requires = [
-    "fastapi-users >= 7.0.0",
+    "fastapi-users >= 9.1.0",
     "sqlmodel",
 ]
 

requirements.txt

@@ -1,3 +1,3 @@
 aiosqlite >= 0.17.0
-fastapi-users >= 8.1.1
+fastapi-users >= 9.1.0
 sqlmodel

tests/conftest.py

@@ -1,4 +1,5 @@
 import asyncio
+import uuid
 from typing import List, Optional
 
 import pytest
@@ -31,15 +32,15 @@ class UserOAuth(User):
 
 
 class UserDBOAuth(SQLModelBaseUserDB, table=True):
-    __tablename__ = "user"
+    __tablename__ = "user_oauth"
     oauth_accounts: List["OAuthAccount"] = Relationship(
         back_populates="user",
         sa_relationship_kwargs={"lazy": "joined", "cascade": "all, delete"},
     )
 
 
 class OAuthAccount(SQLModelBaseOAuthAccount, table=True):
-    user_id: UUID4 = Field(foreign_key="user.id")
+    user_id: UUID4 = Field(foreign_key="user_oauth.id")
     user: Optional[UserDBOAuth] = Relationship(back_populates="oauth_accounts")
 
 
@@ -53,6 +54,7 @@ def event_loop():
 @pytest.fixture
 def oauth_account1() -> OAuthAccount:
     return OAuthAccount(
+        id=uuid.UUID("b9089e5d-2642-406d-a7c0-cbc641aca0ec"),
         oauth_name="service1",
         access_token="TOKEN",
         expires_at=1579000751,
@@ -64,6 +66,7 @@ def oauth_account1() -> OAuthAccount:
 @pytest.fixture
 def oauth_account2() -> OAuthAccount:
     return OAuthAccount(
+        id=uuid.UUID("c9089e5d-2642-406d-a7c0-cbc641aca0ec"),
         oauth_name="service2",
         access_token="TOKEN",
         expires_at=1579000751,

tests/test_access_token.py

@@ -0,0 +1,141 @@
+import uuid
+from datetime import datetime, timedelta, timezone
+from typing import AsyncGenerator
+
+import pytest
+from pydantic import UUID4
+from sqlalchemy import exc
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+from sqlalchemy.orm import sessionmaker
+from sqlmodel import Session, SQLModel, create_engine
+
+from fastapi_users_db_sqlmodel import SQLModelUserDatabase, SQLModelUserDatabaseAsync
+from fastapi_users_db_sqlmodel.access_token import (
+    SQLModelAccessTokenDatabase,
+    SQLModelAccessTokenDatabaseAsync,
+    SQLModelBaseAccessToken,
+)
+from tests.conftest import UserDB
+
+
+class AccessToken(SQLModelBaseAccessToken, table=True):
+    pass
+
+
+@pytest.fixture
+def user_id() -> UUID4:
+    return uuid.UUID("a9089e5d-2642-406d-a7c0-cbc641aca0ec")
+
+
+async def init_sync_session(url: str) -> AsyncGenerator[Session, None]:
+    engine = create_engine(url, connect_args={"check_same_thread": False})
+    SQLModel.metadata.create_all(engine)
+    with Session(engine) as session:
+        yield session
+    SQLModel.metadata.drop_all(engine)
+
+
+async def init_async_session(url: str) -> AsyncGenerator[AsyncSession, None]:
+    engine = create_async_engine(url, connect_args={"check_same_thread": False})
+    make_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
+    async with engine.begin() as conn:
+        await conn.run_sync(SQLModel.metadata.create_all)
+        async with make_session() as session:
+            yield session
+        await conn.run_sync(SQLModel.metadata.drop_all)
+
+
+@pytest.fixture(
+    params=[
+        (
+            init_sync_session,
+            "sqlite:///./test-sqlmodel-access-token.db",
+            SQLModelAccessTokenDatabase,
+            SQLModelUserDatabase,
+        ),
+        (
+            init_async_session,
+            "sqlite+aiosqlite:///./test-sqlmodel-access-token.db",
+            SQLModelAccessTokenDatabaseAsync,
+            SQLModelUserDatabaseAsync,
+        ),
+    ],
+    ids=["sync", "async"],
+)
+async def sqlmodel_access_token_db(
+    request, user_id: UUID4
+) -> AsyncGenerator[SQLModelAccessTokenDatabase, None]:
+    create_session = request.param[0]
+    database_url = request.param[1]
+    access_token_database_class = request.param[2]
+    user_database_class = request.param[3]
+    async for session in create_session(database_url):
+        user = UserDB(
+            id=user_id, email="lancelot@camelot.bt", hashed_password="guinevere"
+        )
+        user_db = user_database_class(UserDB, session)
+        await user_db.create(user)
+        yield access_token_database_class(AccessToken, session)
+
+
+@pytest.mark.asyncio
+@pytest.mark.db
+async def test_queries(
+    sqlmodel_access_token_db: SQLModelAccessTokenDatabase[AccessToken],
+    user_id: UUID4,
+):
+    access_token = AccessToken(token="TOKEN", user_id=user_id)
+
+    # Create
+    access_token_db = await sqlmodel_access_token_db.create(access_token)
+    assert access_token_db.token == "TOKEN"
+    assert access_token_db.user_id == user_id
+
+    # Update
+    access_token_db.created_at = datetime.now(timezone.utc)
+    await sqlmodel_access_token_db.update(access_token_db)
+
+    # Get by token
+    access_token_by_token = await sqlmodel_access_token_db.get_by_token(
+        access_token_db.token
+    )
+    assert access_token_by_token is not None
+
+    # Get by token expired
+    access_token_by_token = await sqlmodel_access_token_db.get_by_token(
+        access_token_db.token, max_age=datetime.now(timezone.utc) + timedelta(hours=1)
+    )
+    assert access_token_by_token is None
+
+    # Get by token not expired
+    access_token_by_token = await sqlmodel_access_token_db.get_by_token(
+        access_token_db.token, max_age=datetime.now(timezone.utc) - timedelta(hours=1)
+    )
+    assert access_token_by_token is not None
+
+    # Get by token unknown
+    access_token_by_token = await sqlmodel_access_token_db.get_by_token(
+        "NOT_EXISTING_TOKEN"
+    )
+    assert access_token_by_token is None
+
+    # Delete token
+    await sqlmodel_access_token_db.delete(access_token_db)
+    deleted_access_token = await sqlmodel_access_token_db.get_by_token(
+        access_token_db.token
+    )
+    assert deleted_access_token is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.db
+async def test_insert_existing_token(
+    sqlmodel_access_token_db: SQLModelAccessTokenDatabase[AccessToken], user_id: UUID4
+):
+    access_token = AccessToken(token="TOKEN", user_id=user_id)
+    await sqlmodel_access_token_db.create(access_token)
+
+    with pytest.raises(exc.IntegrityError):
+        await sqlmodel_access_token_db.create(
+            AccessToken(token="TOKEN", user_id=user_id)
+        )

tests/test_fastapi_users_db_sqlmodel.py renamed to tests/test_users.py

@@ -3,9 +3,9 @@
 
 import pytest
 from sqlalchemy import exc
-from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 from sqlalchemy.orm import sessionmaker
-from sqlmodel import SQLModel, create_engine, Session
+from sqlmodel import Session, SQLModel, create_engine
 
 from fastapi_users_db_sqlmodel import (
     NotSetOAuthAccountTableError,
@@ -14,7 +14,6 @@
 )
 from tests.conftest import OAuthAccount, UserDB, UserDBOAuth
 
-
 safe_uuid = uuid.UUID("a9089e5d-2642-406d-a7c0-cbc641aca0ec")