From c4967c3d682f4a9419ec6d9ceeb892a9c6ffbe68 Mon Sep 17 00:00:00 2001
From: yaellec <yaelle.chaudy@grafana.com>
Date: Mon, 1 Aug 2022 18:08:35 +0200
Subject: [PATCH] Added ARG query functionfor template variables

---
 .../azuremonitor/template-variables.md        |  1 +
 .../VariableEditor/VariableEditor.test.tsx    | 37 +++++++++++++++++++
 .../VariableEditor/VariableEditor.tsx         | 28 ++++++++++++--
 3 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/docs/sources/datasources/azuremonitor/template-variables.md b/docs/sources/datasources/azuremonitor/template-variables.md
index b36e0cd6034..afe5d03c773 100644
--- a/docs/sources/datasources/azuremonitor/template-variables.md
+++ b/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:
 
diff --git a/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/VariableEditor/VariableEditor.test.tsx b/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/VariableEditor/VariableEditor.test.tsx
index 4ec906daf74..7ae4a0c5082 100644
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/VariableEditor/VariableEditor.test.tsx
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/VariableEditor/VariableEditor.test.tsx
@@ -82,6 +82,43 @@ describe('VariableEditor:', () => {
     });
   });
 
+  describe('Azure Resource Graph queries:', () => {
+    const ARGquery = {
+      refId: 'A',
+      queryType: AzureQueryType.AzureResourceGraph,
+      azureResourceGraph: {
+        query: 'Resources | distinct type',
+      },
+      subscriptions: ['Subscription1', 'Subscription2'],
+    };
+
+    it('should render', async () => {
+      render(<VariableEditor {...defaultProps} query={ARGquery} />);
+      await waitFor(() => screen.queryByTestId('mockeditor'));
+      expect(screen.queryByLabelText('Subscriptions')).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 {...defaultProps} query={ARGquery} onChange={onChange} />);
+      await waitFor(() => screen.queryByTestId('mockeditor'));
+      expect(screen.queryByTestId('mockeditor')).toBeInTheDocument();
+      await userEvent.type(screen.getByTestId('mockeditor'), '{backspace}');
+      expect(onChange).toHaveBeenCalledWith({
+        ...ARGquery,
+        azureResourceGraph: {
+          query: 'Resources | distinct typ',
+        },
+      });
+    });
+  });
+
   describe('grafana template variable fn queries:', () => {
     it('should render', async () => {
       const props = {
diff --git a/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/VariableEditor/VariableEditor.tsx b/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/VariableEditor/VariableEditor.tsx
index de08439e49f..5bc5ee07c12 100644
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/VariableEditor/VariableEditor.tsx
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/components/VariableEditor/VariableEditor.tsx
@@ -7,8 +7,9 @@ import { Alert, InlineField, Select } from '@grafana/ui';
 
 import DataSource from '../../datasource';
 import { migrateQuery } from '../../grafanaTemplateVariableFns';
-import { AzureMonitorOption, AzureMonitorQuery, AzureQueryType } from '../../types';
+import { AzureMonitorOption, AzureMonitorQuery, AzureQueryType, AzureResourceGraphQuery } from '../../types';
 import useLastError from '../../utils/useLastError';
+import ArgQueryEditor from '../ArgQueryEditor';
 import LogsQueryEditor from '../LogsQueryEditor';
 import { Space } from '../Space';
 
@@ -32,6 +33,7 @@ const VariableEditor = (props: Props) => {
     { label: 'Metric Names', value: AzureQueryType.MetricNamesQuery },
     { label: 'Workspaces', value: AzureQueryType.WorkspacesQuery },
     { label: 'Logs', value: AzureQueryType.LogAnalytics },
+    { label: 'Resource Graph', value: AzureQueryType.AzureResourceGraph },
   ];
   if (typeof props.query === 'object' && props.query.queryType === AzureQueryType.GrafanaTemplateVariableFn) {
     // Add the option for the GrafanaTemplateVariableFn only if it's already in use
@@ -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={query.subscription}
+            query={query}
+            datasource={datasource}
+            onChange={onQueryChange}
+            variableOptionGroup={variableOptionGroup}
+            setError={setError}
+          />
+          {errorMessage && (
+            <>
+              <Space v={2} />
+              <Alert severity="error" title="An error occurred while requesting data from Azure Monitor">
+                {errorMessage}
+              </Alert>
+            </>
+          )}
+        </>
+      )}
     </>
   );
 };