AzureMonitor: Display variables in Logs resource picker (#34648)

* AzureMonitor: Display variables in Resource Picker

* tests

* hide template variable group when there's no template variables

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

@@ -85,4 +85,25 @@ export const createMockResourcePickerRows = (): ResourceRowGroup => [
       },
     ],
   },
+
+  {
+    id: '$$grafana-templateVariables$$',
+    name: 'Template variables',
+    type: ResourceRowType.VariableGroup,
+    typeLabel: 'Variables',
+    children: [
+      {
+        id: '$machine',
+        name: '$machine',
+        type: ResourceRowType.Variable,
+        typeLabel: 'Variable',
+      },
+      {
+        id: '$workspace',
+        name: '$workspace',
+        type: ResourceRowType.Variable,
+        typeLabel: 'Variable',
+      },
+    ],
+  },
 ];

public/app/plugins/datasource/grafana-azure-monitor-datasource/components/LogsQueryEditor/ResourceField.tsx

@@ -2,6 +2,7 @@ import { css } from '@emotion/css';
 import { GrafanaTheme2 } from '@grafana/data';
 import { Button, Icon, Modal, useStyles2 } from '@grafana/ui';
 import React, { useCallback, useEffect, useState } from 'react';
+import Datasource from '../../datasource';
 import { AzureQueryEditorFieldProps, AzureResourceSummaryItem } from '../../types';
 import { Field } from '../Field';
 import ResourcePicker from '../ResourcePicker';
