Implement access token strategy db adapter

Author: frankie567

Makefile

@@ -8,7 +8,7 @@ format: isort
 
 test:
 	docker stop $(MONGODB_CONTAINER_NAME) || true
-	docker run -d --rm --name $(MONGODB_CONTAINER_NAME) -p 27017:27017 mongo:4.2
+	docker run -d --rm --name $(MONGODB_CONTAINER_NAME) -p 27017:27017 mongo:4.4
 	pytest --cov=fastapi_users_db_mongodb/ --cov-report=term-missing --cov-fail-under=100
 	docker stop $(MONGODB_CONTAINER_NAME)
 

fastapi_users_db_mongodb/__init__.py

@@ -29,31 +29,30 @@ def __init__(
     ):
         super().__init__(user_db_model)
         self.collection = collection
-        self.collection.create_index("id", unique=True)
-        self.collection.create_index("email", unique=True)
+        self.initialized = False
 
         if email_collation:
             self.email_collation = email_collation  # pragma: no cover
         else:
             self.email_collation = Collation("en", strength=2)
 
-        self.collection.create_index(
-            "email",
-            name="case_insensitive_email_index",
-            collation=self.email_collation,
-        )
-
     async def get(self, id: UUID4) -> Optional[UD]:
+        await self._initialize()
+
         user = await self.collection.find_one({"id": id})
         return self.user_db_model(**user) if user else None
 
     async def get_by_email(self, email: str) -> Optional[UD]:
+        await self._initialize()
+
         user = await self.collection.find_one(
             {"email": email}, collation=self.email_collation
         )
         return self.user_db_model(**user) if user else None
 
     async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UD]:
+        await self._initialize()
+
         user = await self.collection.find_one(
             {
                 "oauth_accounts.oauth_name": oauth,
@@ -63,12 +62,29 @@ async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UD
         return self.user_db_model(**user) if user else None
 
     async def create(self, user: UD) -> UD:
+        await self._initialize()
+
         await self.collection.insert_one(user.dict())
         return user
 
     async def update(self, user: UD) -> UD:
+        await self._initialize()
+
         await self.collection.replace_one({"id": user.id}, user.dict())
         return user
 
     async def delete(self, user: UD) -> None:
+        await self._initialize()
+
         await self.collection.delete_one({"id": user.id})
+
+    async def _initialize(self):
+        if not self.initialized:
+            await self.collection.create_index("id", unique=True)
+            await self.collection.create_index("email", unique=True)
+            await self.collection.create_index(
+                "email",
+                name="case_insensitive_email_index",
+                collation=self.email_collation,
+            )
+            self.initialized = True

fastapi_users_db_mongodb/access_token.py

@@ -0,0 +1,57 @@
+from datetime import datetime
+from typing import Any, Dict, Generic, Optional, Type
+
+from fastapi_users.authentication.strategy.db import A, AccessTokenDatabase
+from motor.motor_asyncio import AsyncIOMotorCollection
+
+
+class MongoDBAccessTokenDatabase(AccessTokenDatabase, Generic[A]):
+    """
+    Access token database adapter for MongoDB.
+
+    :param access_token_model: Pydantic model of a DB representation of an access token.
+    :param collection: Collection instance from `motor`.
+    """
+
+    collection: AsyncIOMotorCollection
+
+    def __init__(self, access_token_model: Type[A], collection: AsyncIOMotorCollection):
+        self.access_token_model = access_token_model
+        self.collection = collection
+        self.initialized = False
+
+    async def get_by_token(
+        self, token: str, max_age: Optional[datetime] = None
+    ) -> Optional[A]:
+        await self._initialize()
+
+        query: Dict[str, Any] = {"token": token}
+        if max_age is not None:
+            query["created_at"] = {"$gte": max_age}
+
+        access_token = await self.collection.find_one(query)
+        return self.access_token_model(**access_token) if access_token else None
+
+    async def create(self, access_token: A) -> A:
+        await self._initialize()
+
+        await self.collection.insert_one(access_token.dict())
+        return access_token
+
+    async def update(self, access_token: A) -> A:
+        await self._initialize()
+
+        await self.collection.replace_one(
+            {"token": access_token.token}, access_token.dict()
+        )
+        return access_token
+
+    async def delete(self, access_token: A) -> None:
+        await self._initialize()
+
+        await self.collection.delete_one({"token": access_token.token})
+
+    async def _initialize(self):
+        if not self.initialized:
+            await self.collection.create_index("token", unique=True)
+            self.initialized = True

pyproject.toml

@@ -22,7 +22,7 @@ classifiers = [
 description-file = "README.md"
 requires-python = ">=3.7"
 requires = [
-    "fastapi-users >= 6.1.2",
+    "fastapi-users >= 9.1.0",
     "motor >=2.5.1,<3.0.0"
 ]
 

requirements.txt

@@ -1,2 +1,2 @@
-fastapi-users >= 6.1.2
+fastapi-users >= 9.1.0
 motor >=2.5.1,<3.0.0

tests/test_access_token.py

@@ -0,0 +1,103 @@
+import uuid
+from datetime import datetime, timedelta, timezone
+from typing import AsyncGenerator
+
+import pymongo.errors
+import pytest
+from fastapi_users.authentication.strategy.db.models import BaseAccessToken
+from motor.motor_asyncio import AsyncIOMotorClient
+from pydantic import UUID4
+
+from fastapi_users_db_mongodb.access_token import MongoDBAccessTokenDatabase
+
+
+class AccessToken(BaseAccessToken):
+    pass
+
+
+@pytest.fixture
+def user_id() -> UUID4:
+    return uuid.uuid4()
+
+
+@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
+@pytest.mark.asyncio
+async def mongodb_access_token_db(
+    mongodb_client: AsyncIOMotorClient,
+) -> AsyncGenerator[MongoDBAccessTokenDatabase, None]:
+    db = mongodb_client["test_database_access_token"]
+    collection = db["access_tokens"]
+
+    yield MongoDBAccessTokenDatabase(AccessToken, collection)
+
+    await collection.delete_many({})
+
+
+@pytest.mark.asyncio
+@pytest.mark.db
+async def test_queries(
+    mongodb_access_token_db: MongoDBAccessTokenDatabase[AccessToken],
+    user_id: UUID4,
+):
+    access_token = AccessToken(token="TOKEN", user_id=user_id)
+
+    # Create
+    access_token_db = await mongodb_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 mongodb_access_token_db.update(access_token_db)
+
+    # Get by token
+    access_token_by_token = await mongodb_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 mongodb_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 mongodb_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 mongodb_access_token_db.get_by_token(
+        "NOT_EXISTING_TOKEN"
+    )
+    assert access_token_by_token is None
+
+    # Exception when inserting existing token
+    with pytest.raises(pymongo.errors.DuplicateKeyError):
+        await mongodb_access_token_db.create(access_token_db)
+
+    # Delete token
+    await mongodb_access_token_db.delete(access_token_db)
+    deleted_access_token = await mongodb_access_token_db.get_by_token(
+        access_token_db.token
+    )
+    assert deleted_access_token is None