fix(rbac): review and linting

jordan-umusu · jordan-umusu · commit 4fce9c2b0300 · 2026-02-04T11:25:25.000-05:00
diff --git a/alembic/versions/4bb7e59026f3_drop_role_from_membership_and_.py b/alembic/versions/4bb7e59026f3_drop_role_from_membership_and_.py
@@ -22,7 +22,55 @@
 
 def upgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.add_column("invitation", sa.Column("role_id", sa.UUID(), nullable=False))
+    #
+    # This migration converts enum-based roles to the new Role table.
+    # Existing invitations are migrated by looking up the corresponding role
+    # by slug in the organization's role table.
+    #
+    # Enum to slug mapping:
+    # - WorkspaceRole: VIEWER -> viewer, EDITOR -> editor, ADMIN -> admin
+    # - OrgRole: MEMBER -> member, ADMIN -> admin, OWNER -> owner
+
+    # Step 1: Add role_id columns as nullable
+    op.add_column("invitation", sa.Column("role_id", sa.UUID(), nullable=True))
+    op.add_column(
+        "organization_invitation", sa.Column("role_id", sa.UUID(), nullable=True)
+    )
+
+    # Step 2: Populate role_id for existing workspace invitations
+    # Map old enum values to role slugs and join through workspace to get org
+    op.execute(
+        """
+        UPDATE invitation i
+        SET role_id = r.id
+        FROM workspace w, role r
+        WHERE i.workspace_id = w.id
+          AND r.organization_id = w.organization_id
+          AND r.slug = LOWER(i.role::text)
+        """
+    )
+
+    # Step 3: Populate role_id for existing organization invitations
+    op.execute(
+        """
+        UPDATE organization_invitation oi
+        SET role_id = r.id
+        FROM role r
+        WHERE r.organization_id = oi.organization_id
+          AND r.slug = LOWER(oi.role::text)
+        """
+    )
+
+    # Step 4: Delete any invitations that couldn't be migrated (no matching role)
+    # This handles edge cases where roles weren't seeded yet
+    op.execute("DELETE FROM invitation WHERE role_id IS NULL")
+    op.execute("DELETE FROM organization_invitation WHERE role_id IS NULL")
+
+    # Step 5: Alter columns to NOT NULL
+    op.alter_column("invitation", "role_id", nullable=False)
+    op.alter_column("organization_invitation", "role_id", nullable=False)
+
+    # Step 6: Create indexes and foreign keys
     op.create_index(
         op.f("ix_invitation_role_id"), "invitation", ["role_id"], unique=False
     )
@@ -34,11 +82,6 @@ def upgrade() -> None:
         ["id"],
         ondelete="RESTRICT",
     )
-    op.drop_column("invitation", "role")
-    op.drop_column("membership", "role")
-    op.add_column(
-        "organization_invitation", sa.Column("role_id", sa.UUID(), nullable=False)
-    )
     op.create_index(
         op.f("ix_organization_invitation_role_id"),
         "organization_invitation",
@@ -53,8 +96,14 @@ def upgrade() -> None:
         ["id"],
         ondelete="RESTRICT",
     )
+
+    # Step 7: Drop old enum-based role columns
+    op.drop_column("invitation", "role")
+    op.drop_column("membership", "role")
     op.drop_column("organization_invitation", "role")
     op.drop_column("organization_membership", "role")
+
+    # Step 8: Drop the old enum types
     sa.Enum("MEMBER", "ADMIN", "OWNER", name="orgrole").drop(op.get_bind())
     sa.Enum("VIEWER", "EDITOR", "ADMIN", name="workspacerole").drop(op.get_bind())
     # ### end Alembic commands ###
diff --git a/tracecat/authz/seeding.py b/tracecat/authz/seeding.py
@@ -25,7 +25,7 @@
     ORG_OWNER_SCOPES,
     VIEWER_SCOPES,
 )
