feat(dataframe): Column variable support

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

@@ -193,7 +193,7 @@ it('should migrate queries using columns of arrays of objects', async () => {
 
 it('attempts to replace variables in metadata query', async () => {
   const tableId = '$tableId';
-  replaceMock.mockReturnValue('1');
+  replaceMock.mockReturnValueOnce('1');
 
   await ds.getTableMetadata(tableId);
 
@@ -202,13 +202,24 @@ it('attempts to replace variables in metadata query', async () => {
 });
 
 it('attempts to replace variables in data query', async () => {
-  const query = buildQuery([{ refId: 'A', tableId: '$tableId', columns: ['float'] }]);
-  replaceMock.mockReturnValue('1');
+  const query = buildQuery([{ refId: 'A', tableId: '$tableId', columns: ['$column'] }]);
+  replaceMock.mockReturnValueOnce('1').mockReturnValueOnce('time');
 
   await ds.query(query);
 
-  expect(replaceMock).toHaveBeenCalledTimes(2);
+  expect(replaceMock).toHaveBeenCalledTimes(3);
   expect(replaceMock).toHaveBeenCalledWith(query.targets[0].tableId, expect.anything());
+  expect(replaceMock).toHaveBeenCalledWith(query.targets[0].columns![0], expect.anything());
+});
+
+it('metricFindQuery returns table columns', async () => {
+  const tableId = '12345';
+  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));
 });
 
 const buildQuery = (targets: DataFrameQuery[]): DataQueryRequest<DataFrameQuery> => {

src/datasources/data-frame/DataFrameDataSource.ts

@@ -8,7 +8,8 @@ import {
   FieldType,
   standardTransformers,
   DataFrame,
-  TimeRange
+  TimeRange,
+  MetricFindValue
 } from '@grafana/data';
 import {
   BackendSrv,
@@ -50,6 +51,7 @@ export class DataFrameDataSource extends DataSourceBase<DataFrameQuery> {
   async runQuery(query: DataFrameQuery, { range, scopedVars, maxDataPoints }: DataQueryRequest): Promise<DataFrame> {
     const processedQuery = this.processQuery(query);
     processedQuery.tableId = getTemplateSrv().replace(processedQuery.tableId, scopedVars);
+    processedQuery.columns = processedQuery.columns.map(col => getTemplateSrv().replace(col, scopedVars));
     const tableMetadata = await this.getTableMetadata(processedQuery.tableId);
     const columns = this.getColumnTypes(processedQuery.columns, tableMetadata?.columns ?? []);
     const tableData = await this.getDecimatedTableData(processedQuery, columns, range, maxDataPoints);
@@ -126,6 +128,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 ? 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,58 +1,29 @@
 import React, { useState } from 'react';
 import { useAsync } from 'react-use';
-import { CoreApp, QueryEditorProps, SelectableValue, toOption } from '@grafana/data';
-import { DataFrameDataSource } from '../DataFrameDataSource';
-import { DataFrameQuery } from '../types';
-import { InlineField, InlineSwitch, MultiSelect, Select, AsyncSelect, LoadOptionsCallback } from '@grafana/ui';
+import { SelectableValue, toOption } from '@grafana/data';
+import { InlineField, InlineSwitch, MultiSelect, Select, AsyncSelect } 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 { getWorkspaceName } from 'core/utils';
-
-type Props = QueryEditorProps<DataFrameDataSource, DataFrameQuery>;
+import { DataFrameQueryEditorCommon, Props } from './DataFrameQueryEditorCommon';
 
 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: [] }, false);
-    }
-  };
+  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' }}>
@@ -63,42 +34,42 @@ 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>
       <InlineField label="Columns" shrink={true} tooltip="Specifies the columns to include in the response data.">
         <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="Specifies the method used to decimate the data.">
         <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="Filters out null and NaN values before decimating the data.">
         <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="Queries only for data within the dashboard time range if the table index is a timestamp. Enable when interacting with your data on a graph."
       >
         <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>
       <FloatingError message={errorMsg} />

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

@@ -0,0 +1,51 @@
+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, 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: [] }, false);
+    }
+  };
+
+  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' }));
+});
+

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

@@ -0,0 +1,34 @@
+import React, { useState } from 'react';
+import { AsyncSelect } from '@grafana/ui';
+import { InlineField } from 'core/components/InlineField';
+import { toOption } from '@grafana/data';
+import { isValidId } from '../utils';
+import _ from 'lodash';
+import { FloatingError, parseErrorMessage } from '../errors';
+import { DataFrameQueryEditorCommon, Props } from './DataFrameQueryEditorCommon';
+
+export function DataFrameVariableQueryEditor(props: Props) {
+  const [errorMsg, setErrorMsg] = useState<string>('');
+  const handleError = (error: Error) => setErrorMsg(parseErrorMessage(error));
+  const common = new DataFrameQueryEditorCommon(props, handleError);
+
+  return (
+    <div style={{ position: 'relative' }}>
+      <InlineField label="Id">
+        <AsyncSelect
+          allowCreateWhileLoading
+          allowCustomValue
+          cacheOptions={false}
+          defaultOptions
+          isValidNewOption={isValidId}
+          loadOptions={common.handleLoadOptions}
+          onChange={common.handleIdChange}
+          placeholder="Search by name or enter id"
+          width={30}
+          value={common.query.tableId ? toOption(common.query.tableId) : null}
+        />
+      </InlineField>
+      <FloatingError message={errorMsg} />
+    </div>
+  );
+};

src/datasources/data-frame/module.ts

@@ -2,7 +2,9 @@ import { DataSourcePlugin } from '@grafana/data';
 import { DataFrameDataSource } from './DataFrameDataSource';
 import { DataFrameQueryEditor } from './components/DataFrameQueryEditor';
 import { HttpConfigEditor } from 'core/components/HttpConfigEditor';
+import { DataFrameVariableQueryEditor } from './components/DataFrameVariableQueryEditor';
 
 export const plugin = new DataSourcePlugin(DataFrameDataSource)
   .setConfigEditor(HttpConfigEditor)
-  .setQueryEditor(DataFrameQueryEditor);
+  .setQueryEditor(DataFrameQueryEditor)
+  .setVariableQueryEditor(DataFrameVariableQueryEditor);