Replaced role and organization endpoints with helper functions to return them as context

Author: chriscarrollsmith

routers/organization.py

@@ -5,9 +5,10 @@
 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, UserOrganizationLink, ValidPermissions, RolePermissionLink, Permission
+from utils.models import Organization, User, Role, UserOrganizationLink, ValidPermissions, RolePermissionLink, Permission, utc_time
 from datetime import datetime
 from sqlalchemy import and_
+from typing import List
 
 logger = getLogger("uvicorn.error")
 
@@ -101,22 +102,99 @@ async def as_form(cls, id: int = Form(...), name: str = Form(...)):
         return cls(id=id, name=name)
 
 
-# -- Routes --
+# -- Helper Functions --
 
-def check_user_permission(
+def get_user_organizations(
+    user_id: int,
     session: Session,
-    user: User,
+    include_deleted: bool = False
+) -> List[Organization]:
+    """
+    Retrieve all organizations a user is a member of.
+
+    Args:
+        user_id: ID of the user
+        session: Database session
+        include_deleted: Whether to include soft-deleted organizations
+
+    Returns:
+        List of Organization objects the user belongs to
+    """
+    query = (
+        select(Organization)
+        .join(UserOrganizationLink)
+        .where(UserOrganizationLink.user_id == user_id)
+    )
+
+    if not include_deleted:
+        query = query.where(Organization.deleted == False)
+
+    return list(session.exec(query))
+
+
+def get_organization(
     org_id: int,
-    permission: ValidPermissions
+    user_id: int,
+    session: Session,
+) -> Organization:
+    """
+    Retrieve a specific organization if the user is a member.
+
+    Args:
+        org_id: ID of the organization
+        user_id: ID of the user
+        session: Database session
+
+    Returns:
+        Organization object
+
+    Raises:
+        OrganizationNotFoundError: If organization doesn't exist
+        InsufficientPermissionsError: If user is not a member
+    """
+    # Check if user is a member of the organization
+    user_org = session.exec(
+        select(UserOrganizationLink).where(
+            and_(
+                UserOrganizationLink.user_id == user_id,
+                UserOrganizationLink.organization_id == org_id
+            )
+        )
+    ).first()
+
+    if not user_org:
+        raise InsufficientPermissionsError()
+
+    db_org = session.get(Organization, org_id)
+    if not db_org or db_org.deleted:
+        raise OrganizationNotFoundError()
+
+    return db_org
+
+
+def check_user_permission(
+    user_id: int,
+    org_id: int,
+    permission: ValidPermissions,
+    session: Session,
 ) -> bool:
     """
-    Check if user has the specified permission for the organization
+    Check if user has the specified permission for the organization.
+
+    Args:
+        user_id: ID of the user
+        org_id: ID of the organization
+        permission: Permission to check
+        session: Database session
+
+    Returns:
+        True if user has permission, False otherwise
     """
     # Get user's role in the organization
     user_org = session.exec(
         select(UserOrganizationLink).where(
             and_(
-                UserOrganizationLink.user_id == user.id,
+                UserOrganizationLink.user_id == user_id,
                 UserOrganizationLink.organization_id == org_id
             )
         )
@@ -146,6 +224,8 @@ def check_user_permission(
     return bool(role_permission)
 
 
+# -- Routes --
+
 @router.post("/", response_class=RedirectResponse)
 def create_organization(
     org: OrganizationCreate = Depends(OrganizationCreate.as_form),
@@ -191,43 +271,17 @@ def create_organization(
     return RedirectResponse(url=f"/organizations/{db_org.id}", status_code=303)
 
 
-@router.get("/{org_id}", response_model=OrganizationRead)
-def read_organization(
-    org_id: int,
-    user: User = Depends(get_authenticated_user),
-    session: Session = Depends(get_session)
-):
-    # First check if user is a member of the organization
-    user_org = session.exec(
-        select(UserOrganizationLink).where(
-            and_(
-                UserOrganizationLink.user_id == user.id,
-                UserOrganizationLink.organization_id == org_id
-            )
-        )
-    ).first()
-
-    if not user_org:
-        raise InsufficientPermissionsError()
-
-    db_org = session.get(Organization, org_id)
-    if not db_org:
-        raise OrganizationNotFoundError()
-    return db_org
-
-
 @router.put("/{org_id}", response_class=RedirectResponse)
 def update_organization(
     org: OrganizationUpdate = Depends(OrganizationUpdate.as_form),
     user: User = Depends(get_authenticated_user),
     session: Session = Depends(get_session)
 ) -> RedirectResponse:
-    if not check_user_permission(session, user, org.id, ValidPermissions.EDIT_ORGANIZATION):
-        raise InsufficientPermissionsError()
+    # This will raise appropriate exceptions if org doesn't exist or user lacks access
+    organization = get_organization(org.id, user.id, session)
 
-    db_org = session.get(Organization, org.id)
-    if not db_org:
-        raise OrganizationNotFoundError()
+    if not check_user_permission(user.id, org.id, ValidPermissions.EDIT_ORGANIZATION, session):
+        raise InsufficientPermissionsError()
 
     # Check if new name already exists for another organization
     existing_org = session.exec(
@@ -238,11 +292,11 @@ def update_organization(
     if existing_org:
         raise OrganizationNameTakenError()
 
-    db_org.name = org.name
-    db_org.updated_at = datetime.utcnow()
-    session.add(db_org)
+    organization.name = org.name
+    organization.updated_at = utc_time()
+    session.add(organization)
     session.commit()
-    session.refresh(db_org)
+    session.refresh(organization)
 
     return RedirectResponse(url=f"/organizations/{org.id}", status_code=303)
 
@@ -253,16 +307,15 @@ def delete_organization(
     user: User = Depends(get_authenticated_user),
     session: Session = Depends(get_session)
 ) -> RedirectResponse:
-    if not check_user_permission(session, user, org_id, ValidPermissions.DELETE_ORGANIZATION):
-        raise InsufficientPermissionsError()
+    # This will raise appropriate exceptions if org doesn't exist or user lacks access
+    organization = get_organization(org_id, user.id, session)
 
-    db_org = session.get(Organization, org_id)
-    if not db_org:
-        raise OrganizationNotFoundError()
+    if not check_user_permission(user.id, org_id, ValidPermissions.DELETE_ORGANIZATION, session):
+        raise InsufficientPermissionsError()
 
-    db_org.deleted = True
-    db_org.updated_at = datetime.utcnow()
-    session.add(db_org)
+    organization.deleted = True
+    organization.updated_at = utc_time()
+    session.add(organization)
     session.commit()
 
     return RedirectResponse(url="/organizations", status_code=303)

routers/role.py

@@ -107,6 +107,31 @@ async def as_form(
         )
 
 
+# -- Helper Functions --
+
+def get_organization_roles(
+    organization_id: int,
+    session: Session,
+    include_deleted: bool = False
+) -> List[Role]:
+    """
+    Retrieve all roles for an organization.
+
+    Args:
+        organization_id: ID of the organization
+        session: Database session
+        include_deleted: Whether to include soft-deleted roles
+
+    Returns:
+        List of Role objects with their associated permissions
+    """
+    query = select(Role).where(Role.organization_id == organization_id)
+    if not include_deleted:
+        query = query.where(Role.deleted == False)
+
+    return list(session.exec(query))
+
+
 # -- Routes --
 
 
@@ -117,7 +142,10 @@ def create_role(
     session: Session = Depends(get_session)
 ) -> RedirectResponse:
     # Create role and permissions in a single transaction
-    db_role = Role(name=role.name)
+    db_role = Role(
+        name=role.name,
+        organization_id=user.organization_id  # Add organization ID to role
+    )
 
     # Create RolePermissionLink objects and associate them with the role
     db_role.permissions = [
@@ -131,32 +159,6 @@ def create_role(
     return RedirectResponse(url="/roles", status_code=303)
 
 
-@router.get("/{role_id}", response_model=RoleRead)
-def read_role(
-    role_id: int,
-    user: User = Depends(get_authenticated_user),
-    session: Session = Depends(get_session)
-):
-    db_role: Role | None = session.get(Role, role_id)
-    if not db_role or not db_role.id or db_role.deleted:
-        raise RoleNotFoundError()
-
-    permissions = [
-        ValidPermissions(link.permission.name)
-        for link in db_role.role_permission_links
-        if link.permission is not None
-    ]
-
-    return RoleRead(
-        id=db_role.id,
-        name=db_role.name,
-        created_at=db_role.created_at,
-        updated_at=db_role.updated_at,
-        deleted=db_role.deleted,
-        permissions=permissions
-    )
-
-
 @router.put("/{role_id}", response_class=RedirectResponse)
 def update_role(
     role: RoleUpdate = Depends(RoleUpdate.as_form),