Eagerly load roles, orgs, and permissions with user in endpoints that need them

Author: chriscarrollsmith

main.py

@@ -8,10 +8,9 @@
 from fastapi.exceptions import RequestValidationError, HTTPException, StarletteHTTPException
 from sqlmodel import Session
 from routers import authentication, organization, role, user
-from utils.auth import get_authenticated_user, get_optional_user, NeedsNewTokens, get_user_from_reset_token, PasswordValidationError, AuthenticationError
+from utils.auth import get_authenticated_user, get_user_with_relations, get_optional_user, NeedsNewTokens, get_user_from_reset_token, PasswordValidationError, AuthenticationError
 from utils.models import User
 from utils.db import get_session, set_up_db
-from utils.role_org import get_user_organizations, get_organization_roles
 
 logger = logging.getLogger("uvicorn.error")
 logger.setLevel(logging.DEBUG)
@@ -236,28 +235,27 @@ async def common_authenticated_parameters(
     return {"request": request, "user": user, "error_message": error_message}
 
 
+async def common_authenticated_parameters_with_organizations(
+    request: Request,
+    user: User = Depends(get_user_with_relations),
+    error_message: Optional[str] = None
+) -> dict:
+    return {"request": request, "user": user, "error_message": error_message}
+
+
 # Redirect to home if user is not authenticated
 @app.get("/dashboard")
 async def read_dashboard(
     params: dict = Depends(common_authenticated_parameters)
 ):
-    if not params["user"]:
-        return RedirectResponse(url="/login", status_code=status.HTTP_302_FOUND)
     return templates.TemplateResponse(params["request"], "dashboard/index.html", params)
 
 
 @app.get("/profile")
 async def read_profile(
-    params: dict = Depends(common_authenticated_parameters),
-    session: Session = Depends(get_session)
+    params: dict = Depends(common_authenticated_parameters_with_organizations)
 ):
-    if not params["user"]:
-        return RedirectResponse(url="/login", status_code=status.HTTP_302_FOUND)
-
-    # Get user's organizations
-    params["organizations"] = get_user_organizations(
-        params["user"].id, session)
-
+    params["organizations"] = params["user"].organizations
     return templates.TemplateResponse(params["request"], "users/profile.html", params)
 
 

routers/organization.py

@@ -4,11 +4,9 @@
 from pydantic import BaseModel, ConfigDict, field_validator
 from sqlmodel import Session, select
 from utils.db import get_session
-from utils.auth import get_authenticated_user
-from utils.models import Organization, User, Role, Permission, UserOrganizationLink, ValidPermissions, utc_time
+from utils.auth import get_authenticated_user, get_user_with_relations
+from utils.models import Organization, User, Role, utc_time, default_roles
 from datetime import datetime
-from sqlalchemy import and_
-from utils.role_org import get_organization, check_user_permission
 
 logger = getLogger("uvicorn.error")
 
@@ -23,14 +21,6 @@ def __init__(self):
         )
 
 
