AzureMonitor: Request and concat subsequent resource pages (#36958)

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

@@ -3,7 +3,7 @@ import NestedResourceTable from './NestedResourceTable';
 import { ResourceRow, ResourceRowGroup } from './types';
 import { css } from '@emotion/css';
 import { GrafanaTheme2 } from '@grafana/data';
-import { Button, useStyles2 } from '@grafana/ui';
+import { Button, LoadingPlaceholder, useStyles2 } from '@grafana/ui';
 import ResourcePickerData from '../../resourcePicker/resourcePickerData';
 import { Space } from '../Space';
 import { addResources, findRow, parseResourceURI } from './utils';
@@ -28,6 +28,7 @@ const ResourcePicker = ({
 
   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(() => {
@@ -76,7 +77,9 @@ const ResourcePicker = ({
 
   // Request initial data on first mount
   useEffect(() => {
+    setIsLoading(true);
     resourcePickerData.getResourcePickerData().then((initalRows) => {
+      setIsLoading(false);
       setAzureRows(initalRows);
     });
   }, [resourcePickerData]);
@@ -111,36 +114,44 @@ const ResourcePicker = ({
 
   return (
     <div>
-      <NestedResourceTable
-        rows={rows}
-        requestNestedRows={requestNestedRows}
-        onRowSelectedChange={handleSelectionChanged}
-        selectedRows={selectedResourceRows}
-      />
+      {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}
+                />
+              </>
+            )}
 
-      <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>
+            <Button onClick={handleApply}>Apply</Button>
+            <Space layout="inline" h={1} />
+            <Button onClick={onCancel} variant="secondary">
+              Cancel
+            </Button>
+          </div>
+        </>
+      )}
     </div>
   );
 };
@@ -154,4 +165,10 @@ const getStyles = (theme: GrafanaTheme2) => ({
     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,
+  }),
 });

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

@@ -47,6 +47,52 @@ describe('AzureMonitor resourcePickerData', () => {
         '/subscription/def-456/resourceGroups/qa',
       ]);
     });
+
+    describe('when there is more than one page', () => {
+      beforeEach(() => {
+        const response1 = {
+          ...createMockARGResourceContainersResponse(),
+          $skipToken: 'aaa',
+        };
+        const response2 = createMockARGResourceContainersResponse();
+        postResource = jest.fn();
+        postResource.mockResolvedValueOnce(response1);
+        postResource.mockResolvedValueOnce(response2);
+        resourcePickerData.postResource = postResource;
+      });
+
+      it('should requests additional pages', async () => {
+        await resourcePickerData.getResourcePickerData();
+        expect(postResource).toHaveBeenCalledTimes(2);
+      });
+
+      it('should use the skipToken of the previous page', async () => {
+        await resourcePickerData.getResourcePickerData();
+        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/pre-prod',
+          // second page
+          '/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',
+          // second page
+          '/subscription/def-456/resourceGroups/dev',
+          '/subscription/def-456/resourceGroups/test',
+          '/subscription/def-456/resourceGroups/qa',
+        ]);
+      });
+    });
   });
 
   describe('getResourcesForResourceGroup', () => {

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

@@ -12,6 +12,7 @@ import {
   AzureDataSourceJsonData,
   AzureGraphResponse,
   AzureMonitorQuery,
+  AzureResourceGraphOptions,
   AzureResourceSummaryItem,
   RawAzureResourceGroupItem,
   RawAzureResourceItem,
@@ -54,9 +55,27 @@ export default class ResourcePickerData extends DataSourceWithBackend<AzureMonit
         | order by subscriptionURI asc
     `;
 
-    const response = await this.makeResourceGraphRequest<RawAzureResourceGroupItem[]>(query);
+    let resources: RawAzureResourceGroupItem[] = [];
+    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<RawAzureResourceGroupItem[]>(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;
+    }
 
-    return formatResourceGroupData(response.data);
+    return formatResourceGroupData(resources);
   }
 
   async getResourcesForResourceGroup(resourceGroup: ResourceRow) {
@@ -126,12 +145,17 @@ export default class ResourcePickerData extends DataSourceWithBackend<AzureMonit
     return response[0].id;
   }
 
-  async makeResourceGraphRequest<T = unknown>(query: string, maxRetries = 1): Promise<AzureGraphResponse<T>> {
+  async makeResourceGraphRequest<T = unknown>(
+    query: string,
+    maxRetries = 1,
+    reqOptions?: Partial<AzureResourceGraphOptions>
+  ): Promise<AzureGraphResponse<T>> {
     try {
       return await this.postResource(this.resourcePath + RESOURCE_GRAPH_URL, {
         query: query,
         options: {
           resultFormat: 'objectArray',
+          ...reqOptions,
         },
       });
     } catch (error) {

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

@@ -177,4 +177,15 @@ export interface RawAzureResourceItem {
 
 export interface AzureGraphResponse<T = unknown> {
   data: T;
+  // skipToken is used for pagination, to get the next page
+  $skipToken?: string;
+}
+
+// https://docs.microsoft.com/en-us/rest/api/azureresourcegraph/resourcegraph(2021-03-01)/resources/resources#queryrequestoptions
+export interface AzureResourceGraphOptions {
+  $skip: number;
+  $skipToken: string;
+  $top: number;
+  allowPartialScopes: boolean;
+  resultFormat: 'objectArray' | 'table';
 }