AzureMonitor: Add ResourceGroups template variable (#52141)

2025-02-25 18:55:37 -06:00 · 2022-07-14 09:48:11 +02:00 · 2022-07-14 09:48:11 +02:00 · 99d9c3d0fd
commit 99d9c3d0fd
parent 07e03666ad
7 changed files with 192 additions and 50 deletions
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/mocks/datasource.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/mocks/datasource.ts
@ -28,6 +28,7 @@ export default function createMockDatasource(overrides?: DeepPartial<Datasource>

    getAzureLogAnalyticsWorkspaces: jest.fn().mockResolvedValueOnce([]),

+    getSubscriptions: jest.fn().mockResolvedValue([]),
    getResourceGroups: jest.fn().mockResolvedValueOnce([]),
    getMetricDefinitions: jest.fn().mockResolvedValueOnce([]),
    getResourceNames: jest.fn().mockResolvedValueOnce([]),
@ -43,6 +44,7 @@ export default function createMockDatasource(overrides?: DeepPartial<Datasource>
      getResourceURIFromWorkspace: jest.fn().mockReturnValue(''),
      getResourceURIDisplayProperties: jest.fn().mockResolvedValue({}),
    },
+    getVariablesRaw: jest.fn().mockReturnValue([]),
    ...overrides,
  };

--- 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
@ -40,13 +40,21 @@ beforeEach(() => {

 describe('VariableEditor:', () => {
  it('can select a query type', async () => {
-    render(<VariableEditor {...defaultProps} />);
+    const onChange = jest.fn();
+    const { rerender } = render(<VariableEditor {...defaultProps} onChange={onChange} />);
    await waitFor(() => screen.getByLabelText('select query type'));
    expect(screen.getByLabelText('select query type')).toBeInTheDocument();
    screen.getByLabelText('select query type').click();
    await select(screen.getByLabelText('select query type'), 'Grafana Query Function', {
      container: document.body,
    });
+    expect(onChange).toHaveBeenCalledWith(
+      expect.objectContaining({
+        queryType: AzureQueryType.GrafanaTemplateVariableFn,
+      })
+    );
+    const newQuery = onChange.mock.calls.at(-1)[0];
+    rerender(<VariableEditor {...defaultProps} query={newQuery} />);
    expect(screen.queryByText('Logs')).not.toBeInTheDocument();
    expect(screen.queryByText('Grafana Query Function')).toBeInTheDocument();
  });
@ -58,18 +66,6 @@ describe('VariableEditor:', () => {
      expect(screen.queryByTestId('mockeditor')).toBeInTheDocument();
    });

-    it('should render with legacy query strings', async () => {
-      const props = {
-        query: 'test query',
-        onChange: () => {},
-        datasource: createMockDatasource(),
-      };
-      render(<VariableEditor {...props} />);
-      await waitFor(() => screen.queryByTestId('mockeditor'));
-      expect(screen.queryByText('Resource')).toBeInTheDocument();
-      expect(screen.queryByTestId('mockeditor')).toBeInTheDocument();
-    });
-
    it('should call on change if the query changes', async () => {
      const onChange = jest.fn();
      render(<VariableEditor {...defaultProps} onChange={onChange} />);
@ -156,11 +152,76 @@ describe('VariableEditor:', () => {
    it('should run the query if requesting subscriptions', async () => {
      grafanaRuntime.config.featureToggles.azTemplateVars = true;
      const onChange = jest.fn();
-      render(<VariableEditor {...defaultProps} onChange={onChange} />);
+      const { rerender } = render(<VariableEditor {...defaultProps} onChange={onChange} />);
      openMenu(screen.getByLabelText('select query type'));
      screen.getByText('Subscriptions').click();
+      // Simulate onChange behavior
+      const newQuery = onChange.mock.calls.at(-1)[0];
+      rerender(<VariableEditor {...defaultProps} query={newQuery} onChange={onChange} />);
      await waitFor(() => expect(screen.getByText('Subscriptions')).toBeInTheDocument());
-      expect(onChange).toHaveBeenCalledWith({ queryType: AzureQueryType.SubscriptionsQuery, refId: 'A' });
+      expect(onChange).toHaveBeenCalledWith(
+        expect.objectContaining({ queryType: AzureQueryType.SubscriptionsQuery, refId: 'A' })
+      );
+    });
+
+    it('should run the query if requesting resource groups', async () => {
+      grafanaRuntime.config.featureToggles.azTemplateVars = true;
+      const ds = createMockDatasource({
+        getSubscriptions: jest.fn().mockResolvedValue([{ text: 'Primary Subscription', value: 'sub' }]),
+      });
+      const onChange = jest.fn();
+      const { rerender } = render(<VariableEditor {...defaultProps} onChange={onChange} datasource={ds} />);
+      // wait for initial load
+      await waitFor(() => expect(screen.getByText('Logs')).toBeInTheDocument());
+      // Select RGs variable
+      openMenu(screen.getByLabelText('select query type'));
+      screen.getByText('Resource Groups').click();
+      // Simulate onChange behavior
+      const newQuery = onChange.mock.calls.at(-1)[0];
+      rerender(<VariableEditor {...defaultProps} query={newQuery} onChange={onChange} />);
+      await waitFor(() => expect(screen.getByText('Select subscription')).toBeInTheDocument());
+      // Select a subscription
+      openMenu(screen.getByLabelText('select subscription'));
+      screen.getByText('Primary Subscription').click();
+      expect(onChange).toHaveBeenCalledWith(
+        expect.objectContaining({
+          queryType: AzureQueryType.ResourceGroupsQuery,
+          subscription: 'sub',
+          refId: 'A',
+        })
+      );
+    });
+
+    it('should show template variables as options ', async () => {
+      const onChange = jest.fn();
+      grafanaRuntime.config.featureToggles.azTemplateVars = true;
+      const ds = createMockDatasource({
+        getSubscriptions: jest.fn().mockResolvedValue([{ text: 'Primary Subscription', value: 'sub' }]),
+        getVariablesRaw: jest.fn().mockReturnValue([
+          { label: 'query0', name: 'sub0' },
+          { label: 'query1', name: 'rg', query: { queryType: AzureQueryType.ResourceGroupsQuery } },
+        ]),
+      });
+      const { rerender } = render(<VariableEditor {...defaultProps} datasource={ds} onChange={onChange} />);
+      // wait for initial load
+      await waitFor(() => expect(screen.getByText('Logs')).toBeInTheDocument());
+      // Select RGs variable
+      openMenu(screen.getByLabelText('select query type'));
+      screen.getByText('Resource Groups').click();
+      // Simulate onChange behavior
+      const newQuery = onChange.mock.calls.at(-1)[0];
+      rerender(<VariableEditor {...defaultProps} query={newQuery} onChange={onChange} datasource={ds} />);
+      await waitFor(() => expect(screen.getByText('Select subscription')).toBeInTheDocument());
+      // Select a subscription
+      openMenu(screen.getByLabelText('select subscription'));
+      await waitFor(() => expect(screen.getByText('Primary Subscription')).toBeInTheDocument());
+      screen.getByText('Template Variables').click();
+      // Simulate onChange behavior
+      const lastQuery = onChange.mock.calls.at(-1)[0];
+      rerender(<VariableEditor {...defaultProps} query={lastQuery} onChange={onChange} datasource={ds} />);
+      await waitFor(() => expect(screen.getByText('query0')).toBeInTheDocument());
+      // Template variables of the same type than the current one should not appear
+      expect(screen.queryByText('query1')).not.toBeInTheDocument();
    });
  });
 });
--- 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
@ -1,4 +1,6 @@
+import { get } from 'lodash';
 import React, { useEffect, useState } from 'react';
+import { useEffectOnce } from 'react-use';

 import { SelectableValue } from '@grafana/data';
 import { config } from '@grafana/runtime';
@ -6,7 +8,7 @@ import { Alert, InlineField, Select } from '@grafana/ui';

 import DataSource from '../../datasource';
 import { migrateStringQueriesToObjectQueries } from '../../grafanaTemplateVariableFns';
-import { AzureMonitorQuery, AzureQueryType } from '../../types';
+import { AzureMonitorOption, AzureMonitorQuery, AzureQueryType } from '../../types';
 import useLastError from '../../utils/useLastError';
 import LogsQueryEditor from '../LogsQueryEditor';
 import { Space } from '../Space';
@ -20,53 +22,84 @@ type Props = {
 };

 const VariableEditor = (props: Props) => {
-  const defaultQuery: AzureMonitorQuery = {
-    refId: 'A',
-    queryType: AzureQueryType.GrafanaTemplateVariableFn,
-  };
+  const { query, onChange, datasource } = props;
  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 });
+    AZURE_QUERY_VARIABLE_TYPE_OPTIONS.push({ label: 'Resource Groups', value: AzureQueryType.ResourceGroupsQuery });
  }
-
-  const [query, setQuery] = useState(defaultQuery);
+  const [variableOptionGroup, setVariableOptionGroup] = useState<{ label: string; options: AzureMonitorOption[] }>({
+    label: 'Template Variables',
+    options: [],
+  });
+  const [requireSubscription, setRequireSubscription] = useState(false);
+  const [subscriptions, setSubscriptions] = useState<SelectableValue[]>([]);
+  const [errorMessage, setError] = useLastError();
+  const queryType = typeof query === 'string' ? '' : query.queryType;

  useEffect(() => {
-    migrateStringQueriesToObjectQueries(props.query, { datasource: props.datasource }).then((migratedQuery) => {
-      setQuery(migratedQuery);
+    migrateStringQueriesToObjectQueries(query, { datasource: datasource }).then((migratedQuery) => {
+      onChange(migratedQuery);
    });
-  }, [props.query, props.datasource]);
+  }, [query, datasource, onChange]);
+
+  useEffect(() => {
+    switch (queryType) {
+      case AzureQueryType.ResourceGroupsQuery:
+        setRequireSubscription(true);
+        break;
+      default:
+        setRequireSubscription(false);
+    }
+  }, [queryType]);
+
+  useEffect(() => {
+    const options: AzureMonitorOption[] = [];
+    datasource.getVariablesRaw().forEach((v) => {
+      if (get(v, 'query.queryType') !== queryType) {
+        options.push({ label: v.label || v.name, value: `$${v.name}` });
+      }
+    });
+    setVariableOptionGroup({
+      label: 'Template Variables',
+      options,
+    });
+  }, [datasource, queryType]);
+
+  useEffectOnce(() => {
+    datasource.getSubscriptions().then((subs) => {
+      setSubscriptions(subs.map((s) => ({ label: s.text, value: s.value })));
+    });
+  });
+
+  if (typeof query === 'string') {
+    // still migrating the query
+    return null;
+  }

  const onQueryTypeChange = (selectableValue: SelectableValue) => {
    if (selectableValue.value) {
-      const newQuery = {
+      onChange({
        ...query,
        queryType: selectableValue.value,
-      };
-      setQuery(newQuery);
-      props.onChange(newQuery);
+      });
+    }
+  };
+
+  const onChangeSubscription = (selectableValue: SelectableValue) => {
+    if (selectableValue.value) {
+      onChange({
+        ...query,
+        subscription: selectableValue.value,
+      });
    }
  };

  const onLogsQueryChange = (queryChange: AzureMonitorQuery) => {
-    setQuery(queryChange);
-
-    // only hit backend if there's something to query (prevents error when selecting the resource before pinging a query)
-    if (queryChange.azureLogAnalytics?.query) {
-      props.onChange(queryChange);
-    }
-  };
-
-  const [errorMessage, setError] = useLastError();
-
-  const variableOptionGroup = {
-    label: 'Template Variables',
-    // TODO: figure out a way to filter out the current variable from the variables list
-    // options: props.datasource.getVariables().map((v) => ({ label: v, value: v })),
-    options: [],
+    onChange(queryChange);
  };

  return (
@ -77,15 +110,15 @@ const VariableEditor = (props: Props) => {
          onChange={onQueryTypeChange}
          options={AZURE_QUERY_VARIABLE_TYPE_OPTIONS}
          width={25}
-          value={query.queryType}
+          value={queryType}
        />
      </InlineField>
-      {query.queryType === AzureQueryType.LogAnalytics && (
+      {typeof query === 'object' && query.queryType === AzureQueryType.LogAnalytics && (
        <>
          <LogsQueryEditor
            subscriptionId={query.subscription}
            query={query}
-            datasource={props.datasource}
+            datasource={datasource}
            onChange={onLogsQueryChange}
            variableOptionGroup={variableOptionGroup}
            setError={setError}
@ -101,8 +134,19 @@ const VariableEditor = (props: Props) => {
          )}
        </>
      )}
-      {query.queryType === AzureQueryType.GrafanaTemplateVariableFn && (
-        <GrafanaTemplateVariableFnInput query={query} updateQuery={props.onChange} datasource={props.datasource} />
+      {typeof query === 'object' && query.queryType === AzureQueryType.GrafanaTemplateVariableFn && (
+        <GrafanaTemplateVariableFnInput query={query} updateQuery={props.onChange} datasource={datasource} />
+      )}
+      {typeof query === 'object' && requireSubscription && (
+        <InlineField label="Select subscription" labelWidth={20}>
+          <Select
+            aria-label="select subscription"
+            onChange={onChangeSubscription}
+            options={subscriptions.concat(variableOptionGroup)}
+            width={25}
+            value={query.subscription}
+          />
+        </InlineField>
      )}
    </>
  );
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/datasource.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/datasource.ts
@ -190,6 +190,10 @@ export default class Datasource extends DataSourceWithBackend<AzureMonitorQuery,
  getVariables() {
    return this.templateSrv.getVariables().map((v) => `$${v.name}`);
  }
+
+  getVariablesRaw() {
+    return this.templateSrv.getVariables();
+  }
 }

 function hasQueryForType(query: AzureMonitorQuery): boolean {
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/types/query.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/types/query.ts
@ -7,6 +7,7 @@ export enum AzureQueryType {
  LogAnalytics = 'Azure Log Analytics',
  AzureResourceGraph = 'Azure Resource Graph',
  SubscriptionsQuery = 'Azure Subscriptions',
+  ResourceGroupsQuery = 'Azure Resource Groups',
  /** Deprecated */
  GrafanaTemplateVariableFn = 'Grafana Template Variable Function',
 }
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/variables.test.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/variables.test.ts
@ -514,7 +514,7 @@ describe('VariableSupport', () => {
    });
  });

-  describe('querying for subscriptions', () => {
+  describe('predefined functions', () => {
    it('can fetch subscriptions', (done) => {
      const fakeSubscriptions = ['subscriptionId'];
      const variableSupport = new VariableSupport(
@ -536,5 +536,28 @@ describe('VariableSupport', () => {
        done();
      });
    });
+
+    it('can fetch resourceGroups', (done) => {
+      const expectedResults = ['test'];
+      const variableSupport = new VariableSupport(
+        createMockDatasource({
+          getResourceGroups: jest.fn().mockResolvedValueOnce(expectedResults),
+        })
+      );
+      const mockRequest = {
+        targets: [
+          {
+            refId: 'A',
+            queryType: AzureQueryType.ResourceGroupsQuery,
+            subscription: 'sub',
+          } as AzureMonitorQuery,
+        ],
+      } as DataQueryRequest<AzureMonitorQuery>;
+      const observables = variableSupport.query(mockRequest);
+      observables.subscribe((result: DataQueryResponseData) => {
+        expect(result.data[0].source).toEqual(expectedResults);
+        done();
+      });
+    });
  });
 });
--- a/public/app/plugins/datasource/grafana-azure-monitor-datasource/variables.ts
+++ b/public/app/plugins/datasource/grafana-azure-monitor-datasource/variables.ts
@ -36,6 +36,13 @@ export class VariableSupport extends CustomVariableSupport<DataSource, AzureMoni
            return {
              data: res?.length ? [toDataFrame(res)] : [],
            };
+          case AzureQueryType.ResourceGroupsQuery:
+            if (queryObj.subscription) {
+              const rgs = await this.datasource.getResourceGroups(queryObj.subscription);
+              return {
+                data: rgs?.length ? [toDataFrame(rgs)] : [],
+              };
+            }
          case AzureQueryType.GrafanaTemplateVariableFn:
            if (queryObj.grafanaTemplateVariableFn) {
              const templateVariablesResults = await this.callGrafanaTemplateVariableFn(