Add AccessToken adapter

Author: frankie567

fastapi_users_db_beanie/__init__.py

@@ -7,7 +7,6 @@
 from pymongo import IndexModel
 from pymongo.collation import Collation
 
-
 __version__ = "0.0.0"
 
 

fastapi_users_db_beanie/access_token.py

@@ -0,0 +1,52 @@
+from datetime import datetime, timezone
+from typing import Any, Dict, Generic, Optional, Type, TypeVar
+
+from beanie import Document, PydanticObjectId
+from fastapi_users.authentication.strategy.db import AccessTokenDatabase
+from pydantic import Field
+from pymongo import IndexModel
+
+
+class BeanieBaseAccessToken(Document):
+    token: str
+    user_id: PydanticObjectId
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+
+    class Collection:
+        indexes = [IndexModel("token", unique=True)]
+
+
+AP_BEANIE = TypeVar("AP_BEANIE", bound=BeanieBaseAccessToken)
+
+
+class BeanieAccessTokenDatabase(Generic[AP_BEANIE], AccessTokenDatabase[AP_BEANIE]):
+    """
+    Access token database adapter for Beanie.
+
+    :param access_token_model: Beanie access token model.
+    """
+
+    def __init__(self, access_token_model: Type[AP_BEANIE]):
+        self.access_token_model = access_token_model
+
+    async def get_by_token(
+        self, token: str, max_age: Optional[datetime] = None
+    ) -> Optional[AP_BEANIE]:
+        query: Dict[str, Any] = {"token": token}
+        if max_age is not None:
+            query["created_at"] = {"$gte": max_age}
+        return await self.access_token_model.find_one(query)
+
+    async def create(self, create_dict: Dict[str, Any]) -> AP_BEANIE:
+        access_token = self.access_token_model(**create_dict)
+        return await access_token.save()
+
+    async def update(
+        self, access_token: AP_BEANIE, update_dict: Dict[str, Any]
+    ) -> AP_BEANIE:
+        for key, value in update_dict.items():
+            setattr(access_token, key, value)
+        return await access_token.save()
+
+    async def delete(self, access_token: AP_BEANIE) -> None:
+        await access_token.delete()

tests/test_access_token.py

@@ -0,0 +1,111 @@
+from datetime import datetime, timedelta, timezone
+from typing import AsyncGenerator
+
+import pymongo.errors
+import pytest
+from beanie import PydanticObjectId, init_beanie
+from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
+
+from fastapi_users_db_beanie.access_token import (
+    BeanieAccessTokenDatabase,
+    BeanieBaseAccessToken,
+)
+
+
+class AccessToken(BeanieBaseAccessToken):
+    pass
+
+
+@pytest.fixture(scope="module")
+async def mongodb_client():
+    client = AsyncIOMotorClient(
+        "mongodb://localhost:27017",
+        serverSelectionTimeoutMS=10000,
+        uuidRepresentation="standard",
+    )
+
+    try:
+        await client.server_info()
+        yield client
+        client.close()
+    except pymongo.errors.ServerSelectionTimeoutError:
+        pytest.skip("MongoDB not available", allow_module_level=True)
+        return
+
+
+@pytest.fixture
+async def beanie_access_token_db(
+    mongodb_client: AsyncIOMotorClient,
+) -> AsyncGenerator[BeanieAccessTokenDatabase, None]:
+    database: AsyncIOMotorDatabase = mongodb_client["test_database"]
+    await init_beanie(database=database, document_models=[AccessToken])
+
+    yield BeanieAccessTokenDatabase(AccessToken)
+
+    await mongodb_client.drop_database("test_database")
+
+
+@pytest.fixture
+def user_id() -> PydanticObjectId:
+    return PydanticObjectId()
+
+
+@pytest.mark.asyncio
+async def test_queries(
+    beanie_access_token_db: BeanieAccessTokenDatabase[AccessToken],
+    user_id: PydanticObjectId,
+):
+    access_token_create = {"token": "TOKEN", "user_id": user_id}
+
+    # Create
+    access_token = await beanie_access_token_db.create(access_token_create)
+    assert access_token.token == "TOKEN"
+    assert access_token.user_id == user_id
+
+    # Update
+    update_dict = {"created_at": datetime.now(timezone.utc)}
+    updated_access_token = await beanie_access_token_db.update(
+        access_token, update_dict
+    )
+    assert updated_access_token.created_at == update_dict["created_at"]
+
+    # Get by token
+    access_token_by_token = await beanie_access_token_db.get_by_token(
+        access_token.token
+    )
+    assert access_token_by_token is not None
+
+    # Get by token expired
+    access_token_by_token = await beanie_access_token_db.get_by_token(
+        access_token.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 beanie_access_token_db.get_by_token(
+        access_token.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 beanie_access_token_db.get_by_token(
+        "NOT_EXISTING_TOKEN"
+    )
+    assert access_token_by_token is None
+
+    # Delete token
+    await beanie_access_token_db.delete(access_token)
+    deleted_access_token = await beanie_access_token_db.get_by_token(access_token.token)
+    assert deleted_access_token is None
+
+
+@pytest.mark.asyncio
+async def test_insert_existing_token(
+    beanie_access_token_db: BeanieAccessTokenDatabase[AccessToken],
+    user_id: PydanticObjectId,
+):
+    access_token_create = {"token": "TOKEN", "user_id": user_id}
+    await beanie_access_token_db.create(access_token_create)
+
+    with pytest.raises(pymongo.errors.DuplicateKeyError):
+        await beanie_access_token_db.create(access_token_create)

tests/test_fastapi_users_db_beanie.py

@@ -1,12 +1,12 @@
 from typing import Any, AsyncGenerator, Dict, List, Optional
 
-from beanie import init_beanie
-from pydantic import Field
-import pytest
 import pymongo.errors
+import pytest
+from beanie import init_beanie
 from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
+from pydantic import Field
 
-from fastapi_users_db_beanie import BaseOAuthAccount, BeanieUserDatabase, BeanieBaseUser
+from fastapi_users_db_beanie import BaseOAuthAccount, BeanieBaseUser, BeanieUserDatabase
 
 
 class User(BeanieBaseUser):
@@ -63,7 +63,9 @@ async def beanie_user_db_oauth(
 
 
 @pytest.mark.asyncio
-async def test_queries(beanie_user_db: BeanieUserDatabase[User], oauth_account1: Dict[str, Any]):
+async def test_queries(
+    beanie_user_db: BeanieUserDatabase[User], oauth_account1: Dict[str, Any]
+):
     user_create = {
         "email": "lancelot@camelot.bt",
         "hashed_password": "guinevere",