Azure monitor Logs: Optimize data fetching in resource picker (#44549)

* refactor resource picker * add tests * remove not used code * remove unnecessary check * fix broken tests * update table tests now that we no longer expand second level by default * add missing new line * pr feedback * fix broken test
2025-08-02 21:52:43 +08:00 · 2022-02-01 09:39:48 +01:00
parent c31585c246
commit dc0a2fb55b
12 changed files with 642 additions and 329 deletions
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/mocks/argResourcePickerResponse.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/mocks/argResourcePickerResponse.ts
@ -1,45 +1,67 @@
-import { AzureGraphResponse, RawAzureResourceGroupItem, RawAzureResourceItem } from '../types';
+import {
+  AzureGraphResponse,
+  RawAzureResourceGroupItem,
+  RawAzureResourceItem,
+  RawAzureSubscriptionItem,
+} from '../types';

-export const createMockARGResourceContainersResponse = (): AzureGraphResponse<RawAzureResourceGroupItem[]> => ({
+export const createMockARGSubscriptionResponse = (): AzureGraphResponse<RawAzureSubscriptionItem[]> => ({
  data: [
    {
-      subscriptionURI: '/subscriptions/abc-123',
+      subscriptionId: '1',
      subscriptionName: 'Primary Subscription',
+    },
+    {
+      subscriptionId: '2',
+      subscriptionName: 'Dev Subscription',
+    },
+    {
+      subscriptionId: '3',
+      subscriptionName: 'Dev Subscription',
+    },
+    {
+      subscriptionId: '4',
+      subscriptionName: 'Primary Subscription',
+    },
+    {
+      subscriptionId: '5',
+      subscriptionName: 'Primary Subscription',
+    },
+    {
+      subscriptionId: '6',
+      subscriptionName: 'Dev Subscription',
+    },
+  ],
+});
+
+export const createMockARGResourceGroupsResponse = (): AzureGraphResponse<RawAzureResourceGroupItem[]> => ({
+  data: [
+    {
      resourceGroupURI: '/subscriptions/abc-123/resourceGroups/prod',
      resourceGroupName: 'Production',
    },

    {
-      subscriptionURI: '/subscriptions/def-456',
-      subscriptionName: 'Dev Subscription',
      resourceGroupURI: '/subscriptions/def-456/resourceGroups/dev',
      resourceGroupName: 'Development',
    },

    {
-      subscriptionURI: '/subscriptions/def-456',
-      subscriptionName: 'Dev Subscription',
      resourceGroupURI: '/subscriptions/def-456/resourceGroups/test',
      resourceGroupName: 'Test',
    },

    {
-      subscriptionURI: '/subscriptions/abc-123',
-      subscriptionName: 'Primary Subscription',
      resourceGroupURI: '/subscriptions/abc-123/resourceGroups/test',
      resourceGroupName: 'Test',
    },

    {
-      subscriptionURI: '/subscriptions/abc-123',
-      subscriptionName: 'Primary Subscription',
      resourceGroupURI: '/subscriptions/abc-123/resourceGroups/pre-prod',
      resourceGroupName: 'Pre-production',
    },

    {
-      subscriptionURI: '/subscriptions/def-456',
-      subscriptionName: 'Dev Subscription',
      resourceGroupURI: '/subscriptions/def-456/resourceGroups/qa',
      resourceGroupName: 'QA',
    },
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/mocks/datasource.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/mocks/datasource.ts
@ -37,9 +37,11 @@ export default function createMockDatasource(overrides?: DeepPartial<Datasource>
      getDeprecatedDefaultWorkSpace: () => 'defaultWorkspaceId',
    },
    resourcePickerData: {
-      getResourcePickerData: () => ({}),
-      getResourcesForResourceGroup: () => ({}),
-      getResourceURIFromWorkspace: () => '',
+      getSubscriptions: () => jest.fn().mockResolvedValue([]),
+      getResourceGroupsBySubscriptionId: jest.fn().mockResolvedValue([]),
+      getResourcesForResourceGroup: jest.fn().mockResolvedValue([]),
+      getResourceURIFromWorkspace: jest.fn().mockReturnValue(''),
+      transformVariablesToRow: jest.fn().mockReturnValue({}),
    },
    ...overrides,
  };
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/mocks/resourcePickerData.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/mocks/resourcePickerData.ts
@ -0,0 +1,20 @@
+import ResourcePicker from '../resourcePicker/resourcePickerData';
+
+type DeepPartial<T> = {
+  [P in keyof T]?: DeepPartial<T[P]>;
+};
+
+export default function createMockResourcePickerData(overrides?: DeepPartial<ResourcePicker>) {
+  const _mockResourcePicker: DeepPartial<ResourcePicker> = {
+    getSubscriptions: () => jest.fn().mockResolvedValue([]),
+    getResourceGroupsBySubscriptionId: jest.fn().mockResolvedValue([]),
+    getResourcesForResourceGroup: jest.fn().mockResolvedValue([]),
+    getResourceURIFromWorkspace: jest.fn().mockReturnValue(''),
+    transformVariablesToRow: jest.fn().mockReturnValue({}),
+    ...overrides,
+  };
+
+  const mockDatasource = _mockResourcePicker as ResourcePicker;
+
+  return jest.mocked(mockDatasource, true);
+}
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/mocks/resourcePickerRows.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/mocks/resourcePickerRows.ts
@ -43,7 +43,6 @@ export const createMockResourcePickerRows = (): ResourceRowGroup => [
            type: ResourceRowType.Resource,
            location: 'northeurope',
          },
-
          {
            id: '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/disks/web-server_DataDisk',
            name: 'web-server_DataDisk',
@ -107,3 +106,98 @@ export const createMockResourcePickerRows = (): ResourceRowGroup => [
    ],
  },
 ];
+
+export const createMockSubscriptions = (): ResourceRowGroup => [
+  {
+    id: 'def-123',
+    name: 'Primary Subscription',
+    type: ResourceRowType.Subscription,
+    typeLabel: 'Subscription',
+    children: [],
+  },
+  {
+    id: 'def-456',
+    name: 'Dev Subscription',
+    type: ResourceRowType.Subscription,
+    typeLabel: 'Subscription',
+    children: [],
+  },
+  {
+    id: 'def-789',
+    name: 'Test Subscription',
+    type: ResourceRowType.Subscription,
+    typeLabel: 'Subscription',
+    children: [],
+  },
+];
+
+export const createMockResourceGroupsBySubscription = (): ResourceRowGroup => [
+  {
+    id: '/subscriptions/def-456/resourceGroups/dev-1',
+    name: 'Development',
+    type: ResourceRowType.ResourceGroup,
+    typeLabel: 'Resource Group',
+    children: [],
+  },
+  {
+    id: '/subscriptions/def-456/resourceGroups/dev-2',
+    name: 'Development',
+    type: ResourceRowType.ResourceGroup,
+    typeLabel: 'Resource Group',
+    children: [],
+  },
+  {
+    id: '/subscriptions/def-456/resourceGroups/dev-3',
+    name: 'Development',
+    type: ResourceRowType.ResourceGroup,
+    typeLabel: 'Resource Group',
+    children: [],
+  },
+  {
+    id: '/subscriptions/def-456/resourceGroups/dev-4',
+    name: 'Development',
+    type: ResourceRowType.ResourceGroup,
+    typeLabel: 'Resource Group',
+    children: [],
+  },
+  {
+    id: '/subscriptions/def-456/resourceGroups/dev-5',
+    name: 'Development',
+    type: ResourceRowType.ResourceGroup,
+    typeLabel: 'Resource Group',
+    children: [],
+  },
+];
+
+export const mockResourcesByResourceGroup = (): ResourceRowGroup => [
+  {
+    id: 'Microsoft.Compute/virtualMachines/web-server',
+    name: 'web-server',
+    typeLabel: 'Microsoft.Compute/virtualMachines',
+    type: ResourceRowType.Resource,
+    location: 'northeurope',
+  },
+  {
+    id: 'Microsoft.Compute/disks/web-server_DataDisk',
+    name: 'web-server_DataDisk',
+    typeLabel: 'Microsoft.Compute/disks',
+    type: ResourceRowType.Resource,
+    location: 'northeurope',
+  },
+
+  {
+    id: 'Microsoft.Compute/virtualMachines/db-server',
+    name: 'db-server',
+    typeLabel: 'Microsoft.Compute/virtualMachines',
+    type: ResourceRowType.Resource,
+    location: 'northeurope',
+  },
+
+  {
+    id: 'Microsoft.Compute/disks/db-server_DataDisk',
+    name: 'db-server_DataDisk',
+    typeLabel: 'Microsoft.Compute/disks',
+    type: ResourceRowType.Resource,
+    location: 'northeurope',
+  },
+];
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/NestedResourceTable.test.tsx
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/NestedResourceTable.test.tsx
@ -46,7 +46,7 @@ describe('AzureMonitor NestedResourceTable', () => {
    expect(screen.getByText('web-server_DataDisk')).toBeInTheDocument();
  });

-  it("expands resource groups when they're clicked", async () => {
+  it("expands subscriptions when they're clicked", async () => {
    const rows = createMockResourcePickerRows();
    const promise = Promise.resolve();
    const requestNestedRows = jest.fn().mockReturnValue(promise);
@ -59,20 +59,20 @@ describe('AzureMonitor NestedResourceTable', () => {
      />
    );

-    const expandButton = screen.getAllByLabelText('Expand')[2];
+    const expandButton = screen.getAllByLabelText('Expand')[1];
    userEvent.click(expandButton);

    expect(requestNestedRows).toBeCalledWith(
      expect.objectContaining({
-        id: '/subscriptions/def-456/resourceGroups/dev',
-        name: 'Development',
-        typeLabel: 'Resource Group',
+        id: '/subscriptions/def-456',
+        name: 'Dev Subscription',
+        typeLabel: 'Subscription',
      })
    );

+    expect(screen.queryByText('Development')).not.toBeInTheDocument();
    await act(() => promise);
-
-    expect(screen.getByText('web-server')).toBeInTheDocument();
+    expect(screen.getByText('Development')).toBeInTheDocument();
  });

  it('supports selecting variables', async () => {
@ -89,7 +89,7 @@ describe('AzureMonitor NestedResourceTable', () => {
      />
    );

-    const expandButton = screen.getAllByLabelText('Expand')[5];
+    const expandButton = screen.getAllByLabelText('Expand')[2];
    userEvent.click(expandButton);

    await act(() => promise);
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/NestedRows.tsx
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/NestedRows.tsx
@ -46,8 +46,7 @@ interface NestedRowProps {

 const NestedRow: React.FC<NestedRowProps> = ({ row, selectedRows, level, requestNestedRows, onRowSelectedChange }) => {
  const styles = useStyles2(getStyles);
-  const initialOpenStatus = row.type === ResourceRowType.Subscription ? 'open' : 'closed';
-  const [rowStatus, setRowStatus] = useState<'open' | 'closed' | 'loading'>(initialOpenStatus);
+  const [rowStatus, setRowStatus] = useState<'open' | 'closed' | 'loading'>('closed');

  const isSelected = !!selectedRows.find((v) => v.id === row.id);
  const isDisabled = selectedRows.length > 0 && !isSelected;
@ -59,8 +58,9 @@ const NestedRow: React.FC<NestedRowProps> = ({ row, selectedRows, level, request
      return;
    }
    setRowStatus('loading');
-    await requestNestedRows(row);
-    setRowStatus('open');
+    requestNestedRows(row)
+      .then(() => setRowStatus('open'))
+      .catch(() => setRowStatus('closed'));
  };

  // opens the resource group on load of component if there was a previously saved selection
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/ResourcePicker.test.tsx
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/ResourcePicker.test.tsx
@ -0,0 +1,136 @@
+import { act, render, screen, waitFor } from '@testing-library/react';
+import React from 'react';
+
+import ResourcePicker from '.';
+import createMockResourcePickerData from '../../__mocks__/resourcePickerData';
+import {
+  createMockResourceGroupsBySubscription,
+  createMockSubscriptions,
+  mockResourcesByResourceGroup,
+} from '../../__mocks__/resourcePickerRows';
+
+const noResourceURI = '';
+const singleSubscriptionSelectionURI = 'def-456';
+const singleResourceGroupSelectionURI = '/subscriptions/def-456/resourceGroups/dev-3';
+const singleResourceSelectionURI =
+  '/subscriptions/def-456/resourceGroups/dev-3/providers/Microsoft.Compute/virtualMachines/db-serverproviders/Microsoft.Compute/virtualMachines/db-server';
+
+describe('AzureMonitor ResourcePicker', () => {
+  const noop: any = () => {};
+
+  beforeEach(() => {
+    window.HTMLElement.prototype.scrollIntoView = function () {};
+  });
+  describe('when rendering the resource picker without a selection', () => {
+    it('should load subscriptions', async () => {
+      const resourePickerDataMock = createMockResourcePickerData({
+        getSubscriptions: jest.fn().mockResolvedValue(createMockSubscriptions()),
+        getResourceGroupsBySubscriptionId: jest.fn(),
+        getResourcesForResourceGroup: jest.fn(),
+      });
+      await act(() => {
+        render(
+          <ResourcePicker
+            templateVariables={[]}
+            resourcePickerData={resourePickerDataMock}
+            resourceURI={noResourceURI}
+            onCancel={noop}
+            onApply={noop}
+          />
+        );
+      });
+
+      await waitFor(() => expect(screen.getByText('Primary Subscription')).toBeInTheDocument());
+      expect(resourePickerDataMock.getSubscriptions).toHaveBeenCalledTimes(1);
+      expect(resourePickerDataMock.getResourceGroupsBySubscriptionId).not.toHaveBeenCalled();
+      expect(resourePickerDataMock.getResourcesForResourceGroup).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('when rendering the resource picker with a subscription selected', () => {
+    it('should load subscriptions once', async () => {
+      const resourePickerDataMock = createMockResourcePickerData({
+        getSubscriptions: jest.fn().mockResolvedValue(createMockSubscriptions()),
+        getResourceGroupsBySubscriptionId: jest.fn(),
+        getResourcesForResourceGroup: jest.fn(),
+      });
+      await act(async () => {
+        render(
+          <ResourcePicker
+            templateVariables={[]}
+            resourcePickerData={resourePickerDataMock}
+            resourceURI={singleSubscriptionSelectionURI}
+            onCancel={noop}
+            onApply={noop}
+          />
+        );
+      });
+
+      await waitFor(() => expect(screen.getByText('Primary Subscription')).toBeInTheDocument());
+      expect(resourePickerDataMock.getSubscriptions).toHaveBeenCalledTimes(1);
+      expect(resourePickerDataMock.getResourceGroupsBySubscriptionId).not.toHaveBeenCalled();
+      expect(resourePickerDataMock.getResourcesForResourceGroup).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('when rendering the resource picker with a resource group selected', () => {
+    it('should load subscriptions and resource groups for its parent subscription once', async () => {
+      const resourePickerDataMock = createMockResourcePickerData({
+        getSubscriptions: jest.fn().mockResolvedValue(createMockSubscriptions()),
+        getResourceGroupsBySubscriptionId: jest.fn().mockResolvedValue(createMockResourceGroupsBySubscription()),
+        getResourcesForResourceGroup: jest.fn(),
+      });
+      await act(async () => {
+        render(
+          <ResourcePicker
+            templateVariables={[]}
+            resourcePickerData={resourePickerDataMock}
+            resourceURI={singleResourceGroupSelectionURI}
+            onCancel={noop}
+            onApply={noop}
+          />
+        );
+      });
+
+      await waitFor(() => expect(screen.getByText('Primary Subscription')).toBeInTheDocument());
+      expect(resourePickerDataMock.getSubscriptions).toHaveBeenCalledTimes(1);
+      expect(resourePickerDataMock.getResourceGroupsBySubscriptionId).toHaveBeenCalledTimes(1);
+      expect(resourePickerDataMock.getResourceGroupsBySubscriptionId).toHaveBeenLastCalledWith(
+        singleSubscriptionSelectionURI
+      );
+      expect(resourePickerDataMock.getResourcesForResourceGroup).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('when rendering the resource picker with a resource selected', () => {
+    it('should load subscriptions, resource groups and resources once', async () => {
+      const resourePickerDataMock = createMockResourcePickerData({
+        getSubscriptions: jest.fn().mockResolvedValue(createMockSubscriptions()),
+        getResourceGroupsBySubscriptionId: jest.fn().mockResolvedValue(createMockResourceGroupsBySubscription()),
+        getResourcesForResourceGroup: jest.fn().mockResolvedValue(mockResourcesByResourceGroup()),
+      });
+      await act(async () => {
+        render(
+          <ResourcePicker
+            templateVariables={[]}
+            resourcePickerData={resourePickerDataMock}
+            resourceURI={singleResourceSelectionURI}
+            onCancel={noop}
+            onApply={noop}
+          />
+        );
+      });
+
+      await waitFor(() => expect(screen.getByText('Primary Subscription')).toBeInTheDocument());
+      expect(resourePickerDataMock.getSubscriptions).toHaveBeenCalledTimes(1);
+      expect(resourePickerDataMock.getResourceGroupsBySubscriptionId).toHaveBeenCalledTimes(1);
+      expect(resourePickerDataMock.getResourceGroupsBySubscriptionId).toHaveBeenLastCalledWith(
+        singleSubscriptionSelectionURI
+      );
+      expect(resourePickerDataMock.getResourcesForResourceGroup).toHaveBeenCalledTimes(1);
+      expect(resourePickerDataMock.getResourcesForResourceGroup).toHaveBeenLastCalledWith(
+        singleResourceGroupSelectionURI
+      );
+    });
+  });
+});
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/ResourcePicker.tsx
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/ResourcePicker.tsx
@ -0,0 +1,213 @@
+import { css } from '@emotion/css';
+import { GrafanaTheme2 } from '@grafana/data';
+import { Alert, Button, LoadingPlaceholder, useStyles2 } from '@grafana/ui';
+import React, { useCallback, useEffect, useMemo, useState } from 'react';
+
+import ResourcePickerData from '../../resourcePicker/resourcePickerData';
+import messageFromError from '../../utils/messageFromError';
+import { Space } from '../Space';
+import NestedResourceTable from './NestedResourceTable';
+import { ResourceRow, ResourceRowGroup, ResourceRowType } from './types';
+import { addResources, findRow, parseResourceURI } from './utils';
+
+interface ResourcePickerProps {
+  resourcePickerData: ResourcePickerData;
+  resourceURI: string | undefined;
+  templateVariables: string[];
+
+  onApply: (resourceURI: string | undefined) => void;
+  onCancel: () => void;
+}
+
+const ResourcePicker = ({
+  resourcePickerData,
+  resourceURI,
+  templateVariables,
+  onApply,
+  onCancel,
+}: ResourcePickerProps) => {
+  const styles = useStyles2(getStyles);
+
+  type LoadingStatus = 'NotStarted' | 'Started' | 'Done';
+  const [loadingStatus, setLoadingStatus] = useState<LoadingStatus>('NotStarted');
+  const [azureRows, setAzureRows] = useState<ResourceRowGroup>([]);
+  const [internalSelected, setInternalSelected] = useState<string | undefined>(resourceURI);
+  const [errorMessage, setErrorMessage] = useState<string | undefined>(undefined);
+
+  // Sync the resourceURI prop to internal state
+  useEffect(() => {
+    setInternalSelected(resourceURI);
+  }, [resourceURI]);
+
+  // Request initial data on first mount
+  useEffect(() => {
+    if (loadingStatus === 'NotStarted') {
+      const loadInitialData = async () => {
+        try {
+          setLoadingStatus('Started');
+          let resources = await resourcePickerData.getSubscriptions();
+          if (!internalSelected) {
+            setAzureRows(resources);
+            setLoadingStatus('Done');
+            return;
+          }
+
+          const parsedURI = parseResourceURI(internalSelected ?? '');
+          if (parsedURI) {
+            const resourceGroupURI = `/subscriptions/${parsedURI.subscriptionID}/resourceGroups/${parsedURI.resourceGroup}`;
+
+            // if a resource group was previously selected, but the resource groups under the parent subscription have not been loaded yet
+            if (parsedURI.resourceGroup && !findRow(resources, resourceGroupURI)) {
+              const resourceGroups = await resourcePickerData.getResourceGroupsBySubscriptionId(
+                parsedURI.subscriptionID
+              );
+              resources = addResources(resources, parsedURI.subscriptionID, resourceGroups);
+            }
+
+            // if a resource was previously selected, but the resources under the parent resource group have not been loaded yet
+            if (parsedURI.resource && !findRow(azureRows, parsedURI.resource ?? '')) {
+              const resourcesForResourceGroup = await resourcePickerData.getResourcesForResourceGroup(resourceGroupURI);
+              resources = addResources(resources, resourceGroupURI, resourcesForResourceGroup);
+            }
+          }
+          setAzureRows(resources);
+          setLoadingStatus('Done');
+        } catch (error) {
+          setLoadingStatus('Done');
+          setErrorMessage(messageFromError(error));
+        }
+      };
+
+      loadInitialData();
+    }
+  }, [resourcePickerData, internalSelected, azureRows, loadingStatus]);
+
+  const rows = useMemo(() => {
+    const templateVariableRow = resourcePickerData.transformVariablesToRow(templateVariables);
+    return templateVariables.length ? [...azureRows, templateVariableRow] : azureRows;
+  }, [resourcePickerData, azureRows, templateVariables]);
+
+  // Map the selected item into an array of rows
+  const selectedResourceRows = useMemo(() => {
+    const found = internalSelected && findRow(rows, internalSelected);
+    return found
+      ? [
+          {
+            ...found,
+            children: undefined,
+          },
+        ]
+      : [];
+  }, [internalSelected, rows]);
+
+  // Request resources for a expanded resource group
+  const requestNestedRows = useCallback(
+    async (resourceGroupOrSubscription: ResourceRow) => {
+      // clear error message (also when loading cached resources)
+      setErrorMessage(undefined);
+
+      // If we already have children, we don't need to re-fetch them. Also abort if we're expanding the special
+      // template variable group, though that shouldn't happen in practice
+      if (
+        resourceGroupOrSubscription.children?.length ||
+        resourceGroupOrSubscription.id === ResourcePickerData.templateVariableGroupID
+      ) {
+        return;
+      }
+
+      try {
+        const rows =
+          resourceGroupOrSubscription.type === ResourceRowType.Subscription
+            ? await resourcePickerData.getResourceGroupsBySubscriptionId(resourceGroupOrSubscription.id)
+            : await resourcePickerData.getResourcesForResourceGroup(resourceGroupOrSubscription.id);
+
+        const newRows = addResources(azureRows, resourceGroupOrSubscription.id, rows);
+
+        setAzureRows(newRows);
+      } catch (error) {
+        setErrorMessage(messageFromError(error));
+        throw error;
+      }
+    },
+    [resourcePickerData, azureRows]
+  );
+
+  // Select
+  const handleSelectionChanged = useCallback((row: ResourceRow, isSelected: boolean) => {
+    isSelected ? setInternalSelected(row.id) : setInternalSelected(undefined);
+  }, []);
+
+  const handleApply = useCallback(() => {
+    onApply(internalSelected);
+  }, [internalSelected, onApply]);
+
+  return (
+    <div>
+      {loadingStatus === 'Started' ? (
+        <div className={styles.loadingWrapper}>
+          <LoadingPlaceholder text={'Loading...'} />
+        </div>
+      ) : (
+        <>
+          <NestedResourceTable
+            rows={rows}
+            requestNestedRows={requestNestedRows}
+            onRowSelectedChange={handleSelectionChanged}
+            selectedRows={selectedResourceRows}
+          />
+
+          <div className={styles.selectionFooter}>
+            {selectedResourceRows.length > 0 && (
+              <>
+                <Space v={2} />
+                <h5>Selection</h5>
+                <NestedResourceTable
+                  rows={selectedResourceRows}
+                  requestNestedRows={requestNestedRows}
+                  onRowSelectedChange={handleSelectionChanged}
+                  selectedRows={selectedResourceRows}
+                  noHeader={true}
+                />
+              </>
+            )}
+
+            <Space v={2} />
+
+            <Button disabled={!!errorMessage} onClick={handleApply}>
+              Apply
+            </Button>
+            <Space layout="inline" h={1} />
+            <Button onClick={onCancel} variant="secondary">
+              Cancel
+            </Button>
+          </div>
+        </>
+      )}
+      {errorMessage && (
+        <>
+          <Space v={2} />
+          <Alert severity="error" title="An error occurred while requesting resources from Azure Monitor">
+            {errorMessage}
+          </Alert>
+        </>
+      )}
+    </div>
+  );
+};
+
+export default ResourcePicker;
+
+const getStyles = (theme: GrafanaTheme2) => ({
+  selectionFooter: css({
+    position: 'sticky',
+    bottom: 0,
+    background: theme.colors.background.primary,
+    paddingTop: theme.spacing(2),
+  }),
+  loadingWrapper: css({
+    textAlign: 'center',
+    paddingTop: theme.spacing(2),
+    paddingBottom: theme.spacing(2),
+    color: theme.colors.text.secondary,
+  }),
+});
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/index.tsx
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/index.tsx
@ -1,174 +1 @@
-import React, { useCallback, useEffect, useMemo, useState } from 'react';
-import NestedResourceTable from './NestedResourceTable';
-import { ResourceRow, ResourceRowGroup } from './types';
-import { css } from '@emotion/css';
-import { GrafanaTheme2 } from '@grafana/data';
-import { Button, LoadingPlaceholder, useStyles2 } from '@grafana/ui';
-import ResourcePickerData from '../../resourcePicker/resourcePickerData';
-import { Space } from '../Space';
-import { addResources, findRow, parseResourceURI } from './utils';
-
-interface ResourcePickerProps {
-  resourcePickerData: ResourcePickerData;
-  resourceURI: string | undefined;
-  templateVariables: string[];
-
-  onApply: (resourceURI: string | undefined) => void;
-  onCancel: () => void;
-}
-
-const ResourcePicker = ({
-  resourcePickerData,
-  resourceURI,
-  templateVariables,
-  onApply,
-  onCancel,
-}: ResourcePickerProps) => {
-  const styles = useStyles2(getStyles);
-
-  const [azureRows, setAzureRows] = useState<ResourceRowGroup>([]);
-  const [internalSelected, setInternalSelected] = useState<string | undefined>(resourceURI);
-  const [isLoading, setIsLoading] = useState(false);
-
-  // Sync the resourceURI prop to internal state
-  useEffect(() => {
-    setInternalSelected(resourceURI);
-  }, [resourceURI]);
-
-  const rows = useMemo(() => {
-    const templateVariableRow = resourcePickerData.transformVariablesToRow(templateVariables);
-    return templateVariables.length ? [...azureRows, templateVariableRow] : azureRows;
-  }, [resourcePickerData, azureRows, templateVariables]);
-
-  // Map the selected item into an array of rows
-  const selectedResourceRows = useMemo(() => {
-    const found = internalSelected && findRow(rows, internalSelected);
-    return found
-      ? [
-          {
-            ...found,
-            children: undefined,
-          },
-        ]
-      : [];
-  }, [internalSelected, rows]);
-
-  // Request resources for a expanded resource group
-  const requestNestedRows = useCallback(
-    async (resourceGroup: ResourceRow) => {
-      // If we already have children, we don't need to re-fetch them. Also abort if we're expanding the special
-      // template variable group, though that shouldn't happen in practice
-      if (resourceGroup.children?.length || resourceGroup.id === ResourcePickerData.templateVariableGroupID) {
-        return;
-      }
-
-      // fetch and set nested resources for the resourcegroup into the bigger state object
-      const resources = await resourcePickerData.getResourcesForResourceGroup(resourceGroup);
-      const newRows = addResources(azureRows, resourceGroup.id, resources);
-      setAzureRows(newRows);
-    },
-    [resourcePickerData, azureRows]
-  );
-
-  // Select
-  const handleSelectionChanged = useCallback((row: ResourceRow, isSelected: boolean) => {
-    isSelected ? setInternalSelected(row.id) : setInternalSelected(undefined);
-  }, []);
-
-  // Request initial data on first mount
-  useEffect(() => {
-    setIsLoading(true);
-    resourcePickerData.getResourcePickerData().then((initalRows) => {
-      setIsLoading(false);
-      setAzureRows(initalRows);
-    });
-  }, [resourcePickerData]);
-
-  // Request sibling resources for a selected resource - in practice should only be on first mount
-  useEffect(() => {
-    if (!internalSelected || !rows.length) {
-      return;
-    }
-
-    // If we can find this resource in the rows, then we don't need to load anything
-    const foundResourceRow = findRow(rows, internalSelected);
-    if (foundResourceRow) {
-      return;
-    }
-
-    const parsedURI = parseResourceURI(internalSelected);
-    const resourceGroupURI = `/subscriptions/${parsedURI?.subscriptionID}/resourceGroups/${parsedURI?.resourceGroup}`;
-    const resourceGroupRow = findRow(rows, resourceGroupURI);
-
-    if (!resourceGroupRow) {
-      // We haven't loaded the data from Azure yet
-      return;
-    }
-
-    requestNestedRows(resourceGroupRow);
-  }, [requestNestedRows, internalSelected, rows]);
-
-  const handleApply = useCallback(() => {
-    onApply(internalSelected);
-  }, [internalSelected, onApply]);
-
-  return (
-    <div>
-      {isLoading ? (
-        <div className={styles.loadingWrapper}>
-          <LoadingPlaceholder text={'Loading resources...'} />
-        </div>
-      ) : (
-        <>
-          <NestedResourceTable
-            rows={rows}
-            requestNestedRows={requestNestedRows}
-            onRowSelectedChange={handleSelectionChanged}
-            selectedRows={selectedResourceRows}
-          />
-
-          <div className={styles.selectionFooter}>
-            {selectedResourceRows.length > 0 && (
-              <>
-                <Space v={2} />
-                <h5>Selection</h5>
-                <NestedResourceTable
-                  rows={selectedResourceRows}
-                  requestNestedRows={requestNestedRows}
-                  onRowSelectedChange={handleSelectionChanged}
-                  selectedRows={selectedResourceRows}
-                  noHeader={true}
-                />
-              </>
-            )}
-
-            <Space v={2} />
-
-            <Button onClick={handleApply}>Apply</Button>
-            <Space layout="inline" h={1} />
-            <Button onClick={onCancel} variant="secondary">
-              Cancel
-            </Button>
-          </div>
-        </>
-      )}
-    </div>
-  );
-};
-
-export default ResourcePicker;
-
-const getStyles = (theme: GrafanaTheme2) => ({
-  selectionFooter: css({
-    position: 'sticky',
-    bottom: 0,
-    background: theme.colors.background.primary,
-    paddingTop: theme.spacing(2),
-  }),
-  loadingWrapper: css({
-    textAlign: 'center',
-    paddingTop: theme.spacing(2),
-    paddingBottom: theme.spacing(2),
-    color: theme.colors.text.secondary,
-  }),
-});
+export { default } from './ResourcePicker';
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/resourcePicker/resourcePickerData.test.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/resourcePicker/resourcePickerData.test.ts
@ -1,62 +1,39 @@
-import ResourcePickerData from './resourcePickerData';
 import {
-  createMockARGResourceContainersResponse,
  createARGResourcesResponse,
+  createMockARGResourceGroupsResponse,
+  createMockARGSubscriptionResponse,
 } from '../__mocks__/argResourcePickerResponse';
-import { ResourceRowType } from '../components/ResourcePicker/types';
 import { createMockInstanceSetttings } from '../__mocks__/instanceSettings';
+import { ResourceRowType } from '../components/ResourcePicker/types';
+import ResourcePickerData from './resourcePickerData';

 const instanceSettings = createMockInstanceSetttings();
 const resourcePickerData = new ResourcePickerData(instanceSettings);
 let postResource: jest.Mock;

 describe('AzureMonitor resourcePickerData', () => {
-  describe('getResourcePickerData', () => {
+  describe('getSubscriptions', () => {
    beforeEach(() => {
-      postResource = jest.fn().mockResolvedValue(createMockARGResourceContainersResponse());
+      postResource = jest.fn().mockResolvedValue(createMockARGSubscriptionResponse());
      resourcePickerData.postResource = postResource;
    });

    it('calls ARG API', async () => {
-      await resourcePickerData.getResourcePickerData();
+      await resourcePickerData.getSubscriptions();

      expect(postResource).toHaveBeenCalled();
      const argQuery = postResource.mock.calls[0][1].query;

      expect(argQuery).toContain(`where type == 'microsoft.resources/subscriptions'`);
-      expect(argQuery).toContain(`where type == 'microsoft.resources/subscriptions/resourcegroups'`);
-      expect(argQuery).toContain(`project resourceGroupURI=id, resourceGroupName=name, resourceGroup, subscriptionId`);
-    });
-
-    it('returns only subscriptions at the top level', async () => {
-      const results = await resourcePickerData.getResourcePickerData();
-
-      expect(results.map((v) => v.id)).toEqual(['/subscriptions/abc-123', '/subscriptions/def-456']);
-    });
-
-    it('nests resource groups under their subscriptions', async () => {
-      const results = await resourcePickerData.getResourcePickerData();
-
-      expect(results[0].children?.map((v) => v.id)).toEqual([
-        '/subscriptions/abc-123/resourceGroups/prod',
-        '/subscriptions/abc-123/resourceGroups/test',
-        '/subscriptions/abc-123/resourceGroups/pre-prod',
-      ]);
-
-      expect(results[1].children?.map((v) => v.id)).toEqual([
-        '/subscriptions/def-456/resourceGroups/dev',
-        '/subscriptions/def-456/resourceGroups/test',
-        '/subscriptions/def-456/resourceGroups/qa',
-      ]);
    });

    describe('when there is more than one page', () => {
      beforeEach(() => {
        const response1 = {
-          ...createMockARGResourceContainersResponse(),
+          ...createMockARGSubscriptionResponse(),
          $skipToken: 'aaa',
        };
-        const response2 = createMockARGResourceContainersResponse();
+        const response2 = createMockARGSubscriptionResponse();
        postResource = jest.fn();
        postResource.mockResolvedValueOnce(response1);
        postResource.mockResolvedValueOnce(response2);
@ -64,37 +41,55 @@ describe('AzureMonitor resourcePickerData', () => {
      });

      it('should requests additional pages', async () => {
-        await resourcePickerData.getResourcePickerData();
+        await resourcePickerData.getSubscriptions();
        expect(postResource).toHaveBeenCalledTimes(2);
      });

      it('should use the skipToken of the previous page', async () => {
-        await resourcePickerData.getResourcePickerData();
+        await resourcePickerData.getSubscriptions();
        const secondCall = postResource.mock.calls[1];
        expect(secondCall[1]).toMatchObject({ options: { $skipToken: 'aaa', resultFormat: 'objectArray' } });
      });
+    });
+  });

-      it('should combine responses', async () => {
-        const results = await resourcePickerData.getResourcePickerData();
-        expect(results[0].children?.map((v) => v.id)).toEqual([
-          '/subscriptions/abc-123/resourceGroups/prod',
-          '/subscriptions/abc-123/resourceGroups/test',
-          '/subscriptions/abc-123/resourceGroups/pre-prod',
-          // second page
-          '/subscriptions/abc-123/resourceGroups/prod',
-          '/subscriptions/abc-123/resourceGroups/test',
-          '/subscriptions/abc-123/resourceGroups/pre-prod',
-        ]);
+  describe('getResourcesForResourceGroup', () => {
+    beforeEach(() => {
+      postResource = jest.fn().mockResolvedValue(createMockARGResourceGroupsResponse());
+      resourcePickerData.postResource = postResource;
+    });

-        expect(results[1].children?.map((v) => v.id)).toEqual([
-          '/subscriptions/def-456/resourceGroups/dev',
-          '/subscriptions/def-456/resourceGroups/test',
-          '/subscriptions/def-456/resourceGroups/qa',
-          // second page
-          '/subscriptions/def-456/resourceGroups/dev',
-          '/subscriptions/def-456/resourceGroups/test',
-          '/subscriptions/def-456/resourceGroups/qa',
-        ]);
+    it('calls ARG API', async () => {
+      await resourcePickerData.getResourceGroupsBySubscriptionId('123');
+
+      expect(postResource).toHaveBeenCalled();
+      const argQuery = postResource.mock.calls[0][1].query;
+
+      expect(argQuery).toContain(`| where subscriptionId == '123'`);
+    });
+
+    describe('when there is more than one page', () => {
+      beforeEach(() => {
+        const response1 = {
+          ...createMockARGResourceGroupsResponse(),
+          $skipToken: 'aaa',
+        };
+        const response2 = createMockARGResourceGroupsResponse();
+        postResource = jest.fn();
+        postResource.mockResolvedValueOnce(response1);
+        postResource.mockResolvedValueOnce(response2);
+        resourcePickerData.postResource = postResource;
+      });
+
+      it('should requests additional pages', async () => {
+        await resourcePickerData.getResourceGroupsBySubscriptionId('123');
+        expect(postResource).toHaveBeenCalledTimes(2);
+      });
+
+      it('should use the skipToken of the previous page', async () => {
+        await resourcePickerData.getResourceGroupsBySubscriptionId('123');
+        const secondCall = postResource.mock.calls[1];
+        expect(secondCall[1]).toMatchObject({ options: { $skipToken: 'aaa', resultFormat: 'objectArray' } });
      });
    });
  });
@ -113,7 +108,7 @@ describe('AzureMonitor resourcePickerData', () => {
    });

    it('requests resources for the specified resource row', async () => {
-      await resourcePickerData.getResourcesForResourceGroup(resourceRow);
+      await resourcePickerData.getResourcesForResourceGroup(resourceRow.id);

      expect(postResource).toHaveBeenCalled();
      const argQuery = postResource.mock.calls[0][1].query;
@ -122,7 +117,7 @@ describe('AzureMonitor resourcePickerData', () => {
    });

    it('returns formatted resources', async () => {
-      const results = await resourcePickerData.getResourcesForResourceGroup(resourceRow);
+      const results = await resourcePickerData.getResourcesForResourceGroup(resourceRow.id);

      expect(results.map((v) => v.id)).toEqual([
        '/subscriptions/def-456/resourceGroups/dev/providers/Microsoft.Compute/virtualMachines/web-server',
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/resourcePicker/resourcePickerData.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/resourcePicker/resourcePickerData.ts
@ -1,4 +1,5 @@
 import { DataSourceWithBackend } from '@grafana/runtime';
+
 import { DataSourceInstanceSettings } from '../../../../../../packages/grafana-data/src';
 import {
  locationDisplayNames,
@ -6,7 +7,7 @@ import {
  logsSupportedResourceTypesKusto,
  resourceTypeDisplayNames,
 } from '../azureMetadata';
-import { ResourceRowType, ResourceRow, ResourceRowGroup } from '../components/ResourcePicker/types';
+import { ResourceRow, ResourceRowGroup, ResourceRowType } from '../components/ResourcePicker/types';
 import { parseResourceURI } from '../components/ResourcePicker/utils';
 import {
  AzureDataSourceJsonData,
@ -16,6 +17,7 @@ import {
  AzureResourceSummaryItem,
  RawAzureResourceGroupItem,
  RawAzureResourceItem,
+  RawAzureSubscriptionItem,
 } from '../types';
 import { routeNames } from '../utils/common';

@ -31,29 +33,61 @@ export default class ResourcePickerData extends DataSourceWithBackend<AzureMonit

  static readonly templateVariableGroupID = '$$grafana-templateVariables$$';

-  async getResourcePickerData() {
+  async getSubscriptions(): Promise<ResourceRowGroup> {
    const query = `
-      resources
-        // Put subscription details on each row
-        | join kind=leftouter (
-          ResourceContainers
-            | where type == 'microsoft.resources/subscriptions'
-            | project subscriptionName=name, subscriptionURI=id, subscriptionId
-          ) on subscriptionId
+    resources
+    | join kind=inner (
+              ResourceContainers
+                | where type == 'microsoft.resources/subscriptions'
+                | project subscriptionName=name, subscriptionURI=id, subscriptionId
+              ) on subscriptionId
+    | summarize count() by subscriptionName, subscriptionURI, subscriptionId
+    | order by subscriptionURI asc
+  `;

-        // Put resource group details on each row
-        | join kind=leftouter (
-          ResourceContainers
-            | where type == 'microsoft.resources/subscriptions/resourcegroups'
-            | project resourceGroupURI=id, resourceGroupName=name, resourceGroup, subscriptionId
-          ) on resourceGroup, subscriptionId
+    let resources: RawAzureSubscriptionItem[] = [];

-        | where type in (${logsSupportedResourceTypesKusto})
+    let allFetched = false;
+    let $skipToken = undefined;
+    while (!allFetched) {
+      // The response may include several pages
+      let options: Partial<AzureResourceGraphOptions> = {};
+      if ($skipToken) {
+        options = {
+          $skipToken,
+        };
+      }
+      const resourceResponse = await this.makeResourceGraphRequest<RawAzureSubscriptionItem[]>(query, 1, options);
+      if (!resourceResponse.data.length) {
+        throw new Error('unable to fetch resource details');
+      }
+      resources = resources.concat(resourceResponse.data);
+      $skipToken = resourceResponse.$skipToken;
+      allFetched = !$skipToken;
+    }

-        // Get only unique resource groups and subscriptions. Also acts like a project
-        | summarize count() by resourceGroupName, resourceGroupURI, subscriptionName, subscriptionURI
-        | order by subscriptionURI asc
-    `;
+    return resources.map((subscription) => ({
+      name: subscription.subscriptionName,
+      id: subscription.subscriptionId,
+      typeLabel: 'Subscription',
+      type: ResourceRowType.Subscription,
+      children: [],
+    }));
+  }
+
+  async getResourceGroupsBySubscriptionId(subscriptionId: string) {
+    const query = `
+    resources
+     | join kind=inner (
+       ResourceContainers
+       | where type == 'microsoft.resources/subscriptions/resourcegroups'
+       | project resourceGroupURI=id, resourceGroupName=name, resourceGroup, subscriptionId
+     ) on resourceGroup, subscriptionId
+
+     | where type in (${logsSupportedResourceTypesKusto})
+     | where subscriptionId == '${subscriptionId}'
+     | summarize count() by resourceGroupName, resourceGroupURI
+     | order by resourceGroupURI asc`;

    let resources: RawAzureResourceGroupItem[] = [];
    let allFetched = false;
@ -75,13 +109,19 @@ export default class ResourcePickerData extends DataSourceWithBackend<AzureMonit
      allFetched = !$skipToken;
    }

-    return formatResourceGroupData(resources);
+    return resources.map((r) => ({
+      name: r.resourceGroupName,
+      id: r.resourceGroupURI,
+      type: ResourceRowType.ResourceGroup,
+      typeLabel: 'Resource Group',
+      children: [],
+    }));
  }

-  async getResourcesForResourceGroup(resourceGroup: ResourceRow) {
+  async getResourcesForResourceGroup(resourceGroupId: string) {
    const { data: response } = await this.makeResourceGraphRequest<RawAzureResourceItem[]>(`
      resources
-      | where id hasprefix "${resourceGroup.id}"
+      | where id hasprefix "${resourceGroupId}"
      | where type in (${logsSupportedResourceTypesKusto}) and location in (${logsSupportedLocationsKusto})
    `);

@ -101,13 +141,13 @@ export default class ResourcePickerData extends DataSourceWithBackend<AzureMonit
    const resourceGroupURI = `${subscriptionURI}/resourceGroups/${resourceGroup}`;

    const query = `
-      resourcecontainers
-        | where type == "microsoft.resources/subscriptions"
-        | where id =~ "${subscriptionURI}"
-        | project subscriptionName=name, subscriptionId
+    resourcecontainers
+    | where type == "microsoft.resources/subscriptions"
+    | where id =~ "${subscriptionURI}"
+    | project subscriptionName=name, subscriptionId

-        | join kind=leftouter (
-          resourcecontainers
+    | join kind=leftouter (
+      resourcecontainers            
            | where type == "microsoft.resources/subscriptions/resourcegroups"
            | where id =~ "${resourceGroupURI}"
            | project resourceGroupName=name, resourceGroup, subscriptionId
@ -183,44 +223,6 @@ export default class ResourcePickerData extends DataSourceWithBackend<AzureMonit
  }
 }

-function formatResourceGroupData(rawData: RawAzureResourceGroupItem[]) {
-  // Subscriptions goes into the top level array
-  const rows: ResourceRowGroup = [];
-
-  // Array of all the resource groups, with subscription data on each row
-  for (const row of rawData) {
-    const resourceGroupRow: ResourceRow = {
-      name: row.resourceGroupName,
-      id: row.resourceGroupURI,
-      type: ResourceRowType.ResourceGroup,
-      typeLabel: 'Resource Group',
-      children: [],
-    };
-
-    const subscription = rows.find((v) => v.id === row.subscriptionURI);
-
-    if (subscription) {
-      if (!subscription.children) {
-        subscription.children = [];
-      }
-
-      subscription.children.push(resourceGroupRow);
-    } else {
-      const newSubscriptionRow = {
-        name: row.subscriptionName,
-        id: row.subscriptionURI,
-        typeLabel: 'Subscription',
-        type: ResourceRowType.Subscription,
-        children: [resourceGroupRow],
-      };
-
-      rows.push(newSubscriptionRow);
-    }
-  }
-
-  return rows;
-}
-
 function formatResourceGroupChildren(rawData: RawAzureResourceItem[]): ResourceRowGroup {
  return rawData.map((item) => ({
    name: item.name,
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/types/types.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/types/types.ts
@ -1,6 +1,6 @@
 import { DataSourceInstanceSettings, DataSourceJsonData, DataSourceSettings, TableData } from '@grafana/data';
-import Datasource from '../datasource';

+import Datasource from '../datasource';
 import { AzureMonitorQuery } from './query';

 export type AzureDataSourceSettings = DataSourceSettings<AzureDataSourceJsonData, AzureDataSourceSecureJsonData>;
@ -159,10 +159,12 @@ export interface AzureResourceSummaryItem {
  resourceName: string | undefined;
 }

-export interface RawAzureResourceGroupItem {
-  subscriptionURI: string;
+export interface RawAzureSubscriptionItem {
  subscriptionName: string;
+  subscriptionId: string;
+}

+export interface RawAzureResourceGroupItem {
  resourceGroupURI: string;
  resourceGroupName: string;
 }