Chore: RBAC: Migrate role picker to rtkq (#116571)

What is this feature?

This PR introduce refactoring for role picker for teams and users. It uses RTK Query.
The related task: grafana/identity-access-team#1821

The previous PR with related logic: #113783

Why do we need this feature?

This refactoring is a cleaner way for handing refetching.

.github/CODEOWNERS

@@ -930,6 +930,7 @@ playwright.storybook.config.ts @grafana/grafana-frontend-platform
 /public/app/core/utils/timeRegions* @grafana/dataviz-squad
 /public/app/core/utils/urlToken.ts @grafana/identity-access-team
 /public/app/core/utils/version.ts @grafana/partner-datasources
+/public/app/core/utils/roles.ts @grafana/identity-access-team
 /public/app/dev-utils.ts @grafana/grafana-frontend-platform
 
 /public/app/features/actions/ @grafana/dataviz-squad

public/app/api/clients/roles/index.ts

@@ -0,0 +1,75 @@
+import { FetchBaseQueryError } from '@reduxjs/toolkit/query';
+
+import { generatedAPI, RoleDto } from '@grafana/api-clients/rtkq/legacy';
+import { isFetchError } from '@grafana/runtime';
+import { addDisplayNameForFixedRole, addFilteredDisplayName } from 'app/core/utils/roles';
+import { Role } from 'app/types/accessControl';
+
+const transformRolesResponse = (response: RoleDto[]): Role[] => {
+  if (!response?.length) {
+    return [];
+  }
+  return response.map(addFilteredDisplayName).map(addDisplayNameForFixedRole);
+};
+
+const transformRolesError = (error: FetchBaseQueryError) => {
+  if (isFetchError(error)) {
+    error.isHandled = true;
+  }
+  return error;
+};
+
+// rolesAPI is needed to be overriden to add transformResponse and transformErrorResponse for some endpoints.
+export const rolesAPI = generatedAPI.injectEndpoints({
+  overrideExisting: true,
+  endpoints: (build) => ({
+    listUserRoles: build.query<
+      Role[],
+      { userId: number; includeHidden?: boolean; includeMapped?: boolean; targetOrgId?: number }
+    >({
+      query: (queryArg) => ({
+        url: `/access-control/users/${queryArg.userId}/roles`,
+        params: {
+          includeHidden: queryArg.includeHidden,
+          includeMapped: queryArg.includeMapped,
+          targetOrgId: queryArg.targetOrgId,
+        },
+      }),
+      providesTags: ['access_control', 'enterprise'],
+      transformResponse: transformRolesResponse,
+      transformErrorResponse: transformRolesError,
+    }),
+    listTeamRoles: build.query<Role[], { teamId: number; includeHidden?: boolean; targetOrgId?: number }>({
+      query: (queryArg) => ({
+        url: `/access-control/teams/${queryArg.teamId}/roles`,
+        params: {
+          includeHidden: queryArg.includeHidden,
+          targetOrgId: queryArg.targetOrgId,
+        },
+      }),
+      providesTags: ['access_control', 'enterprise'],
+      transformResponse: transformRolesResponse,
+      transformErrorResponse: transformRolesError,
+    }),
+    listRoles: build.query<Role[], { delegatable?: boolean; includeHidden?: boolean; targetOrgId?: number }>({
+      query: (queryArg) => ({
+        url: '/access-control/roles',
+        params: {
+          delegatable: queryArg.delegatable,
+          includeHidden: queryArg.includeHidden,
+          targetOrgId: queryArg.targetOrgId,
+        },
+      }),
+      providesTags: ['access_control', 'enterprise'],
+      transformResponse: transformRolesResponse,
+    }),
+  }),
+});
+
+export const {
+  useListTeamRolesQuery,
+  useSetTeamRolesMutation,
+  useListUserRolesQuery,
+  useSetUserRolesMutation,
+  useListRolesQuery,
+} = rolesAPI;

public/app/core/components/RolePicker/RoleMenuGroupsSection.tsx

@@ -1,13 +1,13 @@
 import { forwardRef, useCallback, useState } from 'react';
 
 import { useStyles2, getSelectStyles, useTheme2 } from '@grafana/ui';
+import { isNotDelegatable } from 'app/core/utils/roles';
 import { Role } from 'app/types/accessControl';
 
 import { RoleMenuGroupOption } from './RoleMenuGroupOption';
 import { RoleMenuOption } from './RoleMenuOption';
 import { RolePickerSubMenu } from './RolePickerSubMenu';
 import { getStyles } from './styles';
-import { isNotDelegatable } from './utils';
 
 interface RoleMenuGroupsSectionProps {
   roles: Role[];

public/app/core/components/RolePicker/RolePicker.tsx

@@ -2,7 +2,8 @@ import { FormEvent, useCallback, useEffect, useRef, useState, useSyncExternalSto
 
 import { OrgRole } from '@grafana/data';
 import { ClickOutsideWrapper, Portal, useTheme2 } from '@grafana/ui';
-import { Role } from 'app/types/accessControl';
+import { pickerStateStore } from 'app/core/utils/roles';
+import type { Role } from 'app/types/accessControl';
 
 import { RolePickerInput } from './RolePickerInput';
 import { RolePickerMenu } from './RolePickerMenu';
@@ -12,7 +13,6 @@ import {
   ROLE_PICKER_MENU_MAX_WIDTH,
   ROLE_PICKER_WIDTH,
 } from './constants';
-import { pickerStateStore } from './utils';
 
 export interface Props {
   basicRole?: OrgRole;

public/app/core/components/RolePicker/RolePickerSubMenu.tsx

@@ -4,12 +4,12 @@ import type { JSX } from 'react';
 import { Trans, t } from '@grafana/i18n';
 import { Button, ScrollContainer, Stack, useStyles2, useTheme2 } from '@grafana/ui';
 import { getSelectStyles } from '@grafana/ui/internal';
+import { isNotDelegatable } from 'app/core/utils/roles';
 import { Role } from 'app/types/accessControl';
 
 import { RoleMenuOption } from './RoleMenuOption';
 import { MENU_MAX_HEIGHT } from './constants';
 import { getStyles } from './styles';
-import { isNotDelegatable } from './utils';
 
 interface RolePickerSubMenuProps {
   options: Role[];

public/app/core/components/RolePicker/TeamRolePicker.tsx

@@ -1,11 +1,11 @@
-import { useEffect } from 'react';
-import { useAsyncFn } from 'react-use';
+import { skipToken } from '@reduxjs/toolkit/query';
+import { useMemo } from 'react';
 
+import { useListTeamRolesQuery, useSetTeamRolesMutation } from 'app/api/clients/roles';
 import { contextSrv } from 'app/core/services/context_srv';
 import { AccessControlAction, Role } from 'app/types/accessControl';
 
 import { RolePicker } from './RolePicker';
-import { fetchTeamRoles, updateTeamRoles } from './api';
 
 export interface Props {
   teamId: number;
@@ -34,6 +34,7 @@ export interface Props {
 
 export const TeamRolePicker = ({
   teamId,
+  orgId,
   roleOptions,
   disabled,
   roles,
@@ -44,36 +45,44 @@ export const TeamRolePicker = ({
   width,
   isLoading,
 }: Props) => {
-  const [{ loading, value: appliedRoles = roles || [] }, getTeamRoles] = useAsyncFn(
-    async (force = false) => {
-      try {
-        if (!force && roles) {
-          return roles;
-        }
-        if (!force && apply && Boolean(pendingRoles?.length)) {
-          return pendingRoles;
-        }
-        if (contextSrv.hasPermission(AccessControlAction.ActionTeamsRolesList) && teamId > 0) {
-          return await fetchTeamRoles(teamId);
-        }
-      } catch (e) {
-        console.error('Error fetching roles', e);
-      }
-      return [];
-    },
-    [teamId, pendingRoles, roles]
+  const hasPermission = contextSrv.hasPermission(AccessControlAction.ActionTeamsRolesList) && teamId > 0;
+
+  // In non-apply mode, always fetch to ensure we have fresh data after mutations
+  // In apply mode, only skip fetch if we have pendingRoles
+  const shouldFetch = apply ? !Boolean(pendingRoles?.length) && hasPermission : hasPermission;
+
+  const { data: fetchedRoles, isLoading: isFetching } = useListTeamRolesQuery(
+    shouldFetch ? { teamId, targetOrgId: orgId } : skipToken
   );
 
-  useEffect(() => {
-    getTeamRoles();
-  }, [getTeamRoles]);
+  const [updateTeamRoles, { isLoading: isUpdating }] = useSetTeamRolesMutation();
 
-  const onRolesChange = async (roles: Role[]) => {
+  const appliedRoles =
+    useMemo(() => {
+      if (apply && Boolean(pendingRoles?.length)) {
+        return pendingRoles;
+      }
+      // Otherwise prefer fetched data (which is always fresh due to cache invalidation)
+      // Fall back to roles prop if fetched data is not available yet
+      return fetchedRoles || roles || [];
+    }, [roles, pendingRoles, fetchedRoles, apply]) || [];
+
+  const onRolesChange = async (newRoles: Role[]) => {
     if (!apply) {
-      await updateTeamRoles(roles, teamId);
-      await getTeamRoles(true); // Force fetch from backend after update
+      try {
+        const roleUids = newRoles.map((role) => role.uid);
+        await updateTeamRoles({
+          teamId,
+          targetOrgId: orgId,
+          setTeamRolesCommand: {
+            roleUids,
+          },
+        }).unwrap();
+      } catch (error) {
+        console.error('Error updating team roles', error);
+      }
     } else if (onApplyRoles) {
-      onApplyRoles(roles);
+      onApplyRoles(newRoles);
     }
   };
 
@@ -88,7 +97,7 @@ export const TeamRolePicker = ({
       onRolesChange={onRolesChange}
       roleOptions={roleOptions}
       appliedRoles={appliedRoles}
-      isLoading={loading || isLoading}
+      isLoading={isFetching || isUpdating || isLoading}
       disabled={disabled}
       basicRoleDisabled={true}
       canUpdateRoles={canUpdateRoles}

public/app/core/components/RolePicker/UserRolePicker.tsx

@@ -1,12 +1,12 @@
-import { useEffect } from 'react';
-import { useAsyncFn } from 'react-use';
+import { skipToken } from '@reduxjs/toolkit/query';
+import { useMemo } from 'react';
 
 import { OrgRole } from '@grafana/data';
+import { useListUserRolesQuery, useSetUserRolesMutation } from 'app/api/clients/roles';
 import { contextSrv } from 'app/core/services/context_srv';
 import { AccessControlAction, Role } from 'app/types/accessControl';
 
 import { RolePicker } from './RolePicker';
-import { fetchUserRoles, updateUserRoles } from './api';
 
 export interface Props {
   basicRole: OrgRole;
@@ -54,39 +54,46 @@ export const UserRolePicker = ({
   width,
   isLoading,
 }: Props) => {
-  const [{ loading, value: appliedRoles = roles || [] }, getUserRoles] = useAsyncFn(
-    async (force = false) => {
-      try {
-        if (!force && roles) {
-          return roles;
-        }
-        if (!force && apply && Boolean(pendingRoles?.length)) {
-          return pendingRoles;
-        }
-        if (contextSrv.hasPermission(AccessControlAction.ActionUserRolesList) && userId > 0) {
-          return await fetchUserRoles(userId, orgId);
-        }
-      } catch (e) {
-        console.error('Error fetching user roles');
-      }
-      return [];
-    },
-    [orgId, userId, pendingRoles, roles]
+  const hasPermission = contextSrv.hasPermission(AccessControlAction.ActionUserRolesList) && userId > 0 && orgId;
+
+  // Determine when to fetch:
+  // - In apply mode: only fetch if we don't have roles prop AND no pendingRoles (prevents flicker)
+  // - In non-apply mode: always fetch to get fresh data after mutations
+  const shouldFetch = apply ? !roles && !Boolean(pendingRoles?.length) && hasPermission : hasPermission;
+
+  const { data: fetchedRoles, isLoading: isFetching } = useListUserRolesQuery(
+    shouldFetch ? { userId, includeHidden: true, includeMapped: true, targetOrgId: orgId } : skipToken
   );
 
-  useEffect(() => {
-    // only load roles when there is an Org selected
-    if (orgId) {
-      getUserRoles();
-    }
-  }, [getUserRoles, orgId]);
+  const [updateUserRoles, { isLoading: isUpdating }] = useSetUserRolesMutation();
 
-  const onRolesChange = async (roles: Role[]) => {
+  const appliedRoles =
+    useMemo(() => {
+      // In apply mode: prioritize pendingRoles, then roles prop (never use fetched data to prevent flicker)
+      if (apply) {
+        return pendingRoles || roles || [];
+      }
+      // In non-apply mode: prefer fetched data (fresh from cache) over roles prop
+      return fetchedRoles || roles || [];
+    }, [roles, pendingRoles, fetchedRoles, apply]) || [];
+
+  const onRolesChange = async (newRoles: Role[]) => {
     if (!apply) {
-      await updateUserRoles(roles, userId, orgId);
-      await getUserRoles(true); // Force fetch from backend after update
+      try {
+        const filteredRoles = newRoles.filter((role) => !role.mapped);
+        const roleUids = filteredRoles.map((role) => role.uid);
+        await updateUserRoles({
+          userId,
+          targetOrgId: orgId,
+          setUserRolesCommand: {
+            roleUids,
+          },
+        }).unwrap();
+      } catch (error) {
+        console.error('Error updating user roles', error);
+      }
     } else if (onApplyRoles) {
-      onApplyRoles(roles, userId, orgId);
+      onApplyRoles(newRoles, userId, orgId);
     }
   };
 
@@ -102,7 +109,7 @@ export const UserRolePicker = ({
       onRolesChange={onRolesChange}
       onBasicRoleChange={onBasicRoleChange}
       roleOptions={roleOptions}
-      isLoading={loading || isLoading}
+      isLoading={isFetching || isUpdating || isLoading}
       disabled={disabled}
       basicRoleDisabled={basicRoleDisabled}
       basicRoleDisabledMessage={basicRoleDisabledMessage}

public/app/core/components/RolePicker/api.ts

@@ -1,7 +1,16 @@
-import { getBackendSrv, isFetchError } from '@grafana/runtime';
-import { Role } from 'app/types/accessControl';
+/**
+ * @deprecated These functions are legacy API calls. For new code, use the RTK Query API from:
+ * `app/api/clients/roles` which provides:
+ * - useListTeamRolesQuery
+ * - useSetTeamRolesMutation
+ * - useListUserRolesQuery
+ * - useSetUserRolesMutation
+ * - useListRolesQuery
+ */
 
-import { addDisplayNameForFixedRole, addFilteredDisplayName } from './utils';
+import { getBackendSrv } from '@grafana/runtime';
+import { addDisplayNameForFixedRole, addFilteredDisplayName } from 'app/core/utils/roles';
+import { Role } from 'app/types/accessControl';
 
 export const fetchRoleOptions = async (orgId?: number): Promise<Role[]> => {
   let rolesUrl = '/api/access-control/roles?delegatable=true';
@@ -15,25 +24,6 @@ export const fetchRoleOptions = async (orgId?: number): Promise<Role[]> => {
   return roles.map(addDisplayNameForFixedRole).map(addFilteredDisplayName);
 };
 
-export const fetchUserRoles = async (userId: number, orgId?: number): Promise<Role[]> => {
-  let userRolesUrl = `/api/access-control/users/${userId}/roles?includeMapped=true&includeHidden=true`;
-  if (orgId) {
-    userRolesUrl += `&targetOrgId=${orgId}`;
-  }
-  try {
-    const roles = await getBackendSrv().get(userRolesUrl);
-    if (!roles || !roles.length) {
-      return [];
-    }
-    return roles.map(addDisplayNameForFixedRole).map(addFilteredDisplayName);
-  } catch (error) {
-    if (isFetchError(error)) {
-      error.isHandled = true;
-    }
-    return [];
-  }
-};
-
 export const updateUserRoles = (roles: Role[], userId: number, orgId?: number) => {
   let userRolesUrl = `/api/access-control/users/${userId}/roles`;
   if (orgId) {
@@ -46,35 +36,3 @@ export const updateUserRoles = (roles: Role[], userId: number, orgId?: number) =
     roleUids,
   });
 };
-
-export const fetchTeamRoles = async (teamId: number, orgId?: number): Promise<Role[]> => {
-  let teamRolesUrl = `/api/access-control/teams/${teamId}/roles`;
-  if (orgId) {
-    teamRolesUrl += `?targetOrgId=${orgId}`;
-  }
-  try {
-    const roles = await getBackendSrv().get(teamRolesUrl);
-    if (!roles || !roles.length) {
-      return [];
-    }
-    return roles.map(addDisplayNameForFixedRole).map(addFilteredDisplayName);
-  } catch (error) {
-    if (isFetchError(error)) {
-      error.isHandled = true;
-    }
-    return [];
-  }
-};
-
-export const updateTeamRoles = (roles: Role[], teamId: number, orgId?: number) => {
-  let teamRolesUrl = `/api/access-control/teams/${teamId}/roles`;
-  if (orgId) {
-    teamRolesUrl += `?targetOrgId=${orgId}`;
-  }
-  const roleUids = roles.flatMap((x) => x.uid);
-
-  return getBackendSrv().put(teamRolesUrl, {
-    orgId,
-    roleUids,
-  });
-};

public/app/features/teams/hooks.ts

@@ -3,19 +3,19 @@ import { useEffect, useMemo } from 'react';
 
 import { useCreateFolder } from 'app/api/clients/folder/v1beta1/hooks';
 import {
-  useSearchTeamsQuery as useLegacySearchTeamsQuery,
+  CreateTeamCommand,
+  UpdateTeamCommand,
   useCreateTeamMutation,
   useDeleteTeamByIdMutation,
-  useListTeamsRolesQuery,
-  CreateTeamCommand,
-  useSetTeamRolesMutation,
   useGetTeamByIdQuery,
+  useSearchTeamsQuery as useLegacySearchTeamsQuery,
+  useListTeamsRolesQuery,
+  useSetTeamRolesMutation,
   useUpdateTeamMutation,
-  UpdateTeamCommand,
 } from 'app/api/clients/legacy';
-import { addFilteredDisplayName } from 'app/core/components/RolePicker/utils';
 import { updateNavIndex } from 'app/core/reducers/navModel';
 import { contextSrv } from 'app/core/services/context_srv';
+import { addFilteredDisplayName } from 'app/core/utils/roles';
 import { AccessControlAction, Role } from 'app/types/accessControl';
 import { useDispatch } from 'app/types/store';