AzureMonitor: Added ARG query function for template variables (#53059) (#54204)

* Added ARG query functionfor template variables

* removed unused import

* Tweak tests

* Fix tests

* Renamed query function - added 'Azure'

* Revert "Renamed query function - added 'Azure'"

This reverts commit a046bcdd7c.

* Adressed comments

* Fix uncaught error issue

* revert last commit

* Add back errorMessage check

(cherry picked from commit d0d6562f63)

Co-authored-by: Yaelle Chaudy <42030685+yaelleC@users.noreply.github.com>

docs/sources/datasources/azuremonitor/template-variables.md

@@ -34,6 +34,7 @@ The Azure Monitor data source provides the following queries you can specify in
 | Metric Names    | Returns a list of metric names for a resource.                                               |
 | Workspaces      | Returns a list of workspaces for the specified subscription.                                 |
 | Logs            | Use a KQL query to return values.                                                            |
+| Resource Graph  | Use an ARG query to return values.                                                           |
 
 Any Log Analytics KQL query that returns a single list of values can also be used in the Query field. For example:
 

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

@@ -35,7 +35,7 @@ const ArgQueryEditor: React.FC<ArgQueryEditorProps> = ({
     }
 
     fetchedRef.current = true;
-    datasource.azureMonitorDatasource
+    datasource
       .getSubscriptions()
       .then((results) => {
         const fetchedSubscriptions = results.map((v) => ({ label: v.text, value: v.value, description: v.value }));

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

@@ -57,6 +57,7 @@ describe('VariableEditor:', () => {
       })
     );
   });
+
   describe('log queries:', () => {
     it('should render', async () => {
       render(<VariableEditor {...defaultProps} />);
@@ -82,6 +83,49 @@ describe('VariableEditor:', () => {
     });
   });
 
+  describe('Azure Resource Graph queries:', () => {
+    const ARGqueryProps = {
+      ...defaultProps,
+      query: {
+        refId: 'A',
+        queryType: AzureQueryType.AzureResourceGraph,
+        azureResourceGraph: {
+          query: 'Resources | distinct type',
+          resultFormat: 'table',
+        },
+        subscriptions: ['sub'],
+      },
+    };
+
+    it('should render', async () => {
+      render(<VariableEditor {...ARGqueryProps} />);
+      await waitFor(() => screen.queryByTestId('mockeditor'));
+      expect(screen.queryByLabelText('Subscriptions')).toBeInTheDocument();
+      expect(screen.queryByText('Resource Graph')).toBeInTheDocument();
+      expect(screen.queryByLabelText('Select subscription')).not.toBeInTheDocument();
+      expect(screen.queryByLabelText('Select query type')).not.toBeInTheDocument();
+      expect(screen.queryByLabelText('Select resource group')).not.toBeInTheDocument();
+      expect(screen.queryByLabelText('Select namespace')).not.toBeInTheDocument();
+      expect(screen.queryByLabelText('Select resource')).not.toBeInTheDocument();
+      expect(screen.queryByTestId('mockeditor')).toBeInTheDocument();
+    });
+
+    it('should call on change if the query changes', async () => {
+      const onChange = jest.fn();
+      render(<VariableEditor {...ARGqueryProps} onChange={onChange} />);
+      await waitFor(() => screen.queryByTestId('mockeditor'));
+      expect(screen.queryByTestId('mockeditor')).toBeInTheDocument();
+      await userEvent.type(screen.getByTestId('mockeditor'), '{backspace}');
+      expect(onChange).toHaveBeenCalledWith({
+        ...ARGqueryProps.query,
+        azureResourceGraph: {
+          query: 'Resources | distinct typ',
+          resultFormat: 'table',
+        },
+      });
+    });
+  });
+
   describe('grafana template variable fn queries:', () => {
     it('should render', async () => {
       const props = {

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

@@ -9,6 +9,7 @@ import DataSource from '../../datasource';
 import { migrateQuery } from '../../grafanaTemplateVariableFns';
 import { AzureMonitorOption, AzureMonitorQuery, AzureQueryType } from '../../types';
 import useLastError from '../../utils/useLastError';
+import ArgQueryEditor from '../ArgQueryEditor';
 import LogsQueryEditor from '../LogsQueryEditor';
 import { Space } from '../Space';
 
@@ -31,6 +32,7 @@ const VariableEditor = (props: Props) => {
     { label: 'Resource Names', value: AzureQueryType.ResourceNamesQuery },
     { label: 'Metric Names', value: AzureQueryType.MetricNamesQuery },
     { label: 'Workspaces', value: AzureQueryType.WorkspacesQuery },
+    { label: 'Resource Graph', value: AzureQueryType.AzureResourceGraph },
     { label: 'Logs', value: AzureQueryType.LogAnalytics },
   ];
   if (typeof props.query === 'object' && props.query.queryType === AzureQueryType.GrafanaTemplateVariableFn) {
@@ -193,7 +195,7 @@ const VariableEditor = (props: Props) => {
     });
   };
 
-  const onLogsQueryChange = (queryChange: AzureMonitorQuery) => {
+  const onQueryChange = (queryChange: AzureMonitorQuery) => {
     onChange(queryChange);
   };
 
@@ -214,7 +216,7 @@ const VariableEditor = (props: Props) => {
             subscriptionId={query.subscription}
             query={query}
             datasource={datasource}
-            onChange={onLogsQueryChange}
+            onChange={onQueryChange}
             variableOptionGroup={variableOptionGroup}
             setError={setError}
             hideFormatAs={true}
@@ -286,6 +288,26 @@ const VariableEditor = (props: Props) => {
           />
         </InlineField>
       )}
+      {query.queryType === AzureQueryType.AzureResourceGraph && (
+        <>
+          <ArgQueryEditor
+            subscriptionId={datasource.azureLogAnalyticsDatasource.defaultSubscriptionId}
+            query={query}
+            datasource={datasource}
+            onChange={onQueryChange}
+            variableOptionGroup={variableOptionGroup}
+            setError={setError}
+          />
+          {errorMessage && (
+            <>
+              <Space v={2} />
+              <Alert severity="error" title="An error occurred while requesting metadata from Azure Monitor">
+                {errorMessage}
+              </Alert>
+            </>
+          )}
+        </>
+      )}
     </>
   );
 };

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

@@ -91,7 +91,10 @@ export class VariableSupport extends CustomVariableSupport<DataSource, AzureMoni
             }
           default:
             request.targets[0] = queryObj;
-            return lastValueFrom(this.datasource.query(request));
+            const queryResp = await lastValueFrom(this.datasource.query(request));
+            return {
+              data: queryResp.data,
+            };
         }
       } catch (err) {
         return { data: [], error: new Error(messageFromError(err)) };