Add ObjectIDIDMixin

Author: frankie567

fastapi_users_db_beanie/__init__.py

@@ -1,7 +1,9 @@
 """FastAPI Users database adapter for Beanie."""
 from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar
 
+import bson.errors
 from beanie import Document, PydanticObjectId
+from fastapi_users import InvalidID
 from fastapi_users.db.base import BaseUserDatabase
 from fastapi_users.models import ID, OAP
 from pydantic import BaseModel, Field
@@ -127,3 +129,11 @@ async def update_oauth_account(
                     setattr(user.oauth_accounts[i], key, value)  # type: ignore
 
         return await user.save()
+
+
+class ObjectIDIDMixin:
+    def parse_id(self, value: Any) -> PydanticObjectId:
+        try:
+            return PydanticObjectId(value)
+        except (bson.errors.InvalidId, TypeError) as e:
+            raise InvalidID() from e

tests/test_fastapi_users_db_beanie.py

@@ -3,10 +3,16 @@
 import pymongo.errors
 import pytest
 from beanie import PydanticObjectId, init_beanie
+from fastapi_users import InvalidID
 from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
 from pydantic import Field
 
-from fastapi_users_db_beanie import BaseOAuthAccount, BeanieBaseUser, BeanieUserDatabase
+from fastapi_users_db_beanie import (
+    BaseOAuthAccount,
+    BeanieBaseUser,
+    BeanieUserDatabase,
+    ObjectIDIDMixin,
+)
 
 
 class User(BeanieBaseUser[PydanticObjectId]):
@@ -238,3 +244,17 @@ async def test_queries_oauth(
     # Unknown OAuth account
     unknown_oauth_user = await beanie_user_db_oauth.get_by_oauth_account("foo", "bar")
     assert unknown_oauth_user is None
+
+
+def test_objectid_id_mixin():
+    object_id_mixin = ObjectIDIDMixin()
+    object_id = PydanticObjectId("62736e11bae73a7a990f7df1")
+
+    assert object_id_mixin.parse_id("62736e11bae73a7a990f7df1") == object_id
+    assert object_id_mixin.parse_id(object_id) == object_id
+
+    with pytest.raises(InvalidID):
+        object_id_mixin.parse_id("abc")
+
+    with pytest.raises(InvalidID):
+        object_id_mixin.parse_id(12346)