Default permissions will be global, but default roles will be organization-specific

Author: chriscarrollsmith

main.py

@@ -20,7 +20,8 @@
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # Optional startup logic
-    set_up_db(drop=False)
+    # TODO: Set drop=False in production
+    set_up_db(drop=True)
     yield
     # Optional shutdown logic
 

routers/organization.py

@@ -5,7 +5,7 @@
 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, utc_time
+from utils.models import Organization, User, Role, Permission, UserOrganizationLink, ValidPermissions, utc_time
 from datetime import datetime
 from sqlalchemy import and_
 from utils.role_org import get_organization, check_user_permission
@@ -119,6 +119,15 @@ def create_organization(
     session.commit()
     session.refresh(db_org)
 
+    # 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()
+
     owner_role = session.exec(
         select(Role).where(
             and_(

tests/test_models.py

@@ -39,24 +39,51 @@ def get_session():
         yield session
 
 
-def create_roles(session):
+def create_default_roles(session, organization_id: int, check_first: bool = True):
     """
-    Create default roles in the database if they do not exist.
+    Create default roles for an organization in the database if they do not exist.
     """
     roles_in_db = []
     for role_name in default_roles:
         db_role = session.exec(select(Role).where(
-            Role.name == role_name)).first()
+            Role.name == role_name,
+            Role.organization_id == organization_id
+        )).first()
         if not db_role:
-            db_role = Role(name=role_name)
+            db_role = Role(name=role_name, organization_id=organization_id)
             session.add(db_role)
         roles_in_db.append(db_role)
+
+    # Create RolePermissionLink for Owner and Administrator roles
+    for role in roles_in_db[:2]:
+        permissions = session.exec(select(Permission)).all()
+        for permission in permissions:
+            # Check if the role already has the permission
+            if check_first:
+                db_role_permission_link: RolePermissionLink | None = session.exec(select(RolePermissionLink).where(
+                    RolePermissionLink.role_id == role.id,
+                    RolePermissionLink.permission_id == permission.id
+                )).first()
+            else:
+                db_role_permission_link = None
+
+            # Skip giving DELETE_ORGANIZATION permission to Administrator
+            if not db_role_permission_link and not (
+                permission == ValidPermissions.DELETE_ORGANIZATION and
+                role.name == "Administrator"
+            ):
+                role_permission_link = RolePermissionLink(
+                    role_id=role.id,
+                    permission_id=permission.id
+                )
+                session.add(role_permission_link)
+
     return roles_in_db
 
 
-def create_permissions(session, roles_in_db):
+def create_permissions(session):
     """
-    Create default permissions and link them to roles in the database.
+    Create default permissions.
     """
     for permission in ValidPermissions:
         db_permission = session.exec(select(Permission).where(
@@ -65,17 +92,6 @@ def create_permissions(session, roles_in_db):
             db_permission = Permission(name=permission)
             session.add(db_permission)
 
-        # Create RolePermissionLink for Owner and Administrator
-        for role in roles_in_db[:2]:
-            db_role_permission_link = session.exec(select(RolePermissionLink).where(
-                RolePermissionLink.role_id == role.id,
-                RolePermissionLink.permission_id == db_permission.id)).first()
-            if not db_role_permission_link:
-                if not (permission == ValidPermissions.DELETE_ORGANIZATION and role.name == "Administrator"):
-                    role_permission_link = RolePermissionLink(
-                        role_id=role.id, permission_id=db_permission.id)
-                    session.add(role_permission_link)
-
 
 def set_up_db(drop: bool = False):
     """
@@ -85,10 +101,9 @@ def set_up_db(drop: bool = False):
     if drop:
         SQLModel.metadata.drop_all(engine)
     SQLModel.metadata.create_all(engine)
+    # Create default permissions
     with Session(engine) as session:
-        roles_in_db = create_roles(session)
-        session.commit()
-        create_permissions(session, roles_in_db)
+        create_permissions(session)
         session.commit()
     engine.dispose()
 

utils/db.py

@@ -62,13 +62,23 @@ class Organization(SQLModel, table=True):
 
 
 class Role(SQLModel, table=True):
+    """
+    Represents a role within an organization.
+
+    Attributes:
+        id: Primary key.
+        name: The name of the role.
+        organization_id: Foreign key to the associated organization.
+        created_at: Timestamp when the role was created.
+        updated_at: Timestamp when the role was last updated.
+    """
     id: Optional[int] = Field(default=None, primary_key=True)
     name: str
     organization_id: int = Field(foreign_key="organization.id")
     created_at: datetime = Field(default_factory=utc_time)
     updated_at: datetime = Field(default_factory=utc_time)
 
-    organization: Organization = Relationship(back_populates="roles")
+    organization: "Organization" = Relationship(back_populates="roles")
     user_links: List[UserOrganizationLink] = Relationship(
         back_populates="role")
     permissions: List["Permission"] = Relationship(
@@ -78,6 +88,9 @@ class Role(SQLModel, table=True):
 
 
 class Permission(SQLModel, table=True):
+    """
+    Represents a permission that can be assigned to a role.
+    """
     id: Optional[int] = Field(default=None, primary_key=True)
     name: ValidPermissions = Field(
         sa_column=Column(SQLAlchemyEnum(ValidPermissions, create_type=False)))

utils/models.py

@@ -1,6 +1,6 @@
 from typing import List
 from sqlmodel import Session, select
-from sqlalchemy import and_
+from sqlalchemy import and_, or_
 from fastapi import HTTPException
 from utils.models import Organization, Role, UserOrganizationLink, RolePermissionLink, Permission, ValidPermissions
 
@@ -150,6 +150,11 @@ def get_organization_roles(
     Returns:
         List of Role objects with their associated permissions
     """
-    query = select(Role).where(Role.organization_id == organization_id)
+    query = select(Role).where(
+        or_(
+            Role.organization_id == organization_id,
+            Role.organization_id == None
+        )
+    )
 
     return list(session.exec(query))