fix(invitations): create user role assignment on invite

Author: jordan-umusu

alembic/versions/4bb7e59026f3_drop_role_from_membership_and_.py

@@ -111,18 +111,31 @@ def upgrade() -> None:
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
+    #
+    # This downgrade restores the original enum-based role columns by:
+    # 1. Creating the enum types
+    # 2. Adding role columns as nullable
+    # 3. Populating role values from role_id via the role table
+    # 4. Making the columns NOT NULL
+    # 5. Dropping the role_id columns and constraints
+    #
+    # Note: The original schema did NOT have server_default on invitation/org_invitation
+    # role columns, but membership/org_membership did have defaults.
+
+    # Step 1: Create enum types
     sa.Enum("VIEWER", "EDITOR", "ADMIN", name="workspacerole").create(op.get_bind())
     sa.Enum("MEMBER", "ADMIN", "OWNER", name="orgrole").create(op.get_bind())
+
+    # Step 2: Add role columns as nullable
     op.add_column(
         "organization_membership",
         sa.Column(
             "role",
             postgresql.ENUM(
                 "MEMBER", "ADMIN", "OWNER", name="orgrole", create_type=False
             ),
-            server_default=sa.text("'MEMBER'::orgrole"),
             autoincrement=False,
-            nullable=False,
+            nullable=True,
         ),
     )
     op.add_column(
@@ -133,28 +146,18 @@ def downgrade() -> None:
                 "MEMBER", "ADMIN", "OWNER", name="orgrole", create_type=False
             ),
             autoincrement=False,
-            nullable=False,
+            nullable=True,
         ),
     )
-    op.drop_constraint(
-        op.f("fk_organization_invitation_role_id_role"),
-        "organization_invitation",
-        type_="foreignkey",
-    )
-    op.drop_index(
-        op.f("ix_organization_invitation_role_id"), table_name="organization_invitation"
-    )
-    op.drop_column("organization_invitation", "role_id")
     op.add_column(
         "membership",
         sa.Column(
             "role",
             postgresql.ENUM(
                 "VIEWER", "EDITOR", "ADMIN", name="workspacerole", create_type=False
             ),
-            server_default=sa.text("'EDITOR'::workspacerole"),
             autoincrement=False,
-            nullable=False,
+            nullable=True,
         ),
     )
     op.add_column(
@@ -165,9 +168,77 @@ def downgrade() -> None:
                 "VIEWER", "EDITOR", "ADMIN", name="workspacerole", create_type=False
             ),
             autoincrement=False,
-            nullable=False,
+            nullable=True,
         ),
     )
+
+    # Step 3: Populate role values from role_id via the role table
+    # Organization membership: role_id -> role.slug -> uppercase enum
+    op.execute(
+        """
+        UPDATE organization_membership om
+        SET role = UPPER(r.slug)::orgrole
+        FROM role r
+        WHERE om.role_id = r.id
+        """
+    )
+
+    # Organization invitation: role_id -> role.slug -> uppercase enum
+    op.execute(
+        """
+        UPDATE organization_invitation oi
+        SET role = UPPER(r.slug)::orgrole
+        FROM role r
+        WHERE oi.role_id = r.id
+        """
+    )
+
+    # Workspace membership: role_id -> role.slug -> uppercase enum
+    op.execute(
+        """
+        UPDATE membership m
+        SET role = UPPER(r.slug)::workspacerole
+        FROM role r
+        WHERE m.role_id = r.id
+        """
+    )
+
+    # Workspace invitation: role_id -> role.slug -> uppercase enum
+    op.execute(
+        """
+        UPDATE invitation i
+        SET role = UPPER(r.slug)::workspacerole
+        FROM role r
+        WHERE i.role_id = r.id
+        """
+    )
+
+    # Step 4: Make columns NOT NULL and add server_default where original schema had it
+    op.alter_column("organization_membership", "role", nullable=False)
+    op.alter_column(
+        "organization_membership",
+        "role",
+        server_default=sa.text("'MEMBER'::orgrole"),
+    )
+    op.alter_column("organization_invitation", "role", nullable=False)
+    op.alter_column("membership", "role", nullable=False)
+    op.alter_column(
+        "membership",
+        "role",
+        server_default=sa.text("'EDITOR'::workspacerole"),
+    )
+    op.alter_column("invitation", "role", nullable=False)
+
+    # Step 5: Drop role_id columns and constraints
+    op.drop_constraint(
+        op.f("fk_organization_invitation_role_id_role"),
+        "organization_invitation",
+        type_="foreignkey",
+    )
+    op.drop_index(
+        op.f("ix_organization_invitation_role_id"), table_name="organization_invitation"
+    )
+    op.drop_column("organization_invitation", "role_id")
     op.drop_constraint(
         op.f("fk_invitation_role_id_role"), "invitation", type_="foreignkey"
     )