-class OrganizationExistsError(HTTPException):
-    def __init__(self):
-        super().__init__(
-            status_code=400,
-            detail="Organization already exists"
-        )
-
-
 class OrganizationNotFoundError(HTTPException):
     def __init__(self):
         super().__init__(
@@ -109,64 +99,49 @@ def create_organization(
     user: User = Depends(get_authenticated_user),
     session: Session = Depends(get_session)
 ) -> RedirectResponse:
+    # Check if organization already exists
     db_org = session.exec(select(Organization).where(
         Organization.name == org.name)).first()
     if db_org:
-        raise OrganizationExistsError()
+        raise OrganizationNameTakenError()
 
+    # Create organization first
     db_org = Organization(name=org.name)
     session.add(db_org)
-    session.commit()
-    session.refresh(db_org)
+    # This gets us the org ID without committing
+    session.flush()
 
-    # Create default roles
-    default_role_names = ["Owner", "Administrator", "Member"]
-    default_roles = []
-    for role_name in default_role_names:
-        role = Role(name=role_name, organization_id=db_org.id)
-        session.add(role)
-        default_roles.append(role)
-    session.commit()
+    # Create default roles with organization_id
+    initial_roles = [
+        Role(name=name, organization_id=db_org.id)
+        for name in default_roles
+    ]
+    session.add_all(initial_roles)
+    session.flush()
 
-    owner_role = session.exec(
-        select(Role).where(
-            and_(
-                Role.organization_id == db_org.id,
-                Role.name == "Owner"
-            )
-        )
-    ).first()
+    # Get owner role for user assignment
+    owner_role = next(role for role in db_org.roles if role.name == "Owner")
 
-    if not owner_role:
-        owner_role = Role(
-            name="Owner",
-            organization_id=db_org.id
-        )
-        session.add(owner_role)
-        session.commit()
-        session.refresh(owner_role)
-
-    user_org_link = UserOrganizationLink(
-        user_id=user.id,
-        organization_id=db_org.id,
-        role_id=owner_role.id
-    )
-    session.add(user_org_link)
+    # Assign user to owner role
+    user.roles.append(owner_role)
+
+    # Commit changes
     session.commit()
+    session.refresh(db_org)
 
     return RedirectResponse(url=f"/profile", status_code=303)
 
 
 @router.put("/{org_id}", response_class=RedirectResponse)
 def update_organization(
     org: OrganizationUpdate = Depends(OrganizationUpdate.as_form),
-    user: User = Depends(get_authenticated_user),
+    user: User = Depends(get_user_with_relations),
     session: Session = Depends(get_session)
 ) -> RedirectResponse:
     # This will raise appropriate exceptions if org doesn't exist or user lacks access
-    organization = get_organization(org.id, user.id, session)
+    organization: Organization = user.organizations.get(org.id)
 
-    if not check_user_permission(user.id, org.id, ValidPermissions.EDIT_ORGANIZATION, session):
+    if not organization or not any(role.permissions.EDIT_ORGANIZATION for role in organization.roles):
         raise InsufficientPermissionsError()
 
     # Check if new name already exists for another organization
@@ -178,24 +153,27 @@ def update_organization(
     if existing_org:
         raise OrganizationNameTakenError()
 
+    # Update organization name
     organization.name = org.name
     organization.updated_at = utc_time()
     session.add(organization)
     session.commit()
-    session.refresh(organization)
 
     return RedirectResponse(url=f"/profile", status_code=303)
 
 
 @router.delete("/{org_id}", response_class=RedirectResponse)
 def delete_organization(
     org_id: int,
-    user: User = Depends(get_authenticated_user),
+    user: User = Depends(get_user_with_relations),
     session: Session = Depends(get_session)
 ) -> RedirectResponse:
-    # This will raise appropriate exceptions if org doesn't exist or user lacks access
-    organization = get_organization(org_id, user.id, session)
+    # Check if user has permission to delete organization
+    organization: Organization = user.organizations.get(org_id)
+    if not organization or not any(role.permissions.DELETE_ORGANIZATION for role in organization.roles):
+        raise InsufficientPermissionsError()
 
+    # Delete organization
     session.delete(organization)
     session.commit()
 

utils/auth.py

@@ -8,12 +8,13 @@
 from dotenv import load_dotenv
 from pydantic import field_validator, ValidationInfo
 from sqlmodel import Session, select
+from sqlalchemy.orm import selectinload
 from bcrypt import gensalt, hashpw, checkpw
 from datetime import UTC, datetime, timedelta
 from typing import Optional
 from fastapi import Depends, Cookie, HTTPException, status
 from utils.db import get_session
-from utils.models import User, PasswordResetToken
+from utils.models import User, Role, PasswordResetToken
 
 load_dotenv()
 logger = logging.getLogger("uvicorn.error")
@@ -339,3 +340,23 @@ def get_user_from_reset_token(email: str, token: str, session: Session) -> tuple
 
     user, reset_token = result
     return user, reset_token
+
+
+def get_user_with_relations(
+    user: User = Depends(get_authenticated_user),
+    session: Session = Depends(get_session),
+) -> User:
+    """
+    Returns an authenticated user with fully loaded role and organization relationships.
+    """
+    # Refresh the user instance with eagerly loaded relationships
+    eager_user = session.exec(
+        select(User)
+        .where(User.id == user.id)
+        .options(
+            selectinload(User.roles).selectinload(Role.organization),
+            selectinload(User.roles).selectinload(Role.permissions)
+        )
+    ).one()
+
+    return eager_user

utils/models.py

@@ -3,7 +3,7 @@
 from datetime import datetime, UTC, timedelta
 from typing import Optional, List
 from sqlmodel import SQLModel, Field, Relationship
-from sqlalchemy import Column, Enum as SQLAlchemyEnum, ForeignKey
+from sqlalchemy import Column, Enum as SQLAlchemyEnum
 
 
 def utc_time():
@@ -13,6 +13,8 @@ def utc_time():
 default_roles = ["Owner", "Administrator", "Member"]
 
 
+# TODO: User with permission to create/edit roles can only assign permissions
+# they themselves have.
 class ValidPermissions(Enum):
     DELETE_ORGANIZATION = "Delete Organization"
     EDIT_ORGANIZATION = "Edit Organization"
@@ -24,24 +26,17 @@ class ValidPermissions(Enum):
     EDIT_ROLE = "Edit Role"
 
 
-class UserOrganizationLink(SQLModel, table=True):
+class UserRoleLink(SQLModel, table=True):
+    """
+    Associates users with roles. This creates a many-to-many relationship
+    between users and roles.
+    """
     id: Optional[int] = Field(default=None, primary_key=True)
     user_id: int = Field(foreign_key="user.id")
-    organization_id: int = Field(foreign_key="organization.id")
     role_id: int = Field(foreign_key="role.id")
     created_at: datetime = Field(default_factory=utc_time)
     updated_at: datetime = Field(default_factory=utc_time)
 
-    user: "User" = Relationship(
-        back_populates="organization_links"
-    )
-    organization: "Organization" = Relationship(
-        back_populates="user_links"
-    )
-    role: "Role" = Relationship(
-        back_populates="user_links"
-    )
-
 
 class RolePermissionLink(SQLModel, table=True):
     id: Optional[int] = Field(default=None, primary_key=True)
@@ -74,14 +69,20 @@ class Organization(SQLModel, table=True):
     created_at: datetime = Field(default_factory=utc_time)
     updated_at: datetime = Field(default_factory=utc_time)
 
-    user_links: List[UserOrganizationLink] = Relationship(
+    roles: List["Role"] = Relationship(
         back_populates="organization",
         sa_relationship_kwargs={
             "cascade": "all, delete-orphan",
             "passive_deletes": True
         }
     )
-    roles: List["Role"] = Relationship(back_populates="organization")
+
+    @property
+    def users(self) -> List["User"]:
+        """
+        Returns all users in the organization via their roles.
+        """
+        return [role.users for role in self.roles]
 
 
 class Role(SQLModel, table=True):
@@ -101,9 +102,11 @@ class Role(SQLModel, table=True):
     created_at: datetime = Field(default_factory=utc_time)
     updated_at: datetime = Field(default_factory=utc_time)
 
-    organization: "Organization" = Relationship(back_populates="roles")
-    user_links: List[UserOrganizationLink] = Relationship(
-        back_populates="role")
+    organization: Organization = Relationship(back_populates="roles")
+    users: List["User"] = Relationship(
+        back_populates="roles",
+        link_model=UserRoleLink
+    )
     permissions: List["Permission"] = Relationship(
         back_populates="roles",
         link_model=RolePermissionLink
@@ -133,12 +136,9 @@ class User(SQLModel, table=True):
     created_at: datetime = Field(default_factory=utc_time)
     updated_at: datetime = Field(default_factory=utc_time)
 
-    organization_links: List[UserOrganizationLink] = Relationship(
-        back_populates="user",
-        sa_relationship_kwargs={
-            "cascade": "all, delete-orphan",
-            "passive_deletes": True
-        }
+    roles: List[Role] = Relationship(
+        back_populates="users",
+        link_model=UserRoleLink
     )
     password_reset_tokens: List["PasswordResetToken"] = Relationship(
         back_populates="user",
@@ -147,3 +147,10 @@ class User(SQLModel, table=True):
             "passive_deletes": True
         }
     )
+
+    @property
+    def organizations(self) -> List[Organization]:
+        """
+        Returns all organizations the user belongs to via their roles.
+        """
+        return [role.organization for role in self.roles]