Fix _get_members_set to get role name (#1993)

Author: paulnoirel

libs/labelbox/src/labelbox/schema/user_group.py

@@ -230,7 +230,7 @@ def update(self) -> UserGroup:
         Raises:
             ValueError: If group ID or name is not set, or if projects don't exist.
             ResourceNotFoundError: If the group or projects are not found.
-            UnprocessableEntityError: If user validation fails.
+            UnprocessableEntityError: If user validation fails or users have workspace-level org roles.
         """
         if not self.id:
             raise ValueError("Group id is required")
@@ -247,8 +247,11 @@ def update(self) -> UserGroup:
                 )
 
         # Filter eligible users and build user roles
-        eligible_users = self._filter_project_based_users()
-        user_roles = self._build_user_roles(eligible_users)
+        try:
+            eligible_users = self._filter_project_based_users()
+            user_roles = self._build_user_roles(eligible_users)
+        except ValueError as e:
+            raise UnprocessableEntityError(str(e)) from e
 
         query = """
         mutation UpdateUserGroupPyApi($id: ID!, $name: String!, $description: String, $color: String!, $projectIds: [ID!]!, $userRoles: [UserRoleInput!], $notifyMembers: Boolean) {
@@ -312,7 +315,7 @@ def create(self) -> UserGroup:
 
         Raises:
             ValueError: If group already has ID, name is invalid, or projects don't exist.
-            ResourceCreationError: If creation fails or user validation fails.
+            ResourceCreationError: If creation fails, user validation fails, or users have workspace-level org roles.
             ResourceConflict: If a group with the same name already exists.
         """
         if self.id:
@@ -330,8 +333,11 @@ def create(self) -> UserGroup:
                 )
 
         # Filter eligible users and build user roles
-        eligible_users = self._filter_project_based_users()
-        user_roles = self._build_user_roles(eligible_users)
+        try:
+            eligible_users = self._filter_project_based_users()
+            user_roles = self._build_user_roles(eligible_users)
+        except ValueError as e:
+            raise ResourceCreationError(str(e)) from e
 
         query = """
         mutation CreateUserGroupPyApi($name: String!, $description: String, $color: String!, $projectIds: [ID!]!, $userRoles: [UserRoleInput!]!, $notifyMembers: Boolean, $roleId: String, $searchQuery: AlignerrSearchServiceQuery) {
@@ -496,11 +502,14 @@ def get_user_groups(
     def _filter_project_based_users(self) -> Set[User]:
         """Filter users to only include users eligible for UserGroups.
 
-        Filters out users with specific admin organization roles that cannot be
-        added to UserGroups. Most users should be eligible.
+        Only project-based users (org role "NONE") can be added to UserGroups.
+        Users with any workspace-level organization role cannot be added.
 
         Returns:
             Set of users that are eligible to be added to the group.
+
+        Raises:
+            ValueError: If any user has a workspace-level organization role.
         """
         all_users = set()
         for member in self.members:
@@ -509,63 +518,52 @@ def _filter_project_based_users(self) -> Set[User]:
         if not all_users:
             return set()
 
-        user_ids = [user.uid for user in all_users]
-        query = """
-        query CheckUserOrgRolesPyApi($userIds: [ID!]!) {
-            users(where: {id_in: $userIds}) {
-                id
-                orgRole { id name }
-            }
-        }
-        """
+        # Check each user's org role directly
+        invalid_users = []
+        eligible_users = set()
 
-        try:
-            result = self.client.execute(query, {"userIds": user_ids})
-            if not result or "users" not in result:
-                return all_users  # Fallback: let server handle validation
-
-            # Check for users with org roles that cannot be used in UserGroups
-            # Only users with no org role (project-based users) can be assigned to UserGroups
-            eligible_user_ids = set()
-            invalid_users = []
-
-            for user_data in result["users"]:
-                org_role = user_data.get("orgRole")
-                user_id = user_data["id"]
-                user_email = user_data.get("email", "unknown")
-
-                if org_role is None:
-                    # Users with no org role (project-based users) are eligible
-                    eligible_user_ids.add(user_id)
+        for user in all_users:
+            try:
+                # Get the user's organization role directly
+                org_role = user.org_role()
+                if org_role is None or org_role.name.upper() == "NONE":
+                    # Users with no org role or "NONE" role are project-based and eligible
+                    eligible_users.add(user)
                 else:
-                    # Users with ANY workspace org role cannot be assigned to UserGroups
+                    # Users with any workspace org role cannot be assigned to UserGroups
                     invalid_users.append(
                         {
-                            "id": user_id,
-                            "email": user_email,
-                            "org_role": org_role.get("name"),
+                            "id": user.uid,
+                            "email": getattr(user, "email", "unknown"),
+                            "org_role": org_role.name,
                         }
                     )
+            except Exception as e:
+                # If we can't determine the user's role, treat as invalid for safety
+                invalid_users.append(
+                    {
+                        "id": user.uid,
+                        "email": getattr(user, "email", "unknown"),
+                        "org_role": f"unknown (error: {str(e)})",
+                    }
+                )
 
-            # Raise error if any invalid users found
-            if invalid_users:
-                error_details = []
-                for user in invalid_users:
-                    error_details.append(
-                        f"User {user['id']} ({user['email']}) has org role '{user['org_role']}'"
-                    )
-
-                raise ValueError(
-                    f"Cannot create UserGroup with users who have organization roles. "
-                    f"Only project-based users (no org role) can be assigned to UserGroups.\n"
-                    f"Invalid users:\n"
-                    + "\n".join(f"  • {detail}" for detail in error_details)
+        # Raise error if any invalid users found
+        if invalid_users:
+            error_details = []
+            for user_info in invalid_users:
+                error_details.append(
+                    f"User {user_info['id']} ({user_info['email']}) has org role '{user_info['org_role']}'"
                 )
 
-            return {user for user in all_users if user.uid in eligible_user_ids}
+            raise ValueError(
+                f"Cannot create UserGroup with users who have organization roles. "
+                f"Only project-based users (no org role or role 'NONE') can be assigned to UserGroups.\n"
+                f"Invalid users:\n"
+                + "\n".join(f"  • {detail}" for detail in error_details)
+            )
 
-        except Exception:
-            return all_users  # Fallback: let server handle validation
+        return eligible_users
 
     def _build_user_roles(
         self, eligible_users: Set[User]
@@ -679,6 +677,12 @@ def _get_members_set(
         member_nodes = members_data.get("nodes", [])
         user_group_roles = members_data.get("userGroupRoles", [])
 
+        # Get all roles to map IDs to names
+        from labelbox.schema.role import get_roles
+
+        all_roles = get_roles(self.client)
+        role_id_to_role = {role.uid: role for role in all_roles.values()}
+
         # Create a mapping from userId to roleId
         user_role_mapping = {
             role_data["userId"]: role_data["roleId"]
@@ -694,15 +698,9 @@ def _get_members_set(
 
             # Get the role for this user from the mapping
             role_id = user_role_mapping.get(node["id"])
-            if role_id:
-                # We need to fetch the role details since we only have the roleId
-                # For now, create a minimal Role object with just the ID
-                role_values: defaultdict[str, Any] = defaultdict(lambda: None)
-                role_values["id"] = role_id
-                # We don't have the role name from this response, so we'll leave it as None
-                # The Role object will fetch the name when needed
-                role = Role(self.client, role_values)
-
+            if role_id and role_id in role_id_to_role:
+                # Use the actual Role object with proper name resolution
+                role = role_id_to_role[role_id]
                 members.add(UserGroupMember(user=user, role=role))
 
         return members