feat: Let the user set and/or change the auth model (#21)

mslaursen · web-flow · commit 543424850df2 · 2025-04-04T13:57:45.000+02:00
diff --git a/fastapi_forge/dtos.py b/fastapi_forge/dtos.py
@@ -313,6 +313,10 @@ def _validate_models(self) -> Self:
             msg = "Cannot use built-in auth if PostgreSQL is not enabled."
             raise ValueError(msg)
 
+        if self.use_builtin_auth and self.get_auth_model() is None:
+            msg = "Cannot use built-in auth if no auth model is defined."
+            raise ValueError(msg)
+
         for model in self.models:
             for relationship in model.relationships:
                 if relationship.target_model not in model_names_set:
@@ -364,3 +368,11 @@ def has_cycle(node):
                 )
 
         return self
+
+    def get_auth_model(self) -> Model | None:
+        if not self.use_builtin_auth:
+            return None
+        for model in self.models:
+            if model.metadata.is_auth_model:
+                return model
+        return None
diff --git a/fastapi_forge/forge.py b/fastapi_forge/forge.py
@@ -37,17 +37,27 @@ async def build_project(spec: ProjectSpec) -> None:
 
         template_path = str(_get_template_path())
 
+        extra_context = {
+            **spec.model_dump(exclude={"models"}),
+            "models": {
+                "models": [model.model_dump() for model in spec.models],
+            },
+        }
+
+        if spec.use_builtin_auth:
+            auth_user = spec.get_auth_model()
+            if auth_user:
+                extra_context["auth_model"] = auth_user.model_dump()
+            else:
+                logger.warning("No auth model found. Skipping authentication setup.")
+                extra_context["use_builtin_auth"] = False
+
         cookiecutter(
             template_path,
             output_dir=str(Path.cwd()),
             no_input=True,
             overwrite_if_exists=True,
-            extra_context={
-                **spec.model_dump(exclude={"models"}),
-                "models": {
-                    "models": [model.model_dump() for model in spec.models],
-                },
-            },
+            extra_context=extra_context,
         )
         logger.info(f"Project '{spec.project_name}' created successfully.")
 
diff --git a/fastapi_forge/frontend/constants.py b/fastapi_forge/frontend/constants.py
@@ -1,5 +1,8 @@
 from typing import Any
 
+from fastapi_forge.dtos import ModelField
+from fastapi_forge.enums import FieldDataType
+
 SELECTED_MODEL_TEXT_COLOR = "text-black-500 dark:text-amber-300"
 
 FIELD_COLUMNS: list[dict[str, Any]] = [
@@ -46,3 +49,17 @@
     {"name": "index", "label": "Index", "field": "index", "align": "center"},
     {"name": "unique", "label": "Unique", "field": "unique", "align": "center"},
 ]
+
+
+DEFAULT_AUTH_USER_FIELDS: list[ModelField] = [
+    ModelField(
+        name="email",
+        type=FieldDataType.STRING,
+        unique=True,
+        index=True,
+    ),
+    ModelField(
+        name="password",
+        type=FieldDataType.STRING,
+    ),
+]
diff --git a/fastapi_forge/frontend/modals/field_modal.py b/fastapi_forge/frontend/modals/field_modal.py
@@ -122,6 +122,10 @@ def remove_row():
 
     def _show_field_preview(self) -> None:
         if not self.field_name.value:
+            ui.notify("Set a field name first", type="warning")
+            return
+        if not self.field_type.value:
+            ui.notify("Select a field type first", type="warning")
             return
         try:
             with ui.dialog() as modal, ui.card().classes("no-shadow border-[1px]"):
diff --git a/fastapi_forge/frontend/panels/model_editor_panel.py b/fastapi_forge/frontend/panels/model_editor_panel.py
@@ -5,7 +5,11 @@
 
 from fastapi_forge.dtos import Model, ModelField, ModelFieldMetadata, ModelRelationship
 from fastapi_forge.enums import FieldDataType