tracecat/authz/seeding.py

@@ -244,30 +244,22 @@ async def seed_registry_scope(
 
     Creates a scope for `action:{action_key}:execute` if it doesn't exist.
     Registry scopes have organization_id=NULL and source='registry'.
+    Uses upsert (ON CONFLICT DO NOTHING) for concurrency safety.
 
     Args:
         session: Database session
         action_key: The action key (e.g., "tools.okta.list_users")
         description: Optional description for the scope
 
     Returns:
-        The created or existing Scope, or None if upsert had no effect
+        The created or existing Scope
     """
     scope_name = f"action:{action_key}:execute"
+    scope_id = uuid4()
 
-    # Check if scope already exists
-    stmt = select(Scope).where(
-        Scope.name == scope_name, Scope.organization_id.is_(None)
-    )
-    result = await session.execute(stmt)
-    existing = result.scalar_one_or_none()
-
-    if existing:
-        return existing
-
-    # Create new scope
-    scope = Scope(
-        id=uuid4(),
+    # Use upsert for concurrency safety
+    stmt = pg_insert(Scope).values(
+        id=scope_id,
         name=scope_name,
         resource="action",
         action="execute",
@@ -276,10 +268,24 @@ async def seed_registry_scope(
         source_ref=action_key,
         organization_id=None,
     )
-    session.add(scope)
+    stmt = stmt.on_conflict_do_nothing(
+        index_elements=["name"], index_where=Scope.organization_id.is_(None)
+    )
+    result = await session.execute(stmt)
     await session.flush()
 
-    logger.debug("Registry scope created", scope_name=scope_name, action_key=action_key)
+    # Re-query to get the scope (whether newly inserted or already existing)
+    select_stmt = select(Scope).where(
+        Scope.name == scope_name, Scope.organization_id.is_(None)
+    )
+    select_result = await session.execute(select_stmt)
+    scope = select_result.scalar_one_or_none()
+
+    if result.rowcount and result.rowcount > 0:  # pyright: ignore[reportAttributeAccessIssue]
+        logger.debug(
+            "Registry scope created", scope_name=scope_name, action_key=action_key
+        )
+
     return scope
 
 
@@ -343,8 +349,9 @@ async def seed_system_roles_for_org(
 ) -> int:
     """Seed system roles (Admin, Editor, Viewer) for an organization.
 
-    Creates the three system roles with their associated scopes if they don't exist.
+    Creates the system roles with their associated scopes if they don't exist.
     System roles are identified by their well-known slugs.
+    Uses upsert (ON CONFLICT DO NOTHING) for concurrency safety.
 
     Args:
         session: Database session
@@ -366,36 +373,52 @@ async def seed_system_roles_for_org(
 
     roles_created = 0
 
-    for slug, name, description, scope_names in PRESET_ROLE_DEFINITIONS:
-        # Check if role already exists
-        existing_stmt = select(Role.id).where(
-            Role.organization_id == organization_id,
-            Role.slug == slug,
+    # Prepare role values with pre-generated IDs
+    role_values = []
+    role_id_by_slug: dict[str, UUID] = {}
+    for slug, name, description, _ in PRESET_ROLE_DEFINITIONS:
+        role_id = uuid4()
+        role_id_by_slug[slug] = role_id
+        role_values.append(
+            {
+                "id": role_id,
+                "name": name,
+                "slug": slug,
+                "description": description,
+                "organization_id": organization_id,
+                "created_by": None,  # System-created
+            }
         )
-        existing_result = await session.execute(existing_stmt)
-        existing_role_id = existing_result.scalar_one_or_none()
 
-        if existing_role_id is not None:
-            logger.debug(
-                "System role already exists",
+    # Bulk upsert roles - concurrency safe
+    role_stmt = pg_insert(Role).values(role_values)
+    role_stmt = role_stmt.on_conflict_do_nothing(
+        index_elements=["organization_id", "slug"]
+    )
+    result = await session.execute(role_stmt)
+    roles_created = result.rowcount if result.rowcount else 0  # pyright: ignore[reportAttributeAccessIssue]
+
+    # Re-query to get actual role IDs (may differ if roles already existed)
+    existing_roles_stmt = select(Role.id, Role.slug).where(
+        Role.organization_id == organization_id,
+        Role.slug.in_([slug for slug, _, _, _ in PRESET_ROLE_DEFINITIONS]),
+    )
+    existing_roles_result = await session.execute(existing_roles_stmt)
+    actual_role_id_by_slug: dict[str | None, UUID] = {
+        slug: role_id for role_id, slug in existing_roles_result.tuples().all()
+    }
+
+    # Link scopes to roles
+    for slug, _, _, scope_names in PRESET_ROLE_DEFINITIONS:
+        role_id = actual_role_id_by_slug.get(slug)
+        if role_id is None:
+            logger.warning(
+                "Role not found after upsert",
                 slug=slug,
                 organization_id=organization_id,
             )
             continue
 
-        # Create the role
-        role = Role(
-            id=uuid4(),
-            name=name,
-            slug=slug,
-            description=description,
-            organization_id=organization_id,
-            created_by=None,  # System-created
-        )
-        session.add(role)
-        await session.flush()  # Get the role ID
-
-        # Link scopes to the role
         role_scope_values = []
         for scope_name in scope_names:
             scope_id = scope_name_to_id.get(scope_name)
@@ -406,21 +429,13 @@ async def seed_system_roles_for_org(
                     role_slug=slug,
                 )
                 continue
-            role_scope_values.append({"role_id": role.id, "scope_id": scope_id})
+            role_scope_values.append({"role_id": role_id, "scope_id": scope_id})
 
         if role_scope_values:
             role_scope_stmt = pg_insert(RoleScope).values(role_scope_values)
             role_scope_stmt = role_scope_stmt.on_conflict_do_nothing()
             await session.execute(role_scope_stmt)
 
-        roles_created += 1
-        logger.debug(
-            "System role created",
-            slug=slug,
-            organization_id=organization_id,
-            num_scopes=len(role_scope_values),
-        )
-
     await session.commit()
     logger.info(
         "System roles seeded for organization",

tracecat/authz/service.py

@@ -70,6 +70,12 @@ async def list_workspace_members(
 
         Roles are looked up from UserRoleAssignment table.
         """
+        # Get workspace to determine organization_id
+        workspace = await self.session.get(Workspace, workspace_id)
+        if workspace is None:
+            return []
+        organization_id = workspace.organization_id
+
         # Get all members of the workspace
         members_stmt = (
             select(User)
@@ -84,6 +90,7 @@ async def list_workspace_members(
 
         # Get role assignments for these users in this workspace
         # Include both workspace-specific assignments and org-wide assignments (workspace_id IS NULL)
+        # Filter by organization_id to ensure we only get roles from this org
         user_ids = [u.id for u in users]
         role_stmt = (
             select(
@@ -94,6 +101,7 @@ async def list_workspace_members(
             .join(RoleModel, UserRoleAssignment.role_id == RoleModel.id)
             .where(
                 UserRoleAssignment.user_id.in_(user_ids),
+                UserRoleAssignment.organization_id == organization_id,
                 or_(
                     UserRoleAssignment.workspace_id == workspace_id,
                     UserRoleAssignment.workspace_id.is_(None),

tracecat/organization/service.py

@@ -142,14 +142,22 @@ async def accept_invitation_for_user(
             # Shouldn't reach here, but handle gracefully
             raise TracecatAuthorizationError("Invitation is no longer valid")
 
-        # Create membership
+        # Create membership (role is now tracked via UserRoleAssignment)
         membership = OrganizationMembership(
             user_id=user_id,
             organization_id=invitation.organization_id,
-            role=invitation.role,
         )
         session.add(membership)
 
+        # Create role assignment from invitation
+        role_assignment = UserRoleAssignment(
+            organization_id=invitation.organization_id,
+            user_id=user_id,
+            workspace_id=None,  # NULL = org-level assignment
+            role_id=invitation.role_id,
+        )
+        session.add(role_assignment)
+
         await session.commit()
         await session.refresh(membership)
     except TracecatAuthorizationError:
@@ -176,6 +184,8 @@ async def accept_invitation_for_user(
         )
 
     return membership
+
+
 @dataclass(frozen=True)
 class MemberRoleInfo:
     """Role information for an organization member."""