AzureMonitor: ResourceMonitor children refactor (#34591)

* AzureMonitor: Refactor children to be an arrayu

* Add some resourcePickerData tests

* update tests

* some quick NestedResourceTable tests

public/app/plugins/datasource/grafana-azure-monitor-datasource/__mocks__/argResourcePickerResponse.ts

@@ -0,0 +1,80 @@
+import { AzureGraphResponse, RawAzureResourceGroupItem, RawAzureResourceItem } from '../types';
+
+export const createMockARGResourceContainersResponse = (): AzureGraphResponse<RawAzureResourceGroupItem[]> => ({
+  data: [
+    {
+      subscriptionURI: '/subscriptions/abc-123',
+      subscriptionName: 'Primary Subscription',
+      resourceGroupURI: '/subscriptions/abc-123/resourceGroups/prod',
+      resourceGroupName: 'Production',
+    },
+
+    {
+      subscriptionURI: '/subscription/def-456',
+      subscriptionName: 'Dev Subscription',
+      resourceGroupURI: '/subscription/def-456/resourceGroups/dev',
+      resourceGroupName: 'Development',
+    },
+
+    {
+      subscriptionURI: '/subscription/def-456',
+      subscriptionName: 'Dev Subscription',
+      resourceGroupURI: '/subscription/def-456/resourceGroups/test',
+      resourceGroupName: 'Test',
+    },
+
+    {
+      subscriptionURI: '/subscriptions/abc-123',
+      subscriptionName: 'Primary Subscription',
+      resourceGroupURI: '/subscriptions/abc-123/resourceGroups/pre-prod',
+      resourceGroupName: 'Pre-production',
+    },
+
+    {
+      subscriptionURI: '/subscription/def-456',
+      subscriptionName: 'Dev Subscription',
+      resourceGroupURI: '/subscription/def-456/resourceGroups/qa',
+      resourceGroupName: 'QA',
+    },
+  ],
+});
+
+export const createARGResourcesResponse = (): AzureGraphResponse<RawAzureResourceItem[]> => ({
+  data: [
+    {
+      id: '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/virtualMachines/web-server',
+      name: 'web-server',
+      type: 'Microsoft.Compute/virtualMachines',
+      resourceGroup: 'dev',
+      subscriptionId: 'def-456',
+      location: 'northeurope',
+    },
+
+    {
+      id: '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/disks/web-server_DataDisk',
+      name: 'web-server_DataDisk',
+      type: 'Microsoft.Compute/disks',
+      resourceGroup: 'dev',
+      subscriptionId: 'def-456',
+      location: 'northeurope',
+    },
+
+    {
+      id: '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/virtualMachines/db-server',
+      name: 'db-server',
+      type: 'Microsoft.Compute/virtualMachines',
+      resourceGroup: 'dev',
+      subscriptionId: 'def-456',
+      location: 'northeurope',
+    },
+
+    {
+      id: '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/disks/db-server_DataDisk',
+      name: 'db-server_DataDisk',
+      type: 'Microsoft.Compute/disks',
+      resourceGroup: 'dev',
+      subscriptionId: 'def-456',
+      location: 'northeurope',
+    },
+  ],
+});

public/app/plugins/datasource/grafana-azure-monitor-datasource/__mocks__/instanceSettings.ts

@@ -0,0 +1,24 @@
+import { DataSourcePluginMeta } from '@grafana/data';
+import { AzureDataSourceInstanceSettings } from '../types';
+
+export const createMockInstanceSetttings = (): AzureDataSourceInstanceSettings => ({
+  url: '/ds/1',
+  id: 1,
+  uid: 'abc',
+  type: 'azuremonitor',
+  meta: {} as DataSourcePluginMeta,
+  name: 'azure',
+
+  jsonData: {
+    cloudName: 'azuremonitor',
+    azureAuthType: 'clientsecret',
+
+    // monitor
+    tenantId: 'abc-123',
+    clientId: 'def-456',
+    subscriptionId: 'ghi-789',
+
+    // logs
+    azureLogAnalyticsSameAs: true,
+  },
+});

public/app/plugins/datasource/grafana-azure-monitor-datasource/__mocks__/resourcePickerRows.ts

@@ -0,0 +1,88 @@
+import { ResourceRowGroup, ResourceRowType } from '../components/ResourcePicker/types';
+
+export const createMockResourcePickerRows = (): ResourceRowGroup => [
+  {
+    id: '/subscriptions/abc-123',
+    name: 'Primary Subscription',
+    type: ResourceRowType.Subscription,
+    typeLabel: 'Subscription',
+    children: [
+      {
+        id: '/subscriptions/abc-123/resourceGroups/prod',
+        name: 'Production',
+        type: ResourceRowType.ResourceGroup,
+        typeLabel: 'Resource Group',
+        children: [],
+      },
+      {
+        id: '/subscriptions/abc-123/resourceGroups/pre-prod',
+        name: 'Pre-production',
+        type: ResourceRowType.ResourceGroup,
+        typeLabel: 'Resource Group',
+        children: [],
+      },
+    ],
+  },
+
+  {
+    id: '/subscriptions/def-456',
+    name: 'Dev Subscription',
+    type: ResourceRowType.Subscription,
+    typeLabel: 'Subscription',
+    children: [
+      {
+        id: '/subscriptions/def-456/resourceGroups/dev',
+        name: 'Development',
+        type: ResourceRowType.ResourceGroup,
+        typeLabel: 'Resource Group',
+        children: [
+          {
+            id: '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/virtualMachines/web-server',
+            name: 'web-server',
+            typeLabel: 'Microsoft.Compute/virtualMachines',
+            type: ResourceRowType.Resource,
+            location: 'northeurope',
+          },
+
+          {
+            id: '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/disks/web-server_DataDisk',
+            name: 'web-server_DataDisk',
+            typeLabel: 'Microsoft.Compute/disks',
+            type: ResourceRowType.Resource,
+            location: 'northeurope',
+          },
+
+          {
+            id: '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/virtualMachines/db-server',
+            name: 'db-server',
+            typeLabel: 'Microsoft.Compute/virtualMachines',
+            type: ResourceRowType.Resource,
+            location: 'northeurope',
+          },
+
+          {
+            id: '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/disks/db-server_DataDisk',
+            name: 'db-server_DataDisk',
+            typeLabel: 'Microsoft.Compute/disks',
+            type: ResourceRowType.Resource,
+            location: 'northeurope',
+          },
+        ],
+      },
+      {
+        id: '/subscriptions/def-456/resourceGroups/test',
+        name: 'Test',
+        type: ResourceRowType.ResourceGroup,
+        typeLabel: 'Resource Group',
+        children: [],
+      },
+      {
+        id: '/subscriptions/def-456/resourceGroups/qa',
+        name: 'QA',
+        type: ResourceRowType.ResourceGroup,
+        typeLabel: 'Resource Group',
+        children: [],
+      },
+    ],
+  },
+];

public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/NestedResourceTable.test.tsx

@@ -0,0 +1,67 @@
+import { act, render, screen } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+import React from 'react';
+import { createMockResourcePickerRows } from '../../__mocks__/resourcePickerRows';
+
+import NestedResourceTable from './NestedResourceTable';
+import { findRow } from './utils';
+
+describe('AzureMonitor NestedResourceTable', () => {
+  const noop: any = () => {};
+
+  it('renders subscriptions', () => {
+    const rows = createMockResourcePickerRows();
+
+    render(<NestedResourceTable rows={rows} selectedRows={[]} requestNestedRows={noop} onRowSelectedChange={noop} />);
+
+    expect(screen.getByText('Primary Subscription')).toBeInTheDocument();
+    expect(screen.getByText('Dev Subscription')).toBeInTheDocument();
+  });
+
+  it('opens to the selected resource', () => {
+    const rows = createMockResourcePickerRows();
+    const selected = findRow(
+      rows,
+      '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/disks/web-server_DataDisk'
+    );
+
+    if (!selected) {
+      throw new Error("couldn't find row, test data stale");
+    }
+
+    render(
+      <NestedResourceTable rows={rows} selectedRows={[selected]} requestNestedRows={noop} onRowSelectedChange={noop} />
+    );
+
+    expect(screen.getByText('web-server_DataDisk')).toBeInTheDocument();
+  });
+
+  it("expands resource groups when they're clicked", async () => {
+    const rows = createMockResourcePickerRows();
+    const promise = Promise.resolve();
+    const requestNestedRows = jest.fn().mockReturnValue(promise);
+    render(
+      <NestedResourceTable
+        rows={rows}
+        selectedRows={[]}
+        requestNestedRows={requestNestedRows}
+        onRowSelectedChange={noop}
+      />
+    );
+
+    const expandButton = screen.getAllByLabelText('Expand')[2];
+    userEvent.click(expandButton);
+
+    expect(requestNestedRows).toBeCalledWith(
+      expect.objectContaining({
+        id: '/subscriptions/def-456/resourceGroups/dev',
+        name: 'Development',
+        typeLabel: 'Resource Group',
+      })
+    );
+
+    await act(() => promise);
+
+    expect(screen.getByText('web-server')).toBeInTheDocument();
+  });
+});

public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/NestedResourceTable.tsx

@@ -5,14 +5,14 @@ import { useStyles2 } from '@grafana/ui';
 
 import NestedRows from './NestedRows';
 import getStyles from './styles';
-import { Row, RowGroup } from './types';
+import { ResourceRow, ResourceRowGroup } from './types';
 
 interface NestedResourceTableProps {
-  rows: RowGroup;
-  selectedRows: RowGroup;
+  rows: ResourceRowGroup;
+  selectedRows: ResourceRowGroup;
   noHeader?: boolean;
-  requestNestedRows: (row: Row) => Promise<void>;
-  onRowSelectedChange: (row: Row, selected: boolean) => void;
+  requestNestedRows: (row: ResourceRow) => Promise<void>;
+  onRowSelectedChange: (row: ResourceRow, selected: boolean) => void;
 }
 
 const NestedResourceTable: React.FC<NestedResourceTableProps> = ({

public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/NestedRows.tsx

@@ -2,14 +2,15 @@ import { cx } from '@emotion/css';
 import { Checkbox, Icon, IconButton, LoadingPlaceholder, useStyles2, useTheme2, FadeTransition } from '@grafana/ui';
 import React, { useCallback, useEffect, useState } from 'react';
 import getStyles from './styles';
-import { EntryType, Row, RowGroup } from './types';
+import { ResourceRowType, ResourceRow, ResourceRowGroup } from './types';
+import { findRow } from './utils';
 
 interface NestedRowsProps {
-  rows: RowGroup;
+  rows: ResourceRowGroup;
   level: number;
-  selectedRows: RowGroup;
-  requestNestedRows: (row: Row) => Promise<void>;
-  onRowSelectedChange: (row: Row, selected: boolean) => void;
+  selectedRows: ResourceRowGroup;
+  requestNestedRows: (row: ResourceRow) => Promise<void>;
+  onRowSelectedChange: (row: ResourceRow, selected: boolean) => void;
 }
 
 const NestedRows: React.FC<NestedRowsProps> = ({
@@ -20,10 +21,10 @@ const NestedRows: React.FC<NestedRowsProps> = ({
   onRowSelectedChange,
 }) => (
   <>
-    {Object.keys(rows).map((rowId) => (
+    {rows.map((row) => (
       <NestedRow
-        key={rowId}
-        row={rows[rowId]}
+        key={row.id}
+        row={row}
         selectedRows={selectedRows}
         level={level}
         requestNestedRows={requestNestedRows}
@@ -34,39 +35,41 @@ const NestedRows: React.FC<NestedRowsProps> = ({
 );
 
 interface NestedRowProps {
-  row: Row;
+  row: ResourceRow;
   level: number;
-  selectedRows: RowGroup;
-  requestNestedRows: (row: Row) => Promise<void>;
-  onRowSelectedChange: (row: Row, selected: boolean) => void;
+  selectedRows: ResourceRowGroup;
+  requestNestedRows: (row: ResourceRow) => Promise<void>;
+  onRowSelectedChange: (row: ResourceRow, selected: boolean) => void;
 }
 
 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 isSelected = !!selectedRows[row.id];
-  const isDisabled = Object.keys(selectedRows).length > 0 && !isSelected;
-  const initialOpenStatus = row.type === EntryType.Collection ? 'open' : 'closed';
-  const [openStatus, setOpenStatus] = useState<'open' | 'closed' | 'loading'>(initialOpenStatus);
-  const isOpen = openStatus === 'open';
+  const isSelected = !!selectedRows.find((v) => v.id === row.id);
+  const isDisabled = selectedRows.length > 0 && !isSelected;
+  const isOpen = rowStatus === 'open';
 
   const onRowToggleCollapse = async () => {
-    if (openStatus === 'open') {
-      setOpenStatus('closed');
+    if (rowStatus === 'open') {
+      setRowStatus('closed');
       return;
     }
-    setOpenStatus('loading');
+    setRowStatus('loading');
     await requestNestedRows(row);
-    setOpenStatus('open');
+    setRowStatus('open');
   };
 
   // opens the resource group on load of component if there was a previously saved selection
   useEffect(() => {
-    const selectedRow = Object.keys(selectedRows).map((rowId) => selectedRows[rowId])[0];
-    const isSelectedResourceGroup =
-      selectedRow && selectedRow.resourceGroupName && row.name === selectedRow.resourceGroupName;
-    if (isSelectedResourceGroup) {
-      setOpenStatus('open');
+    // Assuming we don't have multi-select yet
+    const selectedRow = selectedRows[0];
+
+    const containsChild = selectedRow && !!findRow(row.children ?? [], selectedRow.id);
+
+    if (containsChild) {
+      setRowStatus('open');
     }
   }, [selectedRows, row]);
 
@@ -100,7 +103,7 @@ const NestedRow: React.FC<NestedRowProps> = ({ row, selectedRows, level, request
         />
       )}
 
-      <FadeTransition visible={openStatus === 'loading'}>
+      <FadeTransition visible={rowStatus === 'loading'}>
         <tr>
           <td className={cx(styles.cell, styles.loadingCell)} colSpan={3}>
             <LoadingPlaceholder text="Loading..." className={styles.spinner} />
@@ -112,19 +115,19 @@ const NestedRow: React.FC<NestedRowProps> = ({ row, selectedRows, level, request
 };
 
 interface EntryIconProps {
-  entry: Row;
+  entry: ResourceRow;
   isOpen: boolean;
 }
 
 const EntryIcon: React.FC<EntryIconProps> = ({ isOpen, entry: { type } }) => {
   switch (type) {
-    case EntryType.Collection:
+    case ResourceRowType.Subscription:
       return <Icon name="layer-group" />;
 
-    case EntryType.SubCollection:
+    case ResourceRowType.ResourceGroup:
       return <Icon name={isOpen ? 'folder-open' : 'folder'} />;
 
-    case EntryType.Resource:
+    case ResourceRowType.Resource:
       return <Icon name="cube" />;
 
     default:
@@ -134,12 +137,12 @@ const EntryIcon: React.FC<EntryIconProps> = ({ isOpen, entry: { type } }) => {
 
 interface NestedEntryProps {
   level: number;
-  entry: Row;
+  entry: ResourceRow;
   isSelected: boolean;
   isOpen: boolean;
   isDisabled: boolean;
-  onToggleCollapse: (row: Row) => void;
-  onSelectedChange: (row: Row, selected: boolean) => void;
+  onToggleCollapse: (row: ResourceRow) => void;
+  onSelectedChange: (row: ResourceRow, selected: boolean) => void;
 }
 
 const NestedEntry: React.FC<NestedEntryProps> = ({
@@ -154,7 +157,7 @@ const NestedEntry: React.FC<NestedEntryProps> = ({
   const theme = useTheme2();
   const styles = useStyles2(getStyles);
   const hasChildren = !!entry.children;
-  const isSelectable = entry.type === EntryType.Resource;
+  const isSelectable = entry.type === ResourceRowType.Resource;
 
   const handleToggleCollapse = useCallback(() => {
     onToggleCollapse(entry);
@@ -168,10 +171,13 @@ const NestedEntry: React.FC<NestedEntryProps> = ({
     [entry, onSelectedChange]
   );
 
+  const checkboxId = `checkbox_${entry.id}`;
+
   return (
     <div className={styles.nestedEntry} style={{ marginLeft: level * (3 * theme.spacing.gridSize) }}>
       {/* When groups are selectable, I *think* we will want to show a 2-wide space instead
             of the collapse button for leaf rows that have no children to get them to align */}
+
       <span className={styles.entryContentItem}>
         {hasChildren && (
           <IconButton
@@ -179,17 +185,22 @@ const NestedEntry: React.FC<NestedEntryProps> = ({
             name={isOpen ? 'angle-down' : 'angle-right'}
             aria-label={isOpen ? 'Collapse' : 'Expand'}
             onClick={handleToggleCollapse}
+            id={entry.id}
           />
         )}
 
-        {isSelectable && <Checkbox onChange={handleSelectedChanged} disabled={isDisabled} value={isSelected} />}
+        {isSelectable && (
+          <Checkbox id={checkboxId} onChange={handleSelectedChanged} disabled={isDisabled} value={isSelected} />
+        )}
       </span>
 
       <span className={styles.entryContentItem}>
         <EntryIcon entry={entry} isOpen={isOpen} />
       </span>
 
-      <span className={cx(styles.entryContentItem, styles.truncated)}>{entry.name}</span>
+      <label htmlFor={checkboxId} className={cx(styles.entryContentItem, styles.truncated)}>
+        {entry.name}
+      </label>
     </div>
   );
 };

public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/index.tsx

@@ -1,13 +1,13 @@
 import React, { useCallback, useEffect, useMemo, useState } from 'react';
 import NestedResourceTable from './NestedResourceTable';
-import { Row, RowGroup } from './types';
+import { ResourceRow, ResourceRowGroup } from './types';
 import { css } from '@emotion/css';
 import { GrafanaTheme2 } from '@grafana/data';
 import { Button, useStyles2 } from '@grafana/ui';
 import ResourcePickerData from '../../resourcePicker/resourcePickerData';
 import { produce } from 'immer';
 import { Space } from '../Space';
-import { parseResourceURI } from './utils';
+import { findRow, parseResourceURI } from './utils';
 
 interface ResourcePickerProps {
   resourcePickerData: Pick<ResourcePickerData, 'getResourcePickerData' | 'getResourcesForResourceGroup'>;
@@ -20,72 +20,80 @@ interface ResourcePickerProps {
 const ResourcePicker = ({ resourcePickerData, resourceURI, onApply, onCancel }: ResourcePickerProps) => {
   const styles = useStyles2(getStyles);
 
-  const [rows, setRows] = useState<RowGroup>({});
+  const [rows, setRows] = useState<ResourceRowGroup>([]);
   const [internalSelected, setInternalSelected] = useState<string | undefined>(resourceURI);
 
+  // Sync the resourceURI prop to internal state
   useEffect(() => {
     setInternalSelected(resourceURI);
   }, [resourceURI]);
 
-  const handleFetchInitialResources = useCallback(async () => {
-    const initalRows = await resourcePickerData.getResourcePickerData();
-    setRows(initalRows);
-  }, [resourcePickerData]);
-
-  useEffect(() => {
-    handleFetchInitialResources();
-  }, [handleFetchInitialResources]);
+  // Map the selected item into an array of rows
+  const selectedResourceRows = useMemo(() => {
+    const found = internalSelected && findRow(rows, internalSelected);
+    return found ? [found] : [];
+  }, [internalSelected, rows]);
 
+  // Request resources for a expanded resource group
   const requestNestedRows = useCallback(
-    async (resourceGroup: Row) => {
-      // if we've already fetched resources for a resource group we don't need to re-fetch them
-      if (resourceGroup.children && Object.keys(resourceGroup.children).length > 0) {
+    async (resourceGroup: ResourceRow) => {
+      // If we already have children, we don't need to re-fetch them
+      if (resourceGroup.children?.length) {
         return;
       }
 
       // fetch and set nested resources for the resourcegroup into the bigger state object
       const resources = await resourcePickerData.getResourcesForResourceGroup(resourceGroup);
-      setRows(
-        produce(rows, (draftState: RowGroup) => {
-          const subscriptionChildren = draftState[resourceGroup.subscriptionId].children;
+      const newRows = produce(rows, (draftState) => {
+        // We want to find the row again from the draftState so we can "mutate" it
+        const draftRow = findRow(draftState, resourceGroup.id);
 
-          if (subscriptionChildren) {
-            subscriptionChildren[resourceGroup.name].children = resources;
-          }
-        })
-      );
+        if (!draftRow) {
+          // This case shouldn't happen
+          throw new Error('Unable to find resource');
+        }
+
+        draftRow.children = resources;
+      });
+
+      setRows(newRows);
     },
     [resourcePickerData, rows]
   );
 
-  const handleSelectionChanged = useCallback((row: Row, isSelected: boolean) => {
+  // Select
+  const handleSelectionChanged = useCallback((row: ResourceRow, isSelected: boolean) => {
     isSelected ? setInternalSelected(row.id) : setInternalSelected(undefined);
   }, []);
 
-  const selectedResource = useMemo(() => {
-    if (internalSelected && Object.keys(rows).length) {
-      const parsed = parseResourceURI(internalSelected);
+  // Request initial data on first mount
+  useEffect(() => {
+    resourcePickerData.getResourcePickerData().then((initalRows) => setRows(initalRows));
+  }, [resourcePickerData]);
 
-      if (parsed) {
-        const { subscriptionID, resourceGroup } = parsed;
-        const allResourceGroups = rows[subscriptionID].children || {};
-        const selectedResourceGroup = allResourceGroups[resourceGroup.toLowerCase()];
-        const allResourcesInResourceGroup = selectedResourceGroup.children;
-
-        if (!allResourcesInResourceGroup || Object.keys(allResourcesInResourceGroup).length === 0) {
-          requestNestedRows(selectedResourceGroup);
-          return {};
-        }
-
-        const matchingResource = allResourcesInResourceGroup[internalSelected];
-
-        return {
-          [internalSelected]: matchingResource,
-        };
-      }
+  // Request sibling resources for a selected resource - in practice should only be on first mount
+  useEffect(() => {
+    if (!internalSelected || !rows.length) {
+      return;
     }
-    return {};
-  }, [internalSelected, rows, requestNestedRows]);
+
+    // 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);
+
+    // TODO: Error UI
+    if (!resourceGroupRow) {
+      throw new Error('Unable to find resource group for resource');
+    }
+
+    requestNestedRows(resourceGroupRow);
+  }, [requestNestedRows, internalSelected, rows]);
 
   const handleApply = useCallback(() => {
     onApply(internalSelected);
@@ -97,20 +105,20 @@ const ResourcePicker = ({ resourcePickerData, resourceURI, onApply, onCancel }:
         rows={rows}
         requestNestedRows={requestNestedRows}
         onRowSelectedChange={handleSelectionChanged}
-        selectedRows={selectedResource}
+        selectedRows={selectedResourceRows}
       />
 
       <div className={styles.selectionFooter}>
-        {internalSelected && (
+        {selectedResourceRows.length > 0 && (
           <>
             <Space v={2} />
             <h5>Selection</h5>
             <NestedResourceTable
-              noHeader={true}
-              rows={selectedResource}
+              rows={selectedResourceRows}
               requestNestedRows={requestNestedRows}
               onRowSelectedChange={handleSelectionChanged}
-              selectedRows={selectedResource}
+              selectedRows={selectedResourceRows}
+              noHeader={true}
             />
           </>
         )}

public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/types.ts

@@ -1,19 +1,16 @@
-export enum EntryType {
-  Collection,
-  SubCollection,
-  Resource,
-}
-export interface Row {
-  id: string;
-  name: string;
-  type: EntryType;
-  typeLabel: string;
-  subscriptionId: string;
-  location?: string;
-  children?: RowGroup;
-  resourceGroupName?: string;
+export enum ResourceRowType {
+  Subscription = 'Subscription',
+  ResourceGroup = 'ResourceGroup',
+  Resource = 'Resource',
 }
 
-export interface RowGroup {
-  [subscriptionIdOrResourceGroupName: string]: Row;
+export interface ResourceRow {
+  id: string;
+  name: string;
+  type: ResourceRowType;
+  typeLabel: string;
+  location?: string;
+  children?: ResourceRowGroup;
 }
+
+export type ResourceRowGroup = ResourceRow[];

public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ResourcePicker/utils.ts

@@ -1,3 +1,5 @@
+import { ResourceRow, ResourceRowGroup } from './types';
+
 const RESOURCE_URI_REGEX = /\/subscriptions\/(?<subscriptionID>.+)\/resourceGroups\/(?<resourceGroup>.+)\/providers.+\/(?<resource>[\w-_]+)/;
 
 export function parseResourceURI(resourceURI: string) {
@@ -14,3 +16,21 @@ export function parseResourceURI(resourceURI: string) {
 export function isGUIDish(input: string) {
   return !!input.match(/^[A-Z0-9]+/i);
 }
+
+export function findRow(rows: ResourceRowGroup, id: string): ResourceRow | undefined {
+  for (const row of rows) {
+    if (row.id === id) {
+      return row;
+    }
+
+    if (row.children) {
+      const result = findRow(row.children, id);
+
+      if (result) {
+        return result;
+      }
+    }
+  }
+
+  return undefined;
+}

public/app/plugins/datasource/grafana-azure-monitor-datasource/resourcePicker/resourcePickerData.test.ts

@@ -0,0 +1,114 @@
+import { of } from 'rxjs';
+
+import { createFetchResponse } from 'test/helpers/createFetchResponse';
+import { backendSrv } from 'app/core/services/backend_srv';
+
+import ResourcePickerData from './resourcePickerData';
+import {
+  createMockARGResourceContainersResponse,
+  createARGResourcesResponse,
+} from '../__mocks__/argResourcePickerResponse';
+import { ResourceRowType } from '../components/ResourcePicker/types';
+import { createMockInstanceSetttings } from '../__mocks__/instanceSettings';
+
+jest.mock('@grafana/runtime', () => ({
+  ...((jest.requireActual('@grafana/runtime') as unknown) as object),
+  getBackendSrv: () => backendSrv,
+}));
+
+const instanceSettings = createMockInstanceSetttings();
+
+describe('AzureMonitor resourcePickerData', () => {
+  describe('getResourcePickerData', () => {
+    let fetchMock: jest.SpyInstance;
+
+    beforeEach(() => {
+      fetchMock = jest.spyOn(backendSrv, 'fetch');
+      fetchMock.mockImplementation(() => {
+        const data = createMockARGResourceContainersResponse();
+        return of(createFetchResponse(data));
+      });
+    });
+
+    afterEach(() => fetchMock.mockReset());
+
+    it('calls ARG API', async () => {
+      const resourcePickerData = new ResourcePickerData(instanceSettings);
+      await resourcePickerData.getResourcePickerData();
+
+      expect(fetchMock).toHaveBeenCalled();
+      const argQuery = fetchMock.mock.calls[0][0].data.query;
+
+      expect(argQuery).toContain(`where type == 'microsoft.resources/subscriptions'`);
+      expect(argQuery).toContain(`where type == 'microsoft.resources/subscriptions/resourcegroups'`);
+    });
+
+    it('returns only subscriptions at the top level', async () => {
+      const resourcePickerData = new ResourcePickerData(instanceSettings);
+      const results = await resourcePickerData.getResourcePickerData();
+
+      expect(results.map((v) => v.id)).toEqual(['/subscriptions/abc-123', '/subscription/def-456']);
+    });
+
+    it('nests resource groups under their subscriptions', async () => {
+      const resourcePickerData = new ResourcePickerData(instanceSettings);
+      const results = await resourcePickerData.getResourcePickerData();
+
+      expect(results[0].children?.map((v) => v.id)).toEqual([
+        '/subscriptions/abc-123/resourceGroups/prod',
+        '/subscriptions/abc-123/resourceGroups/pre-prod',
+      ]);
+
+      expect(results[1].children?.map((v) => v.id)).toEqual([
+        '/subscription/def-456/resourceGroups/dev',
+        '/subscription/def-456/resourceGroups/test',
+        '/subscription/def-456/resourceGroups/qa',
+      ]);
+    });
+  });
+
+  describe('getResourcesForResourceGroup', () => {
+    let fetchMock: jest.SpyInstance;
+
+    const resourceRow = {
+      id: '/subscription/def-456/resourceGroups/dev',
+      name: 'Dev',
+      type: ResourceRowType.ResourceGroup,
+      typeLabel: 'Resource group',
+    };
+
+    beforeEach(() => {
+      fetchMock = jest.spyOn(backendSrv, 'fetch');
+      fetchMock.mockImplementation(() => {
+        const data = createARGResourcesResponse();
+        return of(createFetchResponse(data));
+      });
+    });
+
+    afterEach(() => fetchMock.mockReset());
+
+    it('requests resources for the specified resource row', async () => {
+      const resourcePickerData = new ResourcePickerData(instanceSettings);
+      await resourcePickerData.getResourcesForResourceGroup(resourceRow);
+
+      expect(fetchMock).toHaveBeenCalled();
+      const argQuery = fetchMock.mock.calls[0][0].data.query;
+
+      expect(argQuery).toContain(resourceRow.id);
+    });
+
+    it('returns formatted resources', async () => {
+      const resourcePickerData = new ResourcePickerData(instanceSettings);
+      const results = await resourcePickerData.getResourcesForResourceGroup(resourceRow);
+
+      expect(results.map((v) => v.id)).toEqual([
+        '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/virtualMachines/web-server',
+        '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/disks/web-server_DataDisk',
+        '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/virtualMachines/db-server',
+        '/subscription/def-456/resourceGroups/dev/providers/Microsoft.Compute/disks/db-server_DataDisk',
+      ]);
+
+      results.forEach((v) => expect(v.type).toEqual(ResourceRowType.Resource));
+    });
+  });
+});

public/app/plugins/datasource/grafana-azure-monitor-datasource/resourcePicker/resourcePickerData.ts

@@ -1,32 +1,18 @@
 import { FetchResponse, getBackendSrv } from '@grafana/runtime';
 import { getLogAnalyticsResourcePickerApiRoute } from '../api/routes';
-import { EntryType, Row, RowGroup } from '../components/ResourcePicker/types';
+import { ResourceRowType, ResourceRow, ResourceRowGroup } from '../components/ResourcePicker/types';
 import { getAzureCloud } from '../credentials';
-import { AzureDataSourceInstanceSettings, AzureResourceSummaryItem } from '../types';
+import {
+  AzureDataSourceInstanceSettings,
+  AzureGraphResponse,
+  AzureResourceSummaryItem,
+  RawAzureResourceGroupItem,
+  RawAzureResourceItem,
+} from '../types';
 import { SUPPORTED_LOCATIONS, SUPPORTED_RESOURCE_TYPES } from './supportedResources';
 
 const RESOURCE_GRAPH_URL = '/providers/Microsoft.ResourceGraph/resources?api-version=2020-04-01-preview';
 
-interface RawAzureResourceGroupItem {
-  subscriptionId: string;
-  subscriptionName: string;
-  resourceGroup: string;
-  resourceGroupId: string;
-}
-
-interface RawAzureResourceItem {
-  id: string;
-  name: string;
-  subscriptionId: string;
-  resourceGroup: string;
-  type: string;
-  location: string;
-}
-
-interface AzureGraphResponse<T = unknown> {
-  data: T;
-}
-
 export default class ResourcePickerData {
   private proxyUrl: string;
   private cloud: string;
@@ -37,27 +23,43 @@ export default class ResourcePickerData {
   }
 
   async getResourcePickerData() {
-    const { ok, data: response } = await this.makeResourceGraphRequest<RawAzureResourceGroupItem[]>(
-      `resources
-      | join kind=leftouter (ResourceContainers | where type=='microsoft.resources/subscriptions' | project subscriptionName=name, subscriptionId, resourceGroupId=id) on subscriptionId
-      | where type in (${SUPPORTED_RESOURCE_TYPES})
-      | summarize count() by resourceGroup, subscriptionName, resourceGroupId, subscriptionId
-      | order by resourceGroup asc
-      `
-    );
+    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
+
+        // Put resource group details on each row
+        | join kind=leftouter (
+          ResourceContainers
+            | where type == 'microsoft.resources/subscriptions/resourcegroups'
+            | project resourceGroupURI=id, resourceGroupName=name, resourceGroup
+          ) on resourceGroup
+
+        | where type in (${SUPPORTED_RESOURCE_TYPES})
+
+        // Get only unique resource groups and subscriptions. Also acts like a project
+        | summarize count() by resourceGroupName, resourceGroupURI, subscriptionName, subscriptionURI
+        | order by subscriptionURI asc
+    `;
+
+    const { ok, data: response } = await this.makeResourceGraphRequest<RawAzureResourceGroupItem[]>(query);
 
     // TODO: figure out desired error handling strategy
     if (!ok) {
       throw new Error('unable to fetch resource containers');
     }
 
-    return this.formatResourceGroupData(response.data);
+    return formatResourceGroupData(response.data);
   }
 
-  async getResourcesForResourceGroup(resourceGroup: Row) {
+  async getResourcesForResourceGroup(resourceGroup: ResourceRow) {
     const { ok, data: response } = await this.makeResourceGraphRequest<RawAzureResourceItem[]>(`
       resources
-      | where resourceGroup == "${resourceGroup.name.toLowerCase()}"
+      | where id hasprefix "${resourceGroup.id}"
       | where type in (${SUPPORTED_RESOURCE_TYPES}) and location in (${SUPPORTED_LOCATIONS})
     `);
 
@@ -66,7 +68,7 @@ export default class ResourcePickerData {
       throw new Error('unable to fetch resource containers');
     }
 
-    return this.formatResourceGroupChildren(response.data);
+    return formatResourceGroupChildren(response.data);
   }
 
   async getResource(resourceURI: string) {
@@ -114,57 +116,6 @@ export default class ResourcePickerData {
     return response.data[0].id;
   }
 
-  formatResourceGroupData(rawData: RawAzureResourceGroupItem[]) {
-    const formatedSubscriptionsAndResourceGroups: RowGroup = {};
-
-    rawData.forEach((resourceGroup) => {
-      // if the subscription doesn't exist yet, create it
-      if (!formatedSubscriptionsAndResourceGroups[resourceGroup.subscriptionId]) {
-        formatedSubscriptionsAndResourceGroups[resourceGroup.subscriptionId] = {
-          name: resourceGroup.subscriptionName,
-          id: resourceGroup.subscriptionId,
-          subscriptionId: resourceGroup.subscriptionId,
-          typeLabel: 'Subscription',
-          type: EntryType.Collection,
-          children: {},
-        };
-      }
-
-      // add the resource group to the subscription
-      // store by resourcegroupname not id to match resource uri
-      (formatedSubscriptionsAndResourceGroups[resourceGroup.subscriptionId].children as RowGroup)[
-        resourceGroup.resourceGroup
-      ] = {
-        name: resourceGroup.resourceGroup,
-        id: resourceGroup.resourceGroupId,
-        subscriptionId: resourceGroup.subscriptionId,
-        type: EntryType.SubCollection,
-        typeLabel: 'Resource Group',
-        children: {},
-      };
-    });
-
-    return formatedSubscriptionsAndResourceGroups;
-  }
-
-  formatResourceGroupChildren(rawData: RawAzureResourceItem[]) {
-    const children: RowGroup = {};
-
-    rawData.forEach((item: RawAzureResourceItem) => {
-      children[item.id] = {
-        name: item.name,
-        id: item.id,
-        subscriptionId: item.id,
-        resourceGroupName: item.resourceGroup,
-        type: EntryType.Resource,
-        typeLabel: item.type, // TODO: these types can be quite long, we may wish to format them more
-        location: item.location, // TODO: we may wish to format these locations, by default they are written as 'northeurope' rather than a more human readable "North Europe"
-      };
-    });
-
-    return children;
-  }
-
   async makeResourceGraphRequest<T = unknown>(
     query: string,
     maxRetries = 1
@@ -191,3 +142,52 @@ export default class ResourcePickerData {
     }
   }
 }
+
+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,
+    id: item.id,
+    resourceGroupName: item.resourceGroup,
+    type: ResourceRowType.Resource,
+    typeLabel: item.type, // TODO: these types can be quite long, we may wish to format them more
+    location: item.location, // TODO: we may wish to format these locations, by default they are written as 'northeurope' rather than a more human readable "North Europe"
+  }));
+}

public/app/plugins/datasource/grafana-azure-monitor-datasource/types/index.ts

@@ -230,3 +230,24 @@ export interface AzureResourceSummaryItem {
   subscriptionName: string;
   resourceGroupName: string;
 }
+
+export interface RawAzureResourceGroupItem {
+  subscriptionURI: string;
+  subscriptionName: string;
+
+  resourceGroupURI: string;
+  resourceGroupName: string;
+}
+
+export interface RawAzureResourceItem {
+  id: string;
+  name: string;
+  subscriptionId: string;
+  resourceGroup: string;
+  type: string;
+  location: string;
+}
+
+export interface AzureGraphResponse<T = unknown> {
+  data: T;
+}