Fix #391: put user creation logic in a importable function

fastapi_users/fastapi_users.py

@@ -11,6 +11,7 @@ from fastapi_users.router import (
     get_reset_password_router,
     get_users_router,
 )
+from fastapi_users.user import CreateUserProtocol, get_create_user
 
 try:
     from httpx_oauth.oauth2 import BaseOAuth2
@@ -31,6 +32,7 @@ class FastAPIUsers:
     :param user_update_model: Pydantic model for updating a user.
     :param user_db_model: Pydantic model of a DB representation of a user.
 
+    :attribute create_user: Helper function to create a user programmatically.
     :attribute get_current_user: Dependency callable to inject authenticated user.
     :attribute get_current_active_user: Dependency callable to inject active user.
     :attribute get_current_superuser: Dependency callable to inject superuser.
@@ -38,6 +40,7 @@ class FastAPIUsers:
 
     db: BaseUserDatabase
     authenticator: Authenticator
+    create_user: CreateUserProtocol
     _user_model: Type[models.BaseUser]
     _user_create_model: Type[models.BaseUserCreate]
     _user_update_model: Type[models.BaseUserUpdate]
@@ -61,6 +64,8 @@ class FastAPIUsers:
         self._user_update_model = user_update_model
         self._user_db_model = user_db_model
 
+        self.create_user = get_create_user(db, user_db_model)
+
         self.get_current_user = self.authenticator.get_current_user
         self.get_current_active_user = self.authenticator.get_current_active_user
         self.get_current_superuser = self.authenticator.get_current_superuser
@@ -83,10 +88,9 @@ class FastAPIUsers:
         after a successful registration.
         """
         return get_register_router(
-            self.db,
+            self.create_user,
             self._user_model,
             self._user_create_model,
-            self._user_db_model,
             after_register,
         )
 

fastapi_users/models.py

@@ -31,6 +31,8 @@ class BaseUser(CreateUpdateDictModel):
 class BaseUserCreate(CreateUpdateDictModel):
     email: EmailStr
     password: str
+    is_active: Optional[bool] = True
+    is_superuser: Optional[bool] = False
 
 
 class BaseUserUpdate(BaseUser):

fastapi_users/router/register.py

@@ -1,18 +1,16 @@
-from typing import Callable, Optional, Type, cast
+from typing import Callable, Optional, Type
 
 from fastapi import APIRouter, HTTPException, Request, status
 
 from fastapi_users import models
-from fastapi_users.db import BaseUserDatabase
-from fastapi_users.password import get_password_hash
 from fastapi_users.router.common import ErrorCode, run_handler
+from fastapi_users.user import CreateUserProtocol, UserAlreadyExists
 
 
 def get_register_router(
-    user_db: BaseUserDatabase[models.BaseUserDB],
+    create_user: CreateUserProtocol,
     user_model: Type[models.BaseUser],
     user_create_model: Type[models.BaseUserCreate],
-    user_db_model: Type[models.BaseUserDB],
     after_register: Optional[Callable[[models.UD, Request], None]] = None,
 ) -> APIRouter:
     """Generate a router with the register route."""
@@ -22,21 +20,14 @@ def get_register_router(
         "/register", response_model=user_model, status_code=status.HTTP_201_CREATED
     )
     async def register(request: Request, user: user_create_model):  # type: ignore
-        user = cast(models.BaseUserCreate, user)  # Prevent mypy complain
-        existing_user = await user_db.get_by_email(user.email)
-
-        if existing_user is not None:
+        try:
+            created_user = await create_user(user, safe=True)
+        except UserAlreadyExists:
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,
                 detail=ErrorCode.REGISTER_USER_ALREADY_EXISTS,
             )
 
-        hashed_password = get_password_hash(user.password)
-        db_user = user_db_model(
-            **user.create_update_dict(), hashed_password=hashed_password
-        )
-        created_user = await user_db.create(db_user)
-
         if after_register:
             await run_handler(after_register, created_user, request)
 

fastapi_users/user.py

@@ -0,0 +1,40 @@
+from typing import Awaitable, Type
+
+from typing_extensions import Protocol
+
+from fastapi_users import models
+from fastapi_users.db import BaseUserDatabase
+from fastapi_users.password import get_password_hash
+
+
+class UserAlreadyExists(Exception):
+    pass
+
+
+class CreateUserProtocol(Protocol):  # pragma: no cover
+    def __call__(
+        self, user: models.BaseUserCreate, safe: bool = False
+    ) -> Awaitable[models.BaseUserDB]:
+        pass
+
+
+def get_create_user(
+    user_db: BaseUserDatabase[models.BaseUserDB],
+    user_db_model: Type[models.BaseUserDB],
+) -> CreateUserProtocol:
+    async def create_user(
+        user: models.BaseUserCreate, safe: bool = False
+    ) -> models.BaseUserDB:
+        existing_user = await user_db.get_by_email(user.email)
+
+        if existing_user is not None:
+            raise UserAlreadyExists()
+
+        hashed_password = get_password_hash(user.password)
+        user_dict = (
+            user.create_update_dict() if safe else user.create_update_dict_superuser()
+        )
+        db_user = user_db_model(**user_dict, hashed_password=hashed_password)
+        return await user_db.create(db_user)
+
+    return create_user