-from fastapi_forge.frontend.constants import FIELD_COLUMNS, RELATIONSHIP_COLUMNS
+from fastapi_forge.frontend.constants import (
+    DEFAULT_AUTH_USER_FIELDS,
+    FIELD_COLUMNS,
+    RELATIONSHIP_COLUMNS,
+)
 from fastapi_forge.frontend.modals import (
     AddFieldModal,
     AddRelationModal,
@@ -28,6 +32,7 @@ def __init__(self):
         state.select_model_fn = self.set_selected_model
         state.deselect_model_fn = self.deselect_model
         state.render_model_editor_fn = self.refresh
+        state.render_actions_fn = self._render_action_group
 
         self.add_field_modal: AddFieldModal = AddFieldModal(
             on_add_field=self._handle_modal_add_field,
@@ -56,6 +61,63 @@ def _show_code_preview(self) -> None:
                 ui.code(code).classes("w-full")
                 modal.open()
 
+    def _toggle_auth_model(self) -> None:
+        if not state.selected_model or not state.use_builtin_auth:
+            return
+
+        if not state.selected_model.metadata.is_auth_model and any(
+            m.metadata.is_auth_model for m in state.models
+        ):
+            ui.notify(
+                "Cannot have more than one authentication model.", type="negative"
+            )
+            return
+
+        model = state.selected_model
+        model.metadata.is_auth_model = not model.metadata.is_auth_model
+
+        if not model.metadata.is_auth_model:
+            self._remove_auth_fields(model)
+            if state.render_model_editor_fn:
+                state.render_model_editor_fn()
+
+        if state.render_models_fn:
+            state.render_models_fn()
+        self._render_action_group.refresh()
+
+        if model.metadata.is_auth_model:
+            self._setup_auth_model_fields(model)
+
+    def _remove_auth_fields(self, model: Model) -> None:
+        for field_name in ("email", "password"):
+            if field := next((f for f in model.fields if f.name == field_name), None):
+                model.fields.remove(field)
+
+    def _setup_auth_model_fields(self, model: Model) -> None:
+        self._remove_auth_fields(model)
+        id_index = 0
+        insert_position = id_index + 1 if id_index >= 0 else 0
+        for field in reversed(DEFAULT_AUTH_USER_FIELDS):
+            model.fields.insert(insert_position, field)
+
+        self._refresh_table(model.fields)
+
+    @ui.refreshable
+    def _render_action_group(self) -> None:
+        ui.button(
+            icon="security",
+            on_click=self._toggle_auth_model,
+            color=(
+                "green"
+                if state.use_builtin_auth
+                and state.selected_model
+                and state.selected_model.metadata.is_auth_model
+                else "grey"
+            ),
+        ).tooltip("Authentication model").bind_visibility_from(
+            state, "use_builtin_auth"
+        )
+
     def _build(self) -> None:
         with self:
             with ui.row().classes("w-full justify-between items-center"):
@@ -67,6 +129,7 @@ def _build(self) -> None:
                     ).tooltip("Preview SQLAlchemy model code")
 
                 with ui.row().classes("gap-2 items-center"):
+                    self._render_action_group()
                     with (
                         ui.button(icon="menu").tooltip("Generate"),
                         ui.menu(),
@@ -373,22 +436,35 @@ def _deselect_relation(self) -> None:
         self.relationship_table.selected = []
 
     def _on_select_field(self, selection: list[dict[str, Any]]) -> None:
-        if not state.selected_model:
+        if not state.selected_model or not selection:
+            self._deselect_field()
             return
-        if not selection:
+
+        name = selection[0].get("name")
+
+        if (
+            state.selected_model.metadata.is_auth_model
+            and state.use_builtin_auth
+            and name in ("password", "email")
+        ):
+            ui.notify(
+                f"Cannot edit {name} field in authentication model.",
+                type="warning",
+            )
             self._deselect_field()
             return
-        if selection[0].get("name") == "id":
+
+        if name == "id":
             self._deselect_field()
-        else:
-            state.selected_field = next(
-                (
-                    field
-                    for field in state.selected_model.fields
-                    if field.name == selection[0]["name"]
-                ),
-                None,
+            ui.notify(
+                "Cannot edit the 'id' field, it is automatically generated.",
+                type="warning",
             )
+            return
+
+        state.selected_field = next(
+            (field for field in state.selected_model.fields if field.name == name), None
+        )
 
     def _on_select_relation(self, selection: list[dict[str, Any]]) -> None:
         if not state.selected_model:
@@ -417,11 +493,18 @@ def _handle_update_field(
         default_value: str | None = None,
         extra_kwargs: dict[str, Any] | None = None,
     ) -> None:
+        if not state.selected_model or not state.selected_field:
+            return
+
+        exclude_set = {"id"}
         if (
-            not state.selected_model
-            or not state.selected_field
-            or state.selected_field.name == "id"
+            state.selected_model
+            and state.selected_model.metadata.is_auth_model
+            and state.use_builtin_auth
         ):
+            exclude_set.update({"email", "password"})
+
+        if name in exclude_set:
             return
 
         if state.selected_field.name != name and self._field_name_exists(name):
diff --git a/fastapi_forge/frontend/panels/project_config_panel.py b/fastapi_forge/frontend/panels/project_config_panel.py
@@ -12,6 +12,7 @@
 )
 from fastapi_forge.enums import FieldDataType
 from fastapi_forge.forge import build_project
+from fastapi_forge.frontend.constants import DEFAULT_AUTH_USER_FIELDS
 from fastapi_forge.frontend.notifications import notify_validation_error
 from fastapi_forge.frontend.state import state
 
@@ -166,16 +167,7 @@ def _handle_builtin_auth_change(self, event: ValueChangeEventArguments) -> None:
                             unique=True,
                             index=True,
                         ),
-                        ModelField(
-                            name="email",
-                            type=FieldDataType.STRING,
-                            unique=True,
-                            index=True,
-                        ),
-                        ModelField(
-                            name="password",
-                            type=FieldDataType.STRING,
-                        ),
+                        *DEFAULT_AUTH_USER_FIELDS,
                         ModelField(
                             name="created_at",
                             type=FieldDataType.DATETIME,
diff --git a/fastapi_forge/frontend/state.py b/fastapi_forge/frontend/state.py
@@ -27,6 +27,7 @@ class ProjectState(BaseModel):
 
     render_models_fn: Callable | None = None
     render_model_editor_fn: Callable | None = None
+    render_actions_fn: Callable | None = None
     select_model_fn: Callable[[Model], None] | None = None
     deselect_model_fn: Callable | None = None
 
@@ -177,6 +178,8 @@ def _trigger_ui_refresh(self) -> None:
             self.render_models_fn()
         if self.render_model_editor_fn:
             self.render_model_editor_fn()
+        if self.render_actions_fn:
+            self.render_actions_fn.refresh()
 
 
 state: ProjectState = ProjectState()
diff --git a/fastapi_forge/template/cookiecutter.json b/fastapi_forge/template/cookiecutter.json
@@ -20,6 +20,7 @@
     "models": {
         "models": []
     },
+    "auth_model": null,
     "_extensions": [
         "local_extensions.camel_to_snake"
     ]
diff --git a/fastapi_forge/template/{{cookiecutter.project_name}}/src/__main__.py b/fastapi_forge/template/{{cookiecutter.project_name}}/src/__main__.py
@@ -1,6 +1,7 @@
 from src.settings import settings
 
 
+
 if __name__ == "__main__":
     import uvicorn
 
diff --git a/fastapi_forge/template/{{cookiecutter.project_name}}/src/dependencies/auth_dependencies.py b/fastapi_forge/template/{{cookiecutter.project_name}}/src/dependencies/auth_dependencies.py
@@ -1,11 +1,12 @@
-{%- if cookiecutter.use_builtin_auth %}
+from typing import Annotated
+
 from fastapi import Depends, HTTPException, Request
 from fastapi.security import HTTPBearer as _HTTPBearer
-from src.dtos.auth_user_dtos import AuthUserDTO
-from src.daos import GetDAOs
+
 from src import exceptions
+from src.daos import GetDAOs
+from src.dtos.{{ cookiecutter.auth_model.name }}_dtos import {{ cookiecutter.auth_model.name_cc }}DTO
 from src.utils import auth_utils
-from typing import Annotated
 
 
 class HTTPBearer(_HTTPBearer):
@@ -20,7 +21,8 @@ async def __call__(self, request: Request) -> str | None:  # type: ignore
             obj = await super().__call__(request)
             return obj.credentials if obj else None
         except HTTPException:
-            raise exceptions.Http401("Missing token.")
+            msg = "Missing token."
+            raise exceptions.Http401(msg)
 
 
 auth_scheme = HTTPBearer()
@@ -37,17 +39,17 @@ def get_token(token: str = Depends(auth_scheme)) -> str:
 async def get_current_user(
     token: GetToken,
     daos: GetDAOs,
-) -> AuthUserDTO:
+) -> {{ cookiecutter.auth_model.name_cc }}DTO:
     """Get current user from token data."""
     token_data = auth_utils.decode_token(token)
 
-    user = await daos.auth_user.filter_first(id=token_data.user_id)
+    user = await daos.{{ cookiecutter.auth_model.name }}.filter_first(id=token_data.user_id)
 
     if not user:
-        raise exceptions.Http404("Decoded user not found.")
+        msg = "Decoded user not found."
+        raise exceptions.Http404(msg)
 
-    return AuthUserDTO.model_validate(user)
+    return {{ cookiecutter.auth_model.name_cc }}DTO.model_validate(user)
 
 
-GetCurrentUser = Annotated[AuthUserDTO, Depends(get_current_user)]
-{% endif %}
+GetCurrentUser = Annotated[{{ cookiecutter.auth_model.name_cc }}DTO, Depends(get_current_user)]
diff --git a/fastapi_forge/template/{{cookiecutter.project_name}}/src/routes/auth_routes.py b/fastapi_forge/template/{{cookiecutter.project_name}}/src/routes/auth_routes.py

Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,7 @@`
`20`	`20`	`"models": {`
`21`	`21`	`"models": []`
`22`	`22`	`},`
	`23`	`+ "auth_model": null,`
`23`	`24`	`"_extensions": [`
`24`	`25`	`"local_extensions.camel_to_snake"`
`25`	`26`	`]`