feat(dataframe): Column variable support

src/core/utils.ts

@@ -2,6 +2,7 @@ import { SelectableValue } from '@grafana/data';
 import { useAsync } from 'react-use';
 import { DataSourceBase } from './DataSourceBase';
 import { Workspace } from './types';
+import { TemplateSrv } from '@grafana/runtime';
 
 export function enumToOptions<T>(stringEnum: { [name: string]: T }): Array<SelectableValue<T>> {
   const RESULT = [];
@@ -62,3 +63,22 @@ export function getWorkspaceName(workspaces: Workspace[], id: string) {
 export function sleep(timeout: number) {
   return new Promise(res => window.setTimeout(res, timeout));
 }
+
+/**
+ * Replace variables in an array of values.
+ * Useful for multi-value variables.
+ */
+export function replaceVariables(values: string[], templateSrv: TemplateSrv) {
+  const replaced: string[] = [];
+  values.forEach((col: string, index) => {
+    let value = col;
+    if (templateSrv.containsTemplate(col)) {
+      const variables = templateSrv.getVariables() as any[];
+      const variable = variables.find(v => v.name === col.split('$')[1]);
+      value = variable.current.value;
+    }
+    replaced.push(value);
+  });
+  // Dedupe and flatten
+  return [...new Set(replaced.flat())];
+}

src/datasources/data-frame/DataFrameDataSource.test.ts

@@ -8,11 +8,21 @@ import { DataFrameDataSource } from './DataFrameDataSource';
 jest.mock('@grafana/runtime', () => ({
   ...jest.requireActual('@grafana/runtime'),
   getBackendSrv: () => ({ fetch: fetchMock }),
-  getTemplateSrv: () => ({ replace: replaceMock }),
+  getTemplateSrv: () => ({ replace: replaceMock, containsTemplate: containsTemplateMock, getVariables: getVariablesMock }),
 }));
 
+const mockVariables = [{
+  name: 'tableId',
+  current: { value: '1' }
+}, {
+  name: 'columns',
+  current: { value: ['time', 'int'] }
+}];
+
 const fetchMock = jest.fn<Observable<FetchResponse>, [BackendSrvRequest]>();
 const replaceMock = jest.fn((a: string, ...rest: any) => a);
+const containsTemplateMock = jest.fn((a: string) => mockVariables.map(v => `$${v.name}`).includes(a));
+const getVariablesMock = jest.fn(() => mockVariables);
 
 let ds: DataFrameDataSource;
 
@@ -205,13 +215,25 @@ it('attempts to replace variables in metadata query', async () => {
 });
 
 it('attempts to replace variables in data query', async () => {
-  const query = buildQuery([{ refId: 'A', type: DataFrameQueryType.Data, tableId: '$tableId', columns: ['float'] }]);
+  const query = buildQuery([{ refId: 'A', type: DataFrameQueryType.Data, tableId: '$tableId', columns: ['$columns'] }]);
   replaceMock.mockReturnValueOnce('1');
 
   await ds.query(query);
 
   expect(replaceMock).toHaveBeenCalledTimes(2);
   expect(replaceMock).toHaveBeenCalledWith(query.targets[0].tableId, expect.anything());
+  expect(containsTemplateMock).toHaveBeenCalledTimes(1);
+  expect(containsTemplateMock).toHaveBeenCalledWith(query.targets[0].columns![0]);
+});
+
+it('metricFindQuery returns table columns', async () => {
+  const tableId = '1';
+  const expectedColumns = fakeMetadataResponse.columns.map(col => ({ text: col.name, value: col.name }));
+
+  const columns = await ds.metricFindQuery({ tableId } as DataFrameQuery);
+
+  expect(fetchMock).toHaveBeenCalledWith(expect.objectContaining({ url: `_/nidataframe/v1/tables/${tableId}` }));
+  expect(columns).toEqual(expect.arrayContaining(expectedColumns));
 });
 
 it('returns table properties for metadata query', async () => {

src/datasources/data-frame/DataFrameDataSource.ts

@@ -1,6 +1,6 @@
 import TTLCache from '@isaacs/ttlcache';
 import deepEqual from 'fast-deep-equal';
-import { DataQueryRequest, DataSourceInstanceSettings, FieldType, TimeRange, FieldDTO, dateTime, DataFrameDTO } from '@grafana/data';
+import { DataQueryRequest, DataSourceInstanceSettings, FieldType, TimeRange, FieldDTO, dateTime, DataFrameDTO, MetricFindValue } from '@grafana/data';
 import { BackendSrv, TemplateSrv, TestingStatus, getBackendSrv, getTemplateSrv } from '@grafana/runtime';
 import {
   ColumnDataType,
@@ -17,6 +17,7 @@ import {
 import { metadataCacheTTL } from './constants';
 import _ from 'lodash';
 import { DataSourceBase } from 'core/DataSourceBase';
+import { replaceVariables } from 'core/utils';
 
 export class DataFrameDataSource extends DataSourceBase<DataFrameQuery> {
   private readonly metadataCache: TTLCache<string, TableMetadata> = new TTLCache({ ttl: metadataCacheTTL });
@@ -35,7 +36,8 @@ export class DataFrameDataSource extends DataSourceBase<DataFrameQuery> {
 
   async runQuery(query: DataFrameQuery, { range, scopedVars, maxDataPoints }: DataQueryRequest): Promise<DataFrameDTO> {
     const processedQuery = this.processQuery(query);
-    processedQuery.tableId = getTemplateSrv().replace(processedQuery.tableId, scopedVars);
+    processedQuery.tableId = this.templateSrv.replace(processedQuery.tableId, scopedVars);
+    processedQuery.columns = replaceVariables(processedQuery.columns, this.templateSrv);
     const metadata = await this.getTableMetadata(processedQuery.tableId);
 
     if (processedQuery.type === DataFrameQueryType.Metadata) {
@@ -63,7 +65,7 @@ export class DataFrameDataSource extends DataSourceBase<DataFrameQuery> {
   }
 
   async getTableMetadata(id?: string): Promise<TableMetadata> {
-    const resolvedId = getTemplateSrv().replace(id);
+    const resolvedId = this.templateSrv.replace(id);
     let metadata = this.metadataCache.get(resolvedId);
 
     if (!metadata) {
@@ -119,6 +121,11 @@ export class DataFrameDataSource extends DataSourceBase<DataFrameQuery> {
     return deepEqual(migratedQuery, query) ? (query as ValidDataFrameQuery) : migratedQuery;
   }
 
+  async metricFindQuery(tableQuery: DataFrameQuery): Promise<MetricFindValue[]> {
+    const tableMetadata = await this.getTableMetadata(tableQuery.tableId);
+    return tableMetadata.columns.map(col => ({ text: col.name, value: col.name }));
+  }
+
   private getColumnTypes(columnNames: string[], tableMetadata: Column[]): Column[] {
     return columnNames.map(c => {
       const column = tableMetadata.find(({ name }) => name === c);

src/datasources/data-frame/components/DataFrameQueryEditor.tsx

@@ -1,83 +1,39 @@
 import React, { useState } from 'react';
 import { useAsync } from 'react-use';
-import { CoreApp, QueryEditorProps, SelectableValue, toOption } from '@grafana/data';
-import { DataFrameDataSource } from '../DataFrameDataSource';
-import { DataFrameQuery, DataFrameQueryType } from '../types';
-import {
-  InlineField,
-  InlineSwitch,
-  MultiSelect,
-  Select,
-  AsyncSelect,
-  LoadOptionsCallback,
-  RadioButtonGroup,
-} from '@grafana/ui';
+import { SelectableValue, toOption } from '@grafana/data';
+import { InlineField, InlineSwitch, MultiSelect, Select, AsyncSelect, RadioButtonGroup } from '@grafana/ui';
 import { decimationMethods } from '../constants';
 import _ from 'lodash';
 import { getTemplateSrv } from '@grafana/runtime';
 import { isValidId } from '../utils';
 import { FloatingError, parseErrorMessage } from '../errors';
-import { enumToOptions, getWorkspaceName } from 'core/utils';
-
-type Props = QueryEditorProps<DataFrameDataSource, DataFrameQuery>;
+import { DataFrameQueryEditorCommon, Props } from './DataFrameQueryEditorCommon';
+import { enumToOptions } from 'core/utils';
+import { DataFrameQueryType } from '../types';
 
 export const DataFrameQueryEditor = (props: Props) => {
-  const { datasource, onChange } = props;
-  const query = datasource.processQuery(props.query);
-  const onRunQuery = () => props.app !== CoreApp.Explore && props.onRunQuery();
-
   const [errorMsg, setErrorMsg] = useState<string>('');
   const handleError = (error: Error) => setErrorMsg(parseErrorMessage(error));
-
-  const tableMetadata = useAsync(() => datasource.getTableMetadata(query.tableId).catch(handleError), [query.tableId]);
-
-  const handleQueryChange = (value: DataFrameQuery, runQuery: boolean) => {
-    onChange(value);
-    if (runQuery) {
-      onRunQuery();
-    }
-  };
-
-  const handleIdChange = (item: SelectableValue<string>) => {
-    if (query.tableId !== item.value) {
-      handleQueryChange({ ...query, tableId: item.value, columns: [] }, query.type === DataFrameQueryType.Metadata);
-    }
-  };
+  const common = new DataFrameQueryEditorCommon(props, handleError);
+  const tableMetadata = useAsync(() => common.datasource.getTableMetadata(common.query.tableId).catch(handleError), [common.query.tableId]);
 
   const handleColumnChange = (items: Array<SelectableValue<string>>) => {
-    handleQueryChange({ ...query, columns: items.map(i => i.value!) }, false);
+    common.handleQueryChange({ ...common.query, columns: items.map(i => i.value!) }, false);
   };
 
-  const loadTableOptions = _.debounce((query: string, cb?: LoadOptionsCallback<string>) => {
-    Promise.all([datasource.queryTables(query), datasource.getWorkspaces()])
-      .then(([tables, workspaces]) =>
-        cb?.(
-          tables.map(t => ({
-            label: t.name,
-            value: t.id,
-            title: t.id,
-            description: getWorkspaceName(workspaces, t.workspace),
-          }))
-        )
-      )
-      .catch(handleError);
-  }, 300);
-
-  const handleLoadOptions = (query: string, cb?: LoadOptionsCallback<string>) => {
-    if (!query || query.startsWith('$')) {
-      return cb?.(getVariableOptions().filter(v => v.value?.includes(query)));
-    }
-
-    loadTableOptions(query, cb);
-  };
+  const loadColumnOptions = () => {
+    const columnOptions = (tableMetadata.value?.columns ?? []).map(c => toOption(c.name));
+    columnOptions.unshift(...getVariableOptions());
+    return columnOptions;
+  }
 
   return (
     <div style={{ position: 'relative' }}>
       <InlineField label="Query type" tooltip={tooltips.queryType}>
         <RadioButtonGroup
           options={enumToOptions(DataFrameQueryType)}
-          value={query.type}
-          onChange={value => handleQueryChange({ ...query, type: value }, true)}
+          value={common.query.type}
+          onChange={value => common.handleQueryChange({ ...common.query, type: value }, true)}
         />
       </InlineField>
       <InlineField label="Id">
@@ -87,41 +43,41 @@ export const DataFrameQueryEditor = (props: Props) => {
           cacheOptions={false}
           defaultOptions
           isValidNewOption={isValidId}
-          loadOptions={handleLoadOptions}
-          onChange={handleIdChange}
+          loadOptions={common.handleLoadOptions}
+          onChange={common.handleIdChange}
           placeholder="Search by name or enter id"
           width={30}
-          value={query.tableId ? toOption(query.tableId) : null}
+          value={common.query.tableId ? toOption(common.query.tableId) : null}
         />
       </InlineField>
-      {query.type === DataFrameQueryType.Data && (
+      {common.query.type === DataFrameQueryType.Data && (
         <>
           <InlineField label="Columns" shrink={true} tooltip={tooltips.columns}>
             <MultiSelect
               isLoading={tableMetadata.loading}
-              options={(tableMetadata.value?.columns ?? []).map(c => toOption(c.name))}
+              options={loadColumnOptions()}
               onChange={handleColumnChange}
-              onBlur={onRunQuery}
-              value={query.columns.map(toOption)}
+              onBlur={common.onRunQuery}
+              value={common.query.columns.map(toOption)}
             />
           </InlineField>
           <InlineField label="Decimation" tooltip={tooltips.decimation}>
             <Select
               options={decimationMethods}
-              onChange={item => handleQueryChange({ ...query, decimationMethod: item.value! }, true)}
-              value={query.decimationMethod}
+              onChange={item => common.handleQueryChange({ ...common.query, decimationMethod: item.value! }, true)}
+              value={common.query.decimationMethod}
             />
           </InlineField>
           <InlineField label="Filter nulls" tooltip={tooltips.filterNulls}>
             <InlineSwitch
-              value={query.filterNulls}
-              onChange={event => handleQueryChange({ ...query, filterNulls: event.currentTarget.checked }, true)}
+              value={common.query.filterNulls}
+              onChange={event => common.handleQueryChange({ ...common.query, filterNulls: event.currentTarget.checked }, true)}
             ></InlineSwitch>
           </InlineField>
           <InlineField label="Use time range" tooltip={tooltips.useTimeRange}>
             <InlineSwitch
-              value={query.applyTimeFilters}
-              onChange={event => handleQueryChange({ ...query, applyTimeFilters: event.currentTarget.checked }, true)}
+              value={common.query.applyTimeFilters}
+              onChange={event => common.handleQueryChange({ ...common.query, applyTimeFilters: event.currentTarget.checked }, true)}
             ></InlineSwitch>
           </InlineField>
         </>

src/datasources/data-frame/components/DataFrameQueryEditorCommon.tsx

@@ -0,0 +1,60 @@
+import { QueryEditorProps, CoreApp, SelectableValue } from "@grafana/data";
+import { LoadOptionsCallback } from "@grafana/ui";
+import { getWorkspaceName, getVariableOptions } from "core/utils";
+import _ from "lodash";
+import { DataFrameDataSource } from "../DataFrameDataSource";
+import { DataFrameQuery, DataFrameQueryType, ValidDataFrameQuery } from "../types";
+
+export type Props = QueryEditorProps<DataFrameDataSource, DataFrameQuery>;
+
+export class DataFrameQueryEditorCommon {
+  readonly datasource: DataFrameDataSource;
+  readonly onChange: (value: DataFrameQuery) => void;
+  readonly query: ValidDataFrameQuery;
+  readonly onRunQuery: () => false | void;
+  readonly handleError: (error: Error) => void;
+
+  constructor(readonly props: Props, readonly errorHandler: (error: Error) => void) {
+    this.datasource = props.datasource;
+    this.onChange = props.onChange;
+    this.query = this.datasource.processQuery(props.query);
+    this.onRunQuery = () => props.app !== CoreApp.Explore && props.onRunQuery();
+    this.handleError = errorHandler;
+  }
+
+  readonly handleQueryChange = (value: DataFrameQuery, runQuery: boolean) => {
+    this.onChange(value);
+    if (runQuery) {
+      this.onRunQuery();
+    }
+  };
+
+  readonly handleIdChange = (item: SelectableValue<string>) => {
+    if (this.query.tableId !== item.value) {
+      this.handleQueryChange({ ...this.query, tableId: item.value, columns: [] }, this.query.type === DataFrameQueryType.Metadata);
+    }
+  };
+
+  readonly loadTableOptions = _.debounce((query: string, cb?: LoadOptionsCallback<string>) => {
+    Promise.all([this.datasource.queryTables(query), this.datasource.getWorkspaces()])
+      .then(([tables, workspaces]) =>
+      cb?.(
+        tables.map(t => ({
+          label: t.name,
+          value: t.id,
+          title: t.id,
+          description: getWorkspaceName(workspaces, t.workspace),
+        }))
+      )
+    )
+    .catch(this.handleError);
+  }, 300);
+
+  readonly handleLoadOptions = (query: string, cb?: LoadOptionsCallback<string>) => {
+    if (!query || query.startsWith('$')) {
+      return cb?.(getVariableOptions(this.datasource).filter((v) => v.value?.includes(query)));
+    }
+
+    this.loadTableOptions(query, cb);
+  };
+}

src/datasources/data-frame/components/DataFrameVariableQueryEditor.test.tsx

@@ -0,0 +1,31 @@
+import { render, screen } from '@testing-library/react';
+import React from 'react';
+import { select } from 'react-select-event';
+import { setupDataSource } from 'test/fixtures';
+import { DataFrameDataSource } from '../DataFrameDataSource';
+import { DataFrameVariableQueryEditor } from './DataFrameVariableQueryEditor';
+import { DataFrameQuery } from '../types';
+
+const onChange = jest.fn();
+const onRunQuery = jest.fn();
+const [datasource] = setupDataSource(DataFrameDataSource);
+
+test('renders with no data table selected', async () => {
+  render(<DataFrameVariableQueryEditor {...{ onChange, onRunQuery, datasource, query: '' as unknown as  DataFrameQuery }} />);
+
+  expect(screen.getByRole('combobox')).toHaveAccessibleDescription('Search by name or enter id');
+});
+
+test('populates data table drop-down with variables', async () => {
+  render(<DataFrameVariableQueryEditor {...{ onChange, onRunQuery, datasource, query: { tableId: '$test_var' } as DataFrameQuery }} />);
+
+  expect(screen.getByText('$test_var')).toBeInTheDocument();
+});
+
+test('user selects new data table', async () => {
+  render(<DataFrameVariableQueryEditor {...{ onChange, onRunQuery, datasource, query: '' as unknown as DataFrameQuery }} />);
+
+  await select(screen.getByRole('combobox'), '$test_var', { container: document.body });
+  expect(onChange).toHaveBeenCalledWith(expect.objectContaining({ tableId: '$test_var' }));
+});
+