Merge pull request #1 from LonelyVikingMichael/usermodel-relationships

Usermodel relationships

Author: frankie567

fastapi_users_db_ormar/__init__.py

@@ -44,6 +44,7 @@ class OrmarUserDatabase(BaseUserDatabase[UD]):
     :param user_db_model: Pydantic model of a DB representation of a user.
     :param model: ormar ORM model.
     :param oauth_account_model: Optional ormar ORM model of a OAuth account.
+    :param select_related: Optional list of relationship names to retrieve with User queries.
     """
 
     model: Type[OrmarBaseUserModel]
@@ -54,10 +55,12 @@ def __init__(
         user_db_model: Type[UD],
         model: Type[OrmarBaseUserModel],
         oauth_account_model: Optional[Type[OrmarBaseOAuthAccountModel]] = None,
+        select_related: Optional[List[str]] = None
     ):
         super().__init__(user_db_model)
         self.model = model
         self.oauth_account_model = oauth_account_model
+        self.select_related = select_related
 
     async def get(self, id: UUID4) -> Optional[UD]:
         return await self._get_user(id=id)
@@ -73,6 +76,7 @@ async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UD
     async def create(self, user: UD) -> UD:
         oauth_accounts = getattr(user, "oauth_accounts", [])
         model = await self.model(**user.dict(exclude={"oauth_accounts"})).save()
+        await model.save_related()
         if oauth_accounts and self.oauth_account_model:
             await self._create_oauth_models(model=model, oauth_accounts=oauth_accounts)
         user_db = await self._get_user(id=user.id)
@@ -105,6 +109,9 @@ async def _get_db_user(self, **kwargs: Any) -> OrmarBaseUserModel:
         query = self.model.objects.filter(**kwargs)
         if self.oauth_account_model is not None:
             query = query.select_related("oauth_accounts")
+        if self.select_related is not None:
+            for relation in self.select_related:
+                query = query.select_related(relation)
         return cast(OrmarBaseUserModel, await query.get())
 
     async def _get_user(self, **kwargs: Any) -> Optional[UD]:

tests/conftest.py

@@ -1,16 +1,25 @@
+import uuid
 import asyncio
-from typing import Optional
+from typing import Optional, List
 
 import pytest
 from fastapi_users import models
+from pydantic import BaseModel, UUID4, Field
+
+
+class Role(BaseModel):
+    id: UUID4 = Field(default_factory=uuid.uuid4)
+    name: str
 
 
 class User(models.BaseUser):
     first_name: Optional[str]
+    roles: Optional[List[Role]]
 
 
 class UserCreate(models.BaseUserCreate):
     first_name: Optional[str]
+    roles: Optional[List[Role]]
 
 
 class UserUpdate(models.BaseUserUpdate):

tests/test_fastapi_users_db_ormar.py

@@ -20,12 +20,23 @@
 database = databases.Database(DATABASE_URL)
 
 
+class Role(ormar.Model):
+    class Meta:
+        tablename = "roles"
+        metadata = metadata
+        database = database
+
+    id = ormar.UUID(primary_key=True, uuid_format="string")
+    name = ormar.String(nullable=False, max_length=255)
+
+
 class User(OrmarBaseUserModel):
     class Meta:
         metadata = metadata
         database = database
 
     first_name = ormar.String(nullable=True, max_length=255)
+    roles = ormar.ManyToMany(Role, skip_reverse=True)
 
 
 class OAuthAccount(OrmarBaseOAuthAccountModel):
@@ -68,6 +79,21 @@ async def ormar_user_db_oauth() -> AsyncGenerator[OrmarUserDatabase, None]:
     await database.disconnect()
 
 
+@pytest.fixture
+async def ormar_user_db_with_relations() -> AsyncGenerator[OrmarUserDatabase, None]:
+    engine = sqlalchemy.create_engine(
+        DATABASE_URL, connect_args={"check_same_thread": False}
+    )
+    metadata.create_all(engine)
+
+    await database.connect()
+
+    yield OrmarUserDatabase(user_db_model=UserDB, model=User, select_related=['roles'])
+
+    metadata.drop_all(engine)
+    await database.disconnect()
+
+
 @pytest.mark.asyncio
 @pytest.mark.db
 async def test_queries(ormar_user_db: OrmarUserDatabase[UserDB]):
@@ -196,3 +222,29 @@ async def test_queries_oauth(
     # Unknown OAuth account
     unknown_oauth_user = await ormar_user_db_oauth.get_by_oauth_account("foo", "bar")
     assert unknown_oauth_user is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.db
+async def test_queries_custom_fields_relations(
+    ormar_user_db_with_relations: OrmarUserDatabase[UserDB]
+    ):
+    # Create role to pair with
+    role = await Role.objects.create(
+        id=uuid.uuid4(),
+        name="editor"
+    )
+
+    assert role.id is not None
+
+    user = UserDB(
+        email="[email protected]",
+        hashed_password="guinevere",
+        roles=[role]
+    )
+
+    # Create with relationship
+    user_db = await ormar_user_db_with_relations.create(user)
+    assert user_db.roles is not None
+    assert len(user_db.roles) is not 0
+    assert user_db.roles[0].id == role.id