InfluxDB: Template variable support for SQL language (#77799)

* Only run through with classicQuery if the language is influxql and backend migration is disabled * Add variable editor * Simple template variable support * Show template variables in the drowdowns * better imports * unit tests * it is now 11 just because we add rawSql interpolation in datasource.ts applyVariables method * fix
2025-02-25 18:55:37 -06:00 · 2023-11-08 15:00:13 +01:00
parent a46ff31c4f
commit 3cb92e3460
9 changed files with 287 additions and 42 deletions
--- a/public/app/plugins/datasource/influxdb/components/editor/query/fsql/FSQLEditor.tsx
+++ b/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 {
--- a/public/app/plugins/datasource/influxdb/components/editor/variable/VariableQueryEditor.tsx
+++ b/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 (
--- a/public/app/plugins/datasource/influxdb/datasource.test.ts
+++ b/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);
--- a/public/app/plugins/datasource/influxdb/datasource.ts
+++ b/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';
@@ -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({
--- a/public/app/plugins/datasource/influxdb/datasource_backend_mode.test.ts
+++ b/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);
--- a/public/app/plugins/datasource/influxdb/datasource_sql.test.ts
+++ b/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`
+      );
+    });
+  });
+});
--- a/public/app/plugins/datasource/influxdb/fsql/datasource.flightsql.test.ts
+++ b/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');
+  });
+});
--- a/public/app/plugins/datasource/influxdb/fsql/datasource.flightsql.ts
+++ b/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) {
--- a/public/app/plugins/datasource/influxdb/mocks.ts
+++ b/public/app/plugins/datasource/influxdb/mocks.ts
@@ -17,6 +17,7 @@ import {
  VariableInterpolation,
 } from '@grafana/runtime/src';

+import { SQLQuery } from '../../../features/plugins/sql';
 import { TemplateSrv } from '../../../features/templating/template_srv';

 import InfluxDatasource from './datasource';
@@ -229,7 +230,9 @@ export const mockInfluxRetentionPolicyResponse = [
  },
 ];

-export const mockInfluxQueryRequest = (targets?: InfluxQuery[]): DataQueryRequest<InfluxQuery> => {
+type QueryType = InfluxQuery & SQLQuery;
+
+export const mockInfluxQueryRequest = (targets?: QueryType[]): DataQueryRequest<QueryType> => {
  return {
    app: 'explore',
    interval: '1m',
@@ -251,7 +254,7 @@ export const mockInfluxQueryRequest = (targets?: InfluxQuery[]): DataQueryReques
  };
 };

-export const mockTargets = (): InfluxQuery[] => {
+export const mockTargets = (): QueryType[] => {
  return [
    {
      refId: 'A',
@@ -321,3 +324,112 @@ export const mockInfluxQueryWithTemplateVars = (adhocFilters: AdHocVariableFilte
  ],
  adhocFilters,
 });
+
+export const mockInfluxSQLFetchResponse: FetchResponse<BackendDataSourceResponse> = {
+  config: {
+    url: 'mock-response-url',
+  },
+  headers: new Headers(),
+  ok: false,
+  redirected: false,
+  status: 0,
+  statusText: '',
+  type: 'basic',
+  url: '',
+  data: {
+    results: {
+      A: {
+        status: 200,
+        frames: [
+          {
+            schema: {
+              refId: 'A',
+              meta: {
+                typeVersion: [0, 0],
+                custom: {
+                  headers: {
+                    'content-type': ['application/grpc'],
+                    date: ['Tue, 07 Nov 2023 21:18:27 GMT'],
+                    'strict-transport-security': ['max-age=15724800; includeSubDomains'],
+                    'trace-id': ['05b4f1f285b4bbe2'],
+                    'trace-sampled': ['false'],
+                    'x-envoy-upstream-service-time': ['15'],
+                  },
+                },
+                executedQueryString:
+                  'SELECT "usage_idle", time FROM iox.cpu WHERE time \u003e= cast(\'2023-11-07T21:13:27Z\' as timestamp) ',
+              },
+              fields: [
+                {
+                  name: 'usage_idle',
+                  type: FieldType.number,
+                },
+                {
+                  name: 'time',
+                  type: FieldType.time,
+                },
+              ],
+            },
+            data: {
+              values: [
+                [99.09629480869259, 99.0866204958598, 99.24736578023098, 99.24736578023054, 99.11619965852707],
+                [1699391610000, 1699391620000, 1699391630000, 1699391640000, 1699391650000],
+              ],
+            },
+          },
+        ],
+      },
+    },
+  },
+};
+
+export const mockInfluxSQLVariableFetchResponse: FetchResponse<BackendDataSourceResponse> = {
+  config: {
+    url: 'mock-response-url',
+  },
+  headers: new Headers(),
+  ok: false,
+  redirected: false,
+  status: 0,
+  statusText: '',
+  type: 'basic',
+  url: '',
+  data: {
+    results: {
+      metricFindQuery: {
+        status: 200,
+        frames: [
+          {
+            schema: {
+              refId: 'metricFindQuery',
+              meta: {
+                typeVersion: [0, 0],
+                custom: {
+                  headers: {
+                    'content-type': ['application/grpc'],
+                    date: ['Tue, 07 Nov 2023 22:19:44 GMT'],
+                    'strict-transport-security': ['max-age=15724800; includeSubDomains'],
+                    'trace-id': ['481a45f6066c0a45'],
+                    'trace-sampled': ['false'],
+                    'x-envoy-upstream-service-time': ['8'],
+                  },
+                },
+                executedQueryString:
+                  "SELECT table_name FROM information_schema.tables WHERE table_schema = 'iox' ORDER BY table_name",
+              },
+              fields: [
+                {
+                  name: 'table_name',
+                  type: FieldType.string,
+                },
+              ],
+            },
+            data: {
+              values: [['airSensors', 'cpu', 'disk', 'diskio', 'kernel', 'mem', 'processes', 'swap', 'system']],
+            },
+          },
+        ],
+      },
+    },
+  },
+};