Make Beanie generic in terms of document ID

Author: frankie567

fastapi_users_db_beanie/__init__.py

@@ -3,7 +3,7 @@
 
 from beanie import Document, PydanticObjectId
 from fastapi_users.db.base import BaseUserDatabase
-from fastapi_users.models import OAP
+from fastapi_users.models import ID, OAP
 from pydantic import BaseModel, Field
 from pymongo import IndexModel
 from pymongo.collation import Collation
@@ -41,9 +41,7 @@ class BaseOAuthAccount(BaseModel):
     refresh_token: Optional[str] = None
 
 
-class BeanieUserDatabase(
-    Generic[UP_BEANIE], BaseUserDatabase[UP_BEANIE, PydanticObjectId]
-):
+class BeanieUserDatabase(Generic[UP_BEANIE, ID], BaseUserDatabase[UP_BEANIE, ID]):
     """
     Database adapter for Beanie.
 
@@ -59,9 +57,9 @@ def __init__(
         self.user_model = user_model
         self.oauth_account_model = oauth_account_model
 
-    async def get(self, id: PydanticObjectId) -> Optional[UP_BEANIE]:
+    async def get(self, id: ID) -> Optional[UP_BEANIE]:
         """Get a single user by id."""
-        return await self.user_model.get(id)
+        return await self.user_model.get(id)  # type: ignore
 
     async def get_by_email(self, email: str) -> Optional[UP_BEANIE]:
         """Get a single user by email."""

fastapi_users_db_beanie/access_token.py

@@ -1,15 +1,16 @@
 from datetime import datetime, timezone
 from typing import Any, Dict, Generic, Optional, Type, TypeVar
 
-from beanie import Document, PydanticObjectId
+from beanie import Document
 from fastapi_users.authentication.strategy.db import AccessTokenDatabase
+from fastapi_users.models import ID
 from pydantic import Field
 from pymongo import IndexModel
 
 
-class BeanieBaseAccessToken(Document):
+class BeanieBaseAccessToken(Generic[ID], Document):
     token: str
-    user_id: PydanticObjectId
+    user_id: ID
     created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
 
     class Collection:

tests/test_access_token.py

@@ -12,7 +12,7 @@
 )
 
 
-class AccessToken(BeanieBaseAccessToken):
+class AccessToken(BeanieBaseAccessToken[PydanticObjectId]):
     pass
 
 

tests/test_fastapi_users_db_beanie.py

@@ -2,7 +2,7 @@
 
 import pymongo.errors
 import pytest
-from beanie import init_beanie
+from beanie import PydanticObjectId, init_beanie
 from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
 from pydantic import Field
 
@@ -64,7 +64,8 @@ async def beanie_user_db_oauth(
 
 @pytest.mark.asyncio
 async def test_queries(
-    beanie_user_db: BeanieUserDatabase[User], oauth_account1: Dict[str, Any]
+    beanie_user_db: BeanieUserDatabase[User, PydanticObjectId],
+    oauth_account1: Dict[str, Any],
 ):
     user_create = {
         "email": "lancelot@camelot.bt",
@@ -133,7 +134,10 @@ async def test_queries(
     ],
 )
 async def test_email_query(
-    beanie_user_db: BeanieUserDatabase[User], email: str, query: str, found: bool
+    beanie_user_db: BeanieUserDatabase[User, PydanticObjectId],
+    email: str,
+    query: str,
+    found: bool,
 ):
     user_create = {
         "email": email,
@@ -151,7 +155,9 @@ async def test_email_query(
 
 
 @pytest.mark.asyncio
-async def test_insert_existing_email(beanie_user_db: BeanieUserDatabase[User]):
+async def test_insert_existing_email(
+    beanie_user_db: BeanieUserDatabase[User, PydanticObjectId]
+):
     user_create = {
         "email": "lancelot@camelot.bt",
         "hashed_password": "guinevere",
@@ -164,7 +170,7 @@ async def test_insert_existing_email(beanie_user_db: BeanieUserDatabase[User]):
 
 @pytest.mark.asyncio
 async def test_queries_custom_fields(
-    beanie_user_db: BeanieUserDatabase[User],
+    beanie_user_db: BeanieUserDatabase[User, PydanticObjectId],
 ):
     """It should output custom fields in query result."""
     user_create = {
@@ -183,7 +189,7 @@ async def test_queries_custom_fields(
 
 @pytest.mark.asyncio
 async def test_queries_oauth(
-    beanie_user_db_oauth: BeanieUserDatabase[UserOAuth],
+    beanie_user_db_oauth: BeanieUserDatabase[UserOAuth, PydanticObjectId],
     oauth_account1: Dict[str, Any],
     oauth_account2: Dict[str, Any],
 ):