@@ -27,17 +28,8 @@ const ResourceField: React.FC<AzureQueryEditorFieldProps> = ({ query, datasource
   const styles = useStyles2(getStyles);
   const { resource } = query.azureLogAnalytics;
 
-  const [resourceComponents, setResourceComponents] = useState(parseResourceDetails(resource ?? ''));
   const [pickerIsOpen, setPickerIsOpen] = useState(false);
 
-  useEffect(() => {
-    if (resource && parseResourceDetails(resource)) {
-      datasource.resourcePickerData.getResource(resource).then(setResourceComponents);
-    } else {
-      setResourceComponents(undefined);
-    }
-  }, [datasource.resourcePickerData, resource]);
-
   const handleOpenPicker = useCallback(() => {
     setPickerIsOpen(true);
   }, []);
@@ -60,12 +52,15 @@ const ResourceField: React.FC<AzureQueryEditorFieldProps> = ({ query, datasource
     [closePicker, onQueryChange, query]
   );
 
+  const templateVariables = datasource.getVariables();
+
   return (
     <>
       <Modal className={styles.modal} title="Select a resource" isOpen={pickerIsOpen} onDismiss={closePicker}>
         <ResourcePicker
           resourcePickerData={datasource.resourcePickerData}
           resourceURI={query.azureLogAnalytics.resource!}
+          templateVariables={templateVariables}
           onApply={handleApply}
           onCancel={closePicker}
         />
@@ -73,21 +68,53 @@ const ResourceField: React.FC<AzureQueryEditorFieldProps> = ({ query, datasource
 
       <Field label="Resource">
         <Button variant="secondary" onClick={handleOpenPicker}>
-          {/* Three mutually exclusive states */}
-          {!resource && 'Select a resource'}
-          {resource && resourceComponents && <FormattedResource resource={resourceComponents} />}
-          {resource && !resourceComponents && resource}
+          <ResourceLabel resource={resource} datasource={datasource} />
         </Button>
       </Field>
     </>
   );
 };
 
+interface ResourceLabelProps {
+  resource: string | undefined;
+  datasource: Datasource;
+}
+
+const ResourceLabel = ({ resource, datasource }: ResourceLabelProps) => {
+  const [resourceComponents, setResourceComponents] = useState(parseResourceDetails(resource ?? ''));
+
+  useEffect(() => {
+    if (resource && parseResourceDetails(resource)) {
+      datasource.resourcePickerData.getResource(resource).then(setResourceComponents);
+    } else {
+      setResourceComponents(undefined);
+    }
+  }, [datasource.resourcePickerData, resource]);
+
+  if (!resource) {
+    return <>Select a resource</>;
+  }
+
+  if (resourceComponents) {
+    return <FormattedResource resource={resourceComponents} />;
+  }
+
+  if (resource.startsWith('$')) {
+    return (
+      <span>
+        <Icon name="x" /> {resource}
+      </span>
+    );
+  }
+
+  return <>{resource}</>;
+};
+
 interface FormattedResourceProps {
   resource: AzureResourceSummaryItem;
 }
 
-const FormattedResource: React.FC<FormattedResourceProps> = ({ resource }) => {
+const FormattedResource = ({ resource }: FormattedResourceProps) => {
   return (
     <span>
       <Icon name="layer-group" /> {resource.subscriptionName}

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

@@ -64,4 +64,35 @@ describe('AzureMonitor NestedResourceTable', () => {
 
     expect(screen.getByText('web-server')).toBeInTheDocument();
   });
+
+  it('supports selecting variables', async () => {
+    const rows = createMockResourcePickerRows();
+    const promise = Promise.resolve();
+    const requestNestedRows = jest.fn().mockReturnValue(promise);
+    const onRowSelectedChange = jest.fn();
+    render(
+      <NestedResourceTable
+        rows={rows}
+        selectedRows={[]}
+        requestNestedRows={requestNestedRows}
+        onRowSelectedChange={onRowSelectedChange}
+      />
+    );
+
+    const expandButton = screen.getAllByLabelText('Expand')[5];
+    userEvent.click(expandButton);
+
+    await act(() => promise);
+
+    const checkbox = screen.getByLabelText('$workspace');
+    userEvent.click(checkbox);
+
+    expect(onRowSelectedChange).toHaveBeenCalledWith(
+      expect.objectContaining({
+        id: '$workspace',
+        name: '$workspace',
+      }),
+      true
+    );
+  });
 });

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

@@ -130,6 +130,12 @@ const EntryIcon: React.FC<EntryIconProps> = ({ isOpen, entry: { type } }) => {
     case ResourceRowType.Resource:
       return <Icon name="cube" />;
 
+    case ResourceRowType.VariableGroup:
+      return <Icon name="x" />;
+
+    case ResourceRowType.Variable:
+      return <Icon name="x" />;
+
     default:
       return null;
   }
@@ -157,7 +163,7 @@ const NestedEntry: React.FC<NestedEntryProps> = ({
   const theme = useTheme2();
   const styles = useStyles2(getStyles);
   const hasChildren = !!entry.children;
-  const isSelectable = entry.type === ResourceRowType.Resource;
+  const isSelectable = entry.type === ResourceRowType.Resource || entry.type === ResourceRowType.Variable;
 
   const handleToggleCollapse = useCallback(() => {
     onToggleCollapse(entry);

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

@@ -5,22 +5,28 @@ 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 { findRow, parseResourceURI } from './utils';
+import { addResources, findRow, parseResourceURI } from './utils';
 
 interface ResourcePickerProps {
-  resourcePickerData: Pick<ResourcePickerData, 'getResourcePickerData' | 'getResourcesForResourceGroup'>;
+  resourcePickerData: ResourcePickerData;
   resourceURI: string | undefined;
+  templateVariables: string[];
 
   onApply: (resourceURI: string | undefined) => void;
   onCancel: () => void;
 }
 
-const ResourcePicker = ({ resourcePickerData, resourceURI, onApply, onCancel }: ResourcePickerProps) => {
+const ResourcePicker = ({
+  resourcePickerData,
+  resourceURI,
+  templateVariables,
+  onApply,
+  onCancel,
+}: ResourcePickerProps) => {
   const styles = useStyles2(getStyles);
 
-  const [rows, setRows] = useState<ResourceRowGroup>([]);
+  const [azureRows, setAzureRows] = useState<ResourceRowGroup>([]);
   const [internalSelected, setInternalSelected] = useState<string | undefined>(resourceURI);
 
   // Sync the resourceURI prop to internal state
@@ -28,6 +34,11 @@ const ResourcePicker = ({ resourcePickerData, resourceURI, onApply, onCancel }:
     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);
@@ -37,28 +48,18 @@ const ResourcePicker = ({ resourcePickerData, resourceURI, onApply, onCancel }:
   // 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
-      if (resourceGroup.children?.length) {
+      // 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 = 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 (!draftRow) {
-          // This case shouldn't happen
-          throw new Error('Unable to find resource');
-        }
-
-        draftRow.children = resources;
-      });
-
-      setRows(newRows);
+      const newRows = addResources(azureRows, resourceGroup.id, resources);
+      setAzureRows(newRows);
     },
-    [resourcePickerData, rows]
+    [resourcePickerData, azureRows]
   );
 
   // Select
@@ -68,7 +69,9 @@ const ResourcePicker = ({ resourcePickerData, resourceURI, onApply, onCancel }:
 
   // Request initial data on first mount
   useEffect(() => {
-    resourcePickerData.getResourcePickerData().then((initalRows) => setRows(initalRows));
+    resourcePickerData.getResourcePickerData().then((initalRows) => {
+      setAzureRows(initalRows);
+    });
   }, [resourcePickerData]);
 
   // Request sibling resources for a selected resource - in practice should only be on first mount
@@ -87,9 +90,9 @@ const ResourcePicker = ({ resourcePickerData, resourceURI, onApply, onCancel }:
     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');
+      // We haven't loaded the data from Azure yet
+      return;
     }
 
     requestNestedRows(resourceGroupRow);

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

@@ -2,6 +2,8 @@ export enum ResourceRowType {
   Subscription = 'Subscription',
   ResourceGroup = 'ResourceGroup',
   Resource = 'Resource',
+  VariableGroup = 'TemplateVariableGroup',
+  Variable = 'TemplateVariable',
 }
 
 export interface ResourceRow {

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

@@ -1,3 +1,4 @@
+import produce from 'immer';
 import { ResourceRow, ResourceRowGroup } from './types';
 
 const RESOURCE_URI_REGEX = /\/subscriptions\/(?<subscriptionID>.+)\/resourceGroups\/(?<resourceGroup>.+)\/providers.+\/(?<resource>[\w-_]+)/;
@@ -34,3 +35,16 @@ export function findRow(rows: ResourceRowGroup, id: string): ResourceRow | undef
 
   return undefined;
 }
+
+export function addResources(rows: ResourceRowGroup, targetResourceGroupID: string, newResources: ResourceRowGroup) {
+  return produce(rows, (draftState) => {
+    const draftRow = findRow(draftState, targetResourceGroupID);
+
+    if (!draftRow) {
+      // This case shouldn't happen often because we're usually coming here from a resource we already have
+      throw new Error('Unable to find resource');
+    }
+
+    draftRow.children = newResources;
+  });
+}

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

@@ -22,6 +22,8 @@ export default class ResourcePickerData {
     this.cloud = getAzureCloud(instanceSettings);
   }
 
+  static readonly templateVariableGroupID = '$$grafana-templateVariables$$';
+
   async getResourcePickerData() {
     const query = `
       resources
@@ -141,6 +143,21 @@ export default class ResourcePickerData {
       throw error;
     }
   }
+
+  transformVariablesToRow(templateVariables: string[]): ResourceRow {
+    return {
+      id: ResourcePickerData.templateVariableGroupID,
+      name: 'Template variables',
+      type: ResourceRowType.VariableGroup,
+      typeLabel: 'Variables',
+      children: templateVariables.map((v) => ({
+        id: v,
+        name: v,
+        type: ResourceRowType.Variable,
+        typeLabel: 'Variable',
+      })),
+    };
+  }
 }
 
 function formatResourceGroupData(rawData: RawAzureResourceGroupItem[]) {