InfluxDB: Template variable support for SQL language

public/app/plugins/datasource/influxdb/components/editor/query/fsql/FSQLEditor.tsx

@@ -25,33 +25,36 @@ class UnthemedSQLQueryEditor extends PureComponent<Props> {
     super(props);
     const { datasource: influxDatasource } = props;
 
-    this.datasource = new FlightSQLDatasource({
-      url: influxDatasource.urls[0],
-      access: influxDatasource.access,
-      id: influxDatasource.id,
-
-      jsonData: {
-        // Not applicable to flightSQL? @itsmylife
-        allowCleartextPasswords: false,
-        tlsAuth: false,
-        tlsAuthWithCACert: false,
-        tlsSkipVerify: false,
-        maxIdleConns: 1,
-        maxOpenConns: 1,
-        maxIdleConnsAuto: true,
-        connMaxLifetime: 1,
-        timezone: '',
-        user: '',
-        database: '',
+    this.datasource = new FlightSQLDatasource(
+      {
         url: influxDatasource.urls[0],
-        timeInterval: '',
+        access: influxDatasource.access,
+        id: influxDatasource.id,
+
+        jsonData: {
+          // TODO Clean this
+          allowCleartextPasswords: false,
+          tlsAuth: false,
+          tlsAuthWithCACert: false,
+          tlsSkipVerify: false,
+          maxIdleConns: 1,
+          maxOpenConns: 1,
+          maxIdleConnsAuto: true,
+          connMaxLifetime: 1,
+          timezone: '',
+          user: '',
+          database: '',
+          url: influxDatasource.urls[0],
+          timeInterval: '',
+        },
+        meta: influxDatasource.meta,
+        name: influxDatasource.name,
+        readOnly: false,
+        type: influxDatasource.type,
+        uid: influxDatasource.uid,
       },
-      meta: influxDatasource.meta,
-      name: influxDatasource.name,
-      readOnly: false,
-      type: influxDatasource.type,
-      uid: influxDatasource.uid,
-    });
+      influxDatasource.templateSrv
+    );
   }
 
   transformQuery(query: InfluxQuery & SQLQuery): SQLQuery {

public/app/plugins/datasource/influxdb/components/editor/variable/VariableQueryEditor.tsx

@@ -1,6 +1,6 @@
 import React, { PureComponent } from 'react';
 
-import { InlineFormLabel, TextArea } from '@grafana/ui/src';
+import { Field, FieldSet, InlineFormLabel, TextArea } from '@grafana/ui';
 
 import InfluxDatasource from '../../../datasource';
 import { InfluxVersion } from '../../../types';
@@ -33,11 +33,20 @@ export default class VariableQueryEditor extends PureComponent<Props> {
             onChange={(v) => onChange(v.query)}
           />
         );
-      //@todo add support for SQL
       case InfluxVersion.SQL:
-        return <div className="gf-form-inline">TODO</div>;
-
-      // Influx/default case
+        return (
+          <FieldSet>
+            <Field htmlFor="influx-sql-variable-query">
+              <TextArea
+                id="influx-sql-variable-query"
+                defaultValue={query || ''}
+                placeholder="metric name or tags query"
+                rows={1}
+                onBlur={(e) => onChange(e.currentTarget.value)}
+              />
+            </Field>
+          </FieldSet>
+        );
       case InfluxVersion.InfluxQL:
       default:
         return (

public/app/plugins/datasource/influxdb/datasource.test.ts

@@ -285,7 +285,7 @@ describe('InfluxDataSource Frontend Mode', () => {
     const ds = new InfluxDatasource(getMockDSInstanceSettings(), templateSrv);
 
     function influxChecks(query: InfluxQuery) {
-      expect(templateSrv.replace).toBeCalledTimes(10);
+      expect(templateSrv.replace).toBeCalledTimes(11);
       expect(query.alias).toBe(text);
       expect(query.measurement).toBe(textWithFormatRegex);
       expect(query.policy).toBe(textWithFormatRegex);

public/app/plugins/datasource/influxdb/datasource.ts

@@ -35,6 +35,8 @@ import { CustomFormatterVariable } from '@grafana/scenes';
 import config from 'app/core/config';
 import { getTemplateSrv, TemplateSrv } from 'app/features/templating/template_srv';
 
+import { QueryFormat, SQLQuery } from '../../../features/plugins/sql';
+
 import { AnnotationEditor } from './components/editor/annotation/AnnotationEditor';
 import { FluxQueryEditor } from './components/editor/query/flux/FluxQueryEditor';
 import { BROWSER_MODE_DISABLED_MESSAGE } from './constants';
@@ -138,12 +140,12 @@ export default class InfluxDatasource extends DataSourceWithBackend<InfluxQuery,
       return merge(...streams);
     }
 
-    if (this.isMigrationToggleOnAndIsAccessProxy()) {
-      return super.query(filteredRequest);
+    if (this.version === InfluxVersion.InfluxQL && !this.isMigrationToggleOnAndIsAccessProxy()) {
+      // Fallback to classic query support
+      return this.classicQuery(request);
     }
 
-    // Fallback to classic query support
-    return this.classicQuery(request);
+    return super.query(filteredRequest);
   }
 
   getQueryDisplayText(query: InfluxQuery) {
@@ -169,7 +171,7 @@ export default class InfluxDatasource extends DataSourceWithBackend<InfluxQuery,
     return true;
   }
 
-  applyTemplateVariables(query: InfluxQuery, scopedVars: ScopedVars): InfluxQuery {
+  applyTemplateVariables(query: InfluxQuery, scopedVars: ScopedVars): InfluxQuery & SQLQuery {
     // We want to interpolate these variables on backend
     const { __interval, __interval_ms, ...rest } = scopedVars || {};
 
@@ -224,7 +226,7 @@ export default class InfluxDatasource extends DataSourceWithBackend<InfluxQuery,
     });
   }
 
-  applyVariables(query: InfluxQuery, scopedVars: ScopedVars) {
+  applyVariables(query: InfluxQuery & SQLQuery, scopedVars: ScopedVars) {
     const expandedQuery = { ...query };
     if (query.groupBy) {
       expandedQuery.groupBy = query.groupBy.map((groupBy) => {
@@ -263,6 +265,7 @@ export default class InfluxDatasource extends DataSourceWithBackend<InfluxQuery,
       ...expandedQuery,
       adhocFilters: this.templateSrv.getAdhocFilters(this.name) ?? [],
       query: this.templateSrv.replace(query.query ?? '', scopedVars, this.interpolateQueryExpr), // The raw query text
+      rawSql: this.templateSrv.replace(query.rawSql ?? '', scopedVars, this.interpolateQueryExpr), // The raw query text
       alias: this.templateSrv.replace(query.alias ?? '', scopedVars),
       limit: this.templateSrv.replace(query.limit?.toString() ?? '', scopedVars, this.interpolateQueryExpr),
       measurement: this.templateSrv.replace(query.measurement ?? '', scopedVars, this.interpolateQueryExpr),
@@ -300,11 +303,16 @@ export default class InfluxDatasource extends DataSourceWithBackend<InfluxQuery,
   }
 
   async metricFindQuery(query: string, options?: any): Promise<MetricFindValue[]> {
-    if (this.version === InfluxVersion.Flux || this.isMigrationToggleOnAndIsAccessProxy()) {
-      const target: InfluxQuery = {
+    if (
+      this.version === InfluxVersion.Flux ||
+      this.version === InfluxVersion.SQL ||
+      this.isMigrationToggleOnAndIsAccessProxy()
+    ) {
+      const target: InfluxQuery & SQLQuery = {
         refId: 'metricFindQuery',
         query,
         rawQuery: true,
+        ...(this.version === InfluxVersion.SQL ? { rawSql: query, format: QueryFormat.Table } : {}),
       };
       return lastValueFrom(
         super.query({

public/app/plugins/datasource/influxdb/datasource_backend_mode.test.ts

@@ -176,7 +176,7 @@ describe('InfluxDataSource Backend Mode', () => {
     const ds = new InfluxDatasource(getMockDSInstanceSettings(), templateSrv);
 
     function influxChecks(query: InfluxQuery) {
-      expect(templateSrv.replace).toBeCalledTimes(10);
+      expect(templateSrv.replace).toBeCalledTimes(11);
       expect(query.alias).toBe(text);
       expect(query.measurement).toBe(textWithFormatRegex);
       expect(query.policy).toBe(textWithFormatRegex);

public/app/plugins/datasource/influxdb/datasource_sql.test.ts

@@ -0,0 +1,44 @@
+import { lastValueFrom } from 'rxjs';
+
+import config from 'app/core/config';
+
+import { SQLQuery } from '../../../features/plugins/sql';
+
+import InfluxDatasource from './datasource';
+import {
+  getMockDSInstanceSettings,
+  mockBackendService,
+  mockInfluxQueryRequest,
+  mockInfluxSQLFetchResponse,
+  mockTemplateSrv,
+} from './mocks';
+import { InfluxVersion } from './types';
+
+config.featureToggles.influxdbBackendMigration = true;
+mockBackendService(mockInfluxSQLFetchResponse);
+
+describe('InfluxDB SQL Support', () => {
+  const replaceMock = jest.fn();
+  const templateSrv = mockTemplateSrv(jest.fn(), replaceMock);
+
+  let sqlQuery: SQLQuery;
+
+  beforeEach(() => {
+    sqlQuery = {
+      refId: 'x',
+      rawSql:
+        'SELECT "$interpolationVar2", time FROM iox.$interpolationVar WHERE time >= $__timeFrom AND time <= $__timeTo',
+    };
+  });
+
+  describe('interpolate variables', () => {
+    const ds = new InfluxDatasource(getMockDSInstanceSettings({ version: InfluxVersion.SQL }), templateSrv);
+
+    it('should call replace template variables for rawSql', async () => {
+      await lastValueFrom(ds.query(mockInfluxQueryRequest([sqlQuery])));
+      expect(replaceMock.mock.calls[1][0]).toBe(
+        `SELECT "$interpolationVar2", time FROM iox.$interpolationVar WHERE time >= $__timeFrom AND time <= $__timeTo`
+      );
+    });
+  });
+});

public/app/plugins/datasource/influxdb/fsql/datasource.flightsql.test.ts

@@ -0,0 +1,49 @@
+import { TemplateSrv } from '@grafana/runtime';
+
+import { getMockDSInstanceSettings, mockBackendService, mockInfluxSQLVariableFetchResponse } from '../mocks';
+
+import { FlightSQLDatasource } from './datasource.flightsql';
+
+mockBackendService(mockInfluxSQLVariableFetchResponse);
+describe('flightsql datasource', () => {
+  const templateSrv: TemplateSrv = {
+    containsTemplate: jest.fn(),
+    replace: jest.fn().mockImplementation((val: string) => val),
+    updateTimeRange: jest.fn(),
+    getVariables: jest.fn().mockReturnValue([
+      {
+        name: 'templateVar',
+        text: 'templateVar',
+        value: 'templateVar',
+        type: '',
+        label: 'templateVar',
+      },
+    ]),
+  };
+  const mockInstanceSettings = getMockDSInstanceSettings();
+  const instanceSettings = {
+    ...mockInstanceSettings,
+    jsonData: {
+      ...mockInstanceSettings.jsonData,
+      allowCleartextPasswords: false,
+      tlsAuth: false,
+      tlsAuthWithCACert: false,
+      tlsSkipVerify: false,
+      maxIdleConns: 1,
+      maxOpenConns: 1,
+      maxIdleConnsAuto: true,
+      connMaxLifetime: 1,
+      timezone: '',
+      user: '',
+      database: '',
+      url: '',
+      timeInterval: '',
+    },
+  };
+  const ds = new FlightSQLDatasource(instanceSettings, templateSrv);
+
+  it('should add template variables to the responses', async () => {
+    const fields = await ds.fetchFields({ dataset: 'test', table: 'table' });
+    expect(fields[0].name).toBe('$templateVar');
+  });
+});

public/app/plugins/datasource/influxdb/fsql/datasource.flightsql.ts

@@ -1,5 +1,6 @@
-import { DataSourceInstanceSettings, TimeRange } from '@grafana/data/src';
+import { DataSourceInstanceSettings, TimeRange } from '@grafana/data';
 import { CompletionItemKind, LanguageDefinition, TableIdentifier } from '@grafana/experimental';
+import { getTemplateSrv, TemplateSrv } from '@grafana/runtime';
 import { SqlDatasource } from 'app/features/plugins/sql/datasource/SqlDatasource';
 import { DB, SQLQuery } from 'app/features/plugins/sql/types';
 import { formatSQL } from 'app/features/plugins/sql/utils/formatSQL';
@@ -13,7 +14,10 @@ import { FlightSQLOptions } from './types';
 export class FlightSQLDatasource extends SqlDatasource {
   sqlLanguageDefinition: LanguageDefinition | undefined;
 
-  constructor(private instanceSettings: DataSourceInstanceSettings<FlightSQLOptions>) {
+  constructor(
+    private instanceSettings: DataSourceInstanceSettings<FlightSQLOptions>,
+    protected readonly templateSrv: TemplateSrv = getTemplateSrv()
+  ) {
     super(instanceSettings);
   }
 
@@ -45,14 +49,17 @@ export class FlightSQLDatasource extends SqlDatasource {
   async fetchTables(dataset?: string): Promise<string[]> {
     const query = buildTableQuery(dataset);
     const tables = await this.runSql<string[]>(query, { refId: 'tables' });
-    return tables.map((t) => quoteIdentifierIfNecessary(t[0]));
+    const tableNames = tables.map((t) => quoteIdentifierIfNecessary(t[0]));
+    tableNames.unshift(...this.getTemplateVariables());
+    return tableNames;
   }
 
   async fetchFields(query: Partial<SQLQuery>) {
     if (!query.dataset || !query.table) {
       return [];
     }
-    const queryString = buildColumnQuery(query.table, query.dataset);
+    const interpolatedTable = this.templateSrv.replace(query.table);
+    const queryString = buildColumnQuery(interpolatedTable, query.dataset);
     const frame = await this.runSql<string[]>(queryString, { refId: 'fields' });
     const fields = frame.map((f) => ({
       name: f[0],
@@ -61,9 +68,22 @@ export class FlightSQLDatasource extends SqlDatasource {
       type: f[1],
       label: f[0],
     }));
+    fields.unshift(
+      ...this.getTemplateVariables().map((v) => ({
+        name: v,
+        text: v,
+        value: quoteIdentifierIfNecessary(v),
+        type: '',
+        label: v,
+      }))
+    );
     return mapFieldsToTypes(fields);
   }
 
+  getTemplateVariables() {
+    return this.templateSrv.getVariables().map((v) => `$${v.name}`);
+  }
+
   async fetchMeta(identifier?: TableIdentifier) {
     const defaultDB = this.instanceSettings.jsonData.database;
     if (!identifier?.schema && defaultDB) {