-from tracecat.db.models import Role, RoleScope, Scope
+from tracecat.db.models import Organization, Role, RoleScope, Scope
 from tracecat.logger import logger
 
 if TYPE_CHECKING:
@@ -360,7 +360,9 @@ async def seed_system_roles_for_org(
     # Get all system scope names -> IDs
     scope_stmt = select(Scope.id, Scope.name).where(Scope.organization_id.is_(None))
     scope_result = await session.execute(scope_stmt)
-    scope_name_to_id: dict[str, UUID] = {name: id_ for id_, name in scope_result.all()}
+    scope_name_to_id: dict[str, UUID] = {
+        name: id_ for id_, name in scope_result.tuples().all()
+    }
 
     roles_created = 0
 
@@ -445,8 +447,6 @@ async def seed_system_roles_for_all_orgs(session: AsyncSession) -> dict[UUID, in
     Returns:
         Dict mapping organization_id to number of roles created
     """
-    from tracecat.db.models import Organization
-
     logger.info("Seeding system roles for all organizations")
 
     # Get all organizations
diff --git a/tracecat/authz/service.py b/tracecat/authz/service.py
@@ -3,7 +3,7 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
 
-from sqlalchemy import select
+from sqlalchemy import or_, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from tracecat.auth.types import Role
@@ -83,17 +83,37 @@ async def list_workspace_members(
             return []
 
         # Get role assignments for these users in this workspace
+        # Include both workspace-specific assignments and org-wide assignments (workspace_id IS NULL)
         user_ids = [u.id for u in users]
         role_stmt = (
-            select(UserRoleAssignment.user_id, RoleModel.slug)
+            select(
+                UserRoleAssignment.user_id,
+                UserRoleAssignment.workspace_id,
+                RoleModel.slug,
+            )
             .join(RoleModel, UserRoleAssignment.role_id == RoleModel.id)
             .where(
                 UserRoleAssignment.user_id.in_(user_ids),
-                UserRoleAssignment.workspace_id == workspace_id,
+                or_(
+                    UserRoleAssignment.workspace_id == workspace_id,
+                    UserRoleAssignment.workspace_id.is_(None),
+                ),
             )
         )
         role_result = await self.session.execute(role_stmt)
-        user_role_map = {row[0]: row[1] for row in role_result.all()}
+
+        # Build map preferring workspace-specific assignments over org-wide
+        user_role_map: dict[UserID, str] = {}
+        for uid, ws_id, slug in role_result.tuples().all():
+            if slug is None:
+                # Skip assignments without a slug (custom roles)
+                continue
+            if ws_id is not None:
+                # Workspace-specific assignment takes precedence
+                user_role_map[uid] = slug
+            elif uid not in user_role_map:
+                # Org-wide assignment as fallback
+                user_role_map[uid] = slug
 
         return [
             WorkspaceMember(
diff --git a/tracecat/organization/schemas.py b/tracecat/organization/schemas.py
@@ -1,7 +1,8 @@
 from datetime import datetime
+from typing import Self
 from uuid import UUID
 
-from pydantic import BaseModel, EmailStr
+from pydantic import BaseModel, EmailStr, model_validator
 
 from tracecat.identifiers import OrganizationID, UserID
 from tracecat.invitations.enums import InvitationStatus
@@ -48,6 +49,13 @@ class OrgInvitationCreate(BaseModel):
     role_slug: str | None = None
     """Slug of the role to grant (e.g., 'admin', 'member', 'owner')."""
 
+    @model_validator(mode="after")
+    def validate_role_specified(self) -> Self:
+        """Ensure at least one of role_id or role_slug is provided."""
+        if self.role_id is None and self.role_slug is None:
+            raise ValueError("Either role_id or role_slug must be provided")
+        return self
+
 
 class OrgInvitationRead(BaseModel):
     """Response model for organization invitation."""
diff --git a/tracecat/registry/sync/jobs.py b/tracecat/registry/sync/jobs.py
@@ -9,6 +9,7 @@
 
 import tracecat_registry
 from packaging.version import Version
+from sqlalchemy.exc import DBAPIError
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from tracecat.authz.seeding import seed_registry_scopes_bulk
@@ -196,7 +197,7 @@ async def _seed_registry_scopes(
         inserted = await seed_registry_scopes_bulk(session, action_keys)
         await session.commit()
         logger.info("Registry scopes seeded", inserted=inserted, total=len(action_keys))
-    except Exception as e:
+    except DBAPIError as e:
         logger.warning("Failed to seed registry scopes", error=str(e))
-        # Don't fail the sync if scope seeding fails
+        # Don't fail the sync if scope seeding fails due to DB errors
         await session.rollback()
diff --git a/tracecat/workspaces/router.py b/tracecat/workspaces/router.py
@@ -6,7 +6,7 @@
     HTTPException,
     status,
 )
-from sqlalchemy import select
+from sqlalchemy import or_, select
 from sqlalchemy.exc import IntegrityError, NoResultFound
 
 from tracecat.auth.credentials import RoleACL
@@ -253,17 +253,37 @@ async def list_workspace_memberships(
         return []
 
     # Look up roles from RBAC tables
+    # Include both workspace-specific assignments and org-wide assignments (workspace_id IS NULL)
     user_ids = [m.user_id for m in memberships]
     role_stmt = (
-        select(UserRoleAssignment.user_id, RoleModel.slug)
+        select(
+            UserRoleAssignment.user_id,
+            UserRoleAssignment.workspace_id,
+            RoleModel.slug,
+        )
         .join(RoleModel, UserRoleAssignment.role_id == RoleModel.id)
         .where(
             UserRoleAssignment.user_id.in_(user_ids),
-            UserRoleAssignment.workspace_id == workspace_id,
+            or_(
+                UserRoleAssignment.workspace_id == workspace_id,
+                UserRoleAssignment.workspace_id.is_(None),
+            ),
         )
     )
     role_result = await session.execute(role_stmt)
-    user_role_map = {row[0]: row[1] for row in role_result.all()}
+
+    # Build map preferring workspace-specific assignments over org-wide
+    user_role_map: dict[UserID, str] = {}
+    for uid, ws_id, slug in role_result.tuples().all():
+        if slug is None:
+            # Skip assignments without a slug (custom roles)
+            continue
+        if ws_id is not None:
+            # Workspace-specific assignment takes precedence
+            user_role_map[uid] = slug
+        elif uid not in user_role_map:
+            # Org-wide assignment as fallback
+            user_role_map[uid] = slug
 
     return [
         WorkspaceMembershipRead(
@@ -359,16 +379,30 @@ async def get_workspace_membership(
     membership = membership_with_org.membership
 
     # Look up role from RBAC tables
+    # Include both workspace-specific assignments and org-wide assignments (workspace_id IS NULL)
     role_stmt = (
-        select(RoleModel.slug)
+        select(UserRoleAssignment.workspace_id, RoleModel.slug)
         .join(UserRoleAssignment, UserRoleAssignment.role_id == RoleModel.id)
         .where(
             UserRoleAssignment.user_id == user_id,
-            UserRoleAssignment.workspace_id == workspace_id,
+            or_(
+                UserRoleAssignment.workspace_id == workspace_id,
+                UserRoleAssignment.workspace_id.is_(None),
+            ),
         )
     )
     role_result = await session.execute(role_stmt)
-    slug = role_result.scalar_one_or_none()
+    rows = role_result.tuples().all()
+
+    # Prefer workspace-specific assignment over org-wide
+    slug: str | None = None
+    for ws_id, role_slug in rows:
+        if ws_id is not None:
+            # Workspace-specific assignment takes precedence
+            slug = role_slug
+            break
+        # Org-wide assignment as fallback
+        slug = role_slug
     workspace_role = SLUG_TO_WORKSPACE_ROLE.get(slug or "", WorkspaceRole.VIEWER)
 
     return WorkspaceMembershipRead(
