AzureMonitor: Add support for Subscriptions template variable (#52086)

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

@@ -0,0 +1,59 @@
+import React, { ChangeEvent, useCallback, useEffect, useState } from 'react';
+
+import { InlineField, Input } from '@grafana/ui';
+
+import DataSource from '../../datasource';
+import { migrateStringQueriesToObjectQueries } from '../../grafanaTemplateVariableFns';
+import { AzureMonitorQuery, AzureQueryType } from '../../types';
+
+const GrafanaTemplateVariableFnInput = ({
+  query,
+  updateQuery,
+  datasource,
+}: {
+  query: AzureMonitorQuery;
+  updateQuery: (val: AzureMonitorQuery) => void;
+  datasource: DataSource;
+}) => {
+  const [inputVal, setInputVal] = useState('');
+
+  useEffect(() => {
+    setInputVal(query.grafanaTemplateVariableFn?.rawQuery || '');
+  }, [query.grafanaTemplateVariableFn?.rawQuery]);
+
+  const onRunQuery = useCallback(
+    (newQuery: string) => {
+      migrateStringQueriesToObjectQueries(newQuery, { datasource }).then((updatedQuery) => {
+        if (updatedQuery.queryType === AzureQueryType.GrafanaTemplateVariableFn) {
+          updateQuery(updatedQuery);
+        } else {
+          updateQuery({
+            ...query,
+            grafanaTemplateVariableFn: {
+              kind: 'UnknownQuery',
+              rawQuery: newQuery,
+            },
+          });
+        }
+      });
+    },
+    [datasource, query, updateQuery]
+  );
+
+  const onChange = (event: ChangeEvent<HTMLInputElement>) => {
+    setInputVal(event.target.value);
+  };
+
+  return (
+    <InlineField label="Grafana template variable function">
+      <Input
+        placeholder={'type a grafana template variable function, ex: Subscriptions()'}
+        value={inputVal}
+        onChange={onChange}
+        onBlur={() => onRunQuery(inputVal)}
+      />
+    </InlineField>
+  );
+};
+
+export default GrafanaTemplateVariableFnInput;

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

@@ -1,8 +1,9 @@
 import { render, screen, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
-import { select } from 'react-select-event';
+import { select, openMenu } from 'react-select-event';
 
+import * as grafanaRuntime from '@grafana/runtime';
 import * as ui from '@grafana/ui';
 
 import createMockDatasource from '../../__mocks__/datasource';
@@ -18,23 +19,28 @@ jest.mock('@grafana/ui', () => ({
   },
 }));
 
+const defaultProps = {
+  query: {
+    refId: 'A',
+    queryType: AzureQueryType.LogAnalytics,
+    azureLogAnalytics: {
+      query: 'test query',
+    },
+    subscription: 'id',
+  },
+  onChange: jest.fn(),
+  datasource: createMockDatasource(),
+};
+
+const originalConfigValue = grafanaRuntime.config.featureToggles.azTemplateVars;
+beforeEach(() => {
+  // reset config
+  grafanaRuntime.config.featureToggles.azTemplateVars = originalConfigValue;
+});
+
 describe('VariableEditor:', () => {
   it('can select a query type', async () => {
-    const onChange = jest.fn();
+    render(<VariableEditor {...defaultProps} />);
-
-    const props = {
-      query: {
-        refId: 'A',
-        queryType: AzureQueryType.LogAnalytics,
-        azureLogAnalytics: {
-          query: 'test query',
-        },
-        subscription: 'id',
-      },
-      onChange,
-      datasource: createMockDatasource(),
-    };
-    render(<VariableEditor {...props} />);
     await waitFor(() => screen.getByLabelText('select query type'));
     expect(screen.getByLabelText('select query type')).toBeInTheDocument();
     screen.getByLabelText('select query type').click();
@@ -46,19 +52,7 @@ describe('VariableEditor:', () => {
   });
   describe('log queries:', () => {
     it('should render', async () => {
-      const props = {
+      render(<VariableEditor {...defaultProps} />);
-        query: {
-          refId: 'A',
-          queryType: AzureQueryType.LogAnalytics,
-          azureLogAnalytics: {
-            query: 'test query',
-          },
-          subscription: 'id',
-        },
-        onChange: () => {},
-        datasource: createMockDatasource(),
-      };
-      render(<VariableEditor {...props} />);
       await waitFor(() => screen.queryByTestId('mockeditor'));
       expect(screen.queryByText('Resource')).toBeInTheDocument();
       expect(screen.queryByTestId('mockeditor')).toBeInTheDocument();
@@ -75,24 +69,14 @@ describe('VariableEditor:', () => {
       expect(screen.queryByText('Resource')).toBeInTheDocument();
       expect(screen.queryByTestId('mockeditor')).toBeInTheDocument();
     });
+
     it('should call on change if the query changes', async () => {
-      const props = {
+      const onChange = jest.fn();
-        query: {
+      render(<VariableEditor {...defaultProps} onChange={onChange} />);
-          refId: 'A',
-          queryType: AzureQueryType.LogAnalytics,
-          azureLogAnalytics: {
-            query: 'test query',
-          },
-          subscription: 'id',
-        },
-        onChange: jest.fn(),
-        datasource: createMockDatasource(),
-      };
-      render(<VariableEditor {...props} />);
       await waitFor(() => screen.queryByTestId('mockeditor'));
       expect(screen.queryByTestId('mockeditor')).toBeInTheDocument();
       await userEvent.type(screen.getByTestId('mockeditor'), '{backspace}');
-      expect(props.onChange).toHaveBeenCalledWith({
+      expect(onChange).toHaveBeenCalledWith({
         azureLogAnalytics: {
           query: 'test quer',
         },
@@ -106,6 +90,7 @@ describe('VariableEditor:', () => {
   describe('grafana template variable fn queries:', () => {
     it('should render', async () => {
       const props = {
+        ...defaultProps,
         query: {
           refId: 'A',
           queryType: AzureQueryType.GrafanaTemplateVariableFn,
@@ -115,8 +100,6 @@ describe('VariableEditor:', () => {
           },
           subscription: 'id',
         } as AzureMonitorQuery,
-        onChange: () => {},
-        datasource: createMockDatasource(),
       };
       render(<VariableEditor {...props} />);
       await waitFor(() => screen.queryByText('Grafana template variable function'));
@@ -126,6 +109,7 @@ describe('VariableEditor:', () => {
 
     it('should call on change if the query changes', async () => {
       const props = {
+        ...defaultProps,
         query: {
           refId: 'A',
           queryType: AzureQueryType.GrafanaTemplateVariableFn,
@@ -135,8 +119,6 @@ describe('VariableEditor:', () => {
           },
           subscription: 'subscriptionId',
         } as AzureMonitorQuery,
-        onChange: jest.fn(),
-        datasource: createMockDatasource(),
       };
       render(<VariableEditor {...props} />);
       await waitFor(() => screen.queryByText('Grafana template variable function'));
@@ -155,4 +137,30 @@ describe('VariableEditor:', () => {
       });
     });
   });
+
+  describe('predefined queries:', () => {
+    it('should show the new query types if feature gate is enabled', async () => {
+      grafanaRuntime.config.featureToggles.azTemplateVars = true;
+      render(<VariableEditor {...defaultProps} />);
+      openMenu(screen.getByLabelText('select query type'));
+      await waitFor(() => expect(screen.getByText('Subscriptions')).toBeInTheDocument());
+    });
+
+    it('should not show the new query types if feature gate is disabled', async () => {
+      grafanaRuntime.config.featureToggles.azTemplateVars = false;
+      render(<VariableEditor {...defaultProps} />);
+      openMenu(screen.getByLabelText('select query type'));
+      await waitFor(() => expect(screen.queryByText('Subscriptions')).not.toBeInTheDocument());
+    });
+
+    it('should run the query if requesting subscriptions', async () => {
+      grafanaRuntime.config.featureToggles.azTemplateVars = true;
+      const onChange = jest.fn();
+      render(<VariableEditor {...defaultProps} onChange={onChange} />);
+      openMenu(screen.getByLabelText('select query type'));
+      screen.getByText('Subscriptions').click();
+      await waitFor(() => expect(screen.getByText('Subscriptions')).toBeInTheDocument());
+      expect(onChange).toHaveBeenCalledWith({ queryType: AzureQueryType.SubscriptionsQuery, refId: 'A' });
+    });
+  });
 });

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

@@ -1,7 +1,8 @@
-import React, { ChangeEvent, useCallback, useEffect, useState } from 'react';
+import React, { useEffect, useState } from 'react';
 
 import { SelectableValue } from '@grafana/data';
-import { Alert, InlineField, Input, Select } from '@grafana/ui';
+import { config } from '@grafana/runtime';
+import { Alert, InlineField, Select } from '@grafana/ui';
 
 import DataSource from '../../datasource';
 import { migrateStringQueriesToObjectQueries } from '../../grafanaTemplateVariableFns';
@@ -10,59 +11,7 @@ import useLastError from '../../utils/useLastError';
 import LogsQueryEditor from '../LogsQueryEditor';
 import { Space } from '../Space';
 
-const AZURE_QUERY_VARIABLE_TYPE_OPTIONS = [
+import GrafanaTemplateVariableFnInput from './GrafanaTemplateVariableFn';
-  { label: 'Grafana Query Function', value: AzureQueryType.GrafanaTemplateVariableFn },
-  { label: 'Logs', value: AzureQueryType.LogAnalytics },
-];
-
-const GrafanaTemplateVariableFnInput = ({
-  query,
-  updateQuery,
-  datasource,
-}: {
-  query: AzureMonitorQuery;
-  updateQuery: (val: AzureMonitorQuery) => void;
-  datasource: DataSource;
-}) => {
-  const [inputVal, setInputVal] = useState('');
-  useEffect(() => {
-    setInputVal(query.grafanaTemplateVariableFn?.rawQuery || '');
-  }, [query.grafanaTemplateVariableFn?.rawQuery]);
-
-  const onRunQuery = useCallback(
-    (newQuery: string) => {
-      migrateStringQueriesToObjectQueries(newQuery, { datasource }).then((updatedQuery) => {
-        if (updatedQuery.queryType === AzureQueryType.GrafanaTemplateVariableFn) {
-          updateQuery(updatedQuery);
-        } else {
-          updateQuery({
-            ...query,
-            grafanaTemplateVariableFn: {
-              kind: 'UnknownQuery',
-              rawQuery: newQuery,
-            },
-          });
-        }
-      });
-    },
-    [datasource, query, updateQuery]
-  );
-
-  const onChange = (event: ChangeEvent<HTMLInputElement>) => {
-    setInputVal(event.target.value);
-  };
-
-  return (
-    <InlineField label="Grafana template variable function">
-      <Input
-        placeholder={'type a grafana template variable function, ex: Subscriptions()'}
-        value={inputVal}
-        onChange={onChange}
-        onBlur={() => onRunQuery(inputVal)}
-      />
-    </InlineField>
-  );
-};
 
 type Props = {
   query: AzureMonitorQuery | string;
@@ -75,6 +24,14 @@ const VariableEditor = (props: Props) => {
     refId: 'A',
     queryType: AzureQueryType.GrafanaTemplateVariableFn,
   };
+  const AZURE_QUERY_VARIABLE_TYPE_OPTIONS = [
+    { label: 'Grafana Query Function', value: AzureQueryType.GrafanaTemplateVariableFn },
+    { label: 'Logs', value: AzureQueryType.LogAnalytics },
+  ];
+  if (config.featureToggles.azTemplateVars) {
+    AZURE_QUERY_VARIABLE_TYPE_OPTIONS.push({ label: 'Subscriptions', value: AzureQueryType.SubscriptionsQuery });
+  }
+
   const [query, setQuery] = useState(defaultQuery);
 
   useEffect(() => {
@@ -85,12 +42,15 @@ const VariableEditor = (props: Props) => {
 
   const onQueryTypeChange = (selectableValue: SelectableValue) => {
     if (selectableValue.value) {
-      setQuery({
+      const newQuery = {
         ...query,
         queryType: selectableValue.value,
-      });
+      };
+      setQuery(newQuery);
+      props.onChange(newQuery);
     }
   };
+
   const onLogsQueryChange = (queryChange: AzureMonitorQuery) => {
     setQuery(queryChange);
 
@@ -111,7 +71,7 @@ const VariableEditor = (props: Props) => {
 
   return (
     <>
-      <InlineField label="Select query type">
+      <InlineField label="Select query type" labelWidth={20}>
         <Select
           aria-label="select query type"
           onChange={onQueryTypeChange}

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

@@ -6,6 +6,8 @@ export enum AzureQueryType {
   AzureMonitor = 'Azure Monitor',
   LogAnalytics = 'Azure Log Analytics',
   AzureResourceGraph = 'Azure Resource Graph',
+  SubscriptionsQuery = 'Azure Subscriptions',
+  /** Deprecated */
   GrafanaTemplateVariableFn = 'Grafana Template Variable Function',
 }
 

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

@@ -513,4 +513,28 @@ describe('VariableSupport', () => {
       done();
     });
   });
+
+  describe('querying for subscriptions', () => {
+    it('can fetch subscriptions', (done) => {
+      const fakeSubscriptions = ['subscriptionId'];
+      const variableSupport = new VariableSupport(
+        createMockDatasource({
+          getSubscriptions: jest.fn().mockResolvedValueOnce(fakeSubscriptions),
+        })
+      );
+      const mockRequest = {
+        targets: [
+          {
+            refId: 'A',
+            queryType: AzureQueryType.SubscriptionsQuery,
+          } as AzureMonitorQuery,
+        ],
+      } as DataQueryRequest<AzureMonitorQuery>;
+      const observables = variableSupport.query(mockRequest);
+      observables.subscribe((result: DataQueryResponseData) => {
+        expect(result.data[0].source).toEqual(fakeSubscriptions);
+        done();
+      });
+    });
+  });
 });

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

@@ -29,18 +29,29 @@ export class VariableSupport extends CustomVariableSupport<DataSource, AzureMoni
     const promisedResults = async () => {
       const queryObj = await migrateStringQueriesToObjectQueries(request.targets[0], { datasource: this.datasource });
 
-      if (queryObj.queryType === AzureQueryType.GrafanaTemplateVariableFn && queryObj.grafanaTemplateVariableFn) {
+      try {
-        try {
+        switch (queryObj.queryType) {
-          const templateVariablesResults = await this.callGrafanaTemplateVariableFn(queryObj.grafanaTemplateVariableFn);
+          case AzureQueryType.SubscriptionsQuery:
-          return {
+            const res = await this.datasource.getSubscriptions();
-            data: templateVariablesResults?.length ? [toDataFrame(templateVariablesResults)] : [],
+            return {
-          };
+              data: res?.length ? [toDataFrame(res)] : [],
-        } catch (err) {
+            };
-          return { data: [], error: { message: messageFromError(err) } };
+          case AzureQueryType.GrafanaTemplateVariableFn:
+            if (queryObj.grafanaTemplateVariableFn) {
+              const templateVariablesResults = await this.callGrafanaTemplateVariableFn(
+                queryObj.grafanaTemplateVariableFn
+              );
+              return {
+                data: templateVariablesResults?.length ? [toDataFrame(templateVariablesResults)] : [],
+              };
+            }
+          default:
+            request.targets[0] = queryObj;
+            return lastValueFrom(this.datasource.query(request));
         }
+      } catch (err) {
+        return { data: [], error: { message: messageFromError(err) } };
       }
-      request.targets[0] = queryObj;
-      return lastValueFrom(this.datasource.query(request));
     };
 
     return from(promisedResults());