From 3cb92e34603ab1ee661729984988e0cfce796bb7 Mon Sep 17 00:00:00 2001
From: ismail simsek <ismailsimsek09@gmail.com>
Date: Wed, 8 Nov 2023 15:00:13 +0100
Subject: [PATCH] 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
---
 .../editor/query/fsql/FSQLEditor.tsx          |  53 ++++----
 .../editor/variable/VariableQueryEditor.tsx   |  19 ++-
 .../datasource/influxdb/datasource.test.ts    |   2 +-
 .../plugins/datasource/influxdb/datasource.ts |  16 ++-
 .../influxdb/datasource_backend_mode.test.ts  |   2 +-
 .../influxdb/datasource_sql.test.ts           |  44 +++++++
 .../fsql/datasource.flightsql.test.ts         |  49 ++++++++
 .../influxdb/fsql/datasource.flightsql.ts     |  28 ++++-
 .../app/plugins/datasource/influxdb/mocks.ts  | 116 +++++++++++++++++-
 9 files changed, 287 insertions(+), 42 deletions(-)
 create mode 100644 public/app/plugins/datasource/influxdb/datasource_sql.test.ts
 create mode 100644 public/app/plugins/datasource/influxdb/fsql/datasource.flightsql.test.ts

diff --git a/public/app/plugins/datasource/influxdb/components/editor/query/fsql/FSQLEditor.tsx b/public/app/plugins/datasource/influxdb/components/editor/query/fsql/FSQLEditor.tsx
index cc79bdc0e44..5b33e20a8d3 100644
--- 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 {
diff --git a/public/app/plugins/datasource/influxdb/components/editor/variable/VariableQueryEditor.tsx b/public/app/plugins/datasource/influxdb/components/editor/variable/VariableQueryEditor.tsx
index 1df0f2ff78b..e771e54a982 100644
--- 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 (
diff --git a/public/app/plugins/datasource/influxdb/datasource.test.ts b/public/app/plugins/datasource/influxdb/datasource.test.ts
index 6462fbd09b2..be4d4a53c68 100644
--- 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);
diff --git a/public/app/plugins/datasource/influxdb/datasource.ts b/public/app/plugins/datasource/influxdb/datasource.ts
index e5a617a79fc..18805064d9f 100644
--- 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({
diff --git a/public/app/plugins/datasource/influxdb/datasource_backend_mode.test.ts b/public/app/plugins/datasource/influxdb/datasource_backend_mode.test.ts
index 993cdbebe8e..7665a83954d 100644
--- 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);
diff --git a/public/app/plugins/datasource/influxdb/datasource_sql.test.ts b/public/app/plugins/datasource/influxdb/datasource_sql.test.ts
new file mode 100644
index 00000000000..fa0af1ed35d
--- /dev/null
+++ 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`
+      );
+    });
+  });
+});
diff --git a/public/app/plugins/datasource/influxdb/fsql/datasource.flightsql.test.ts b/public/app/plugins/datasource/influxdb/fsql/datasource.flightsql.test.ts
new file mode 100644
index 00000000000..30b0855893a
--- /dev/null
+++ 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');
+  });
+});
diff --git a/public/app/plugins/datasource/influxdb/fsql/datasource.flightsql.ts b/public/app/plugins/datasource/influxdb/fsql/datasource.flightsql.ts
index ec47c71f0b7..179842c6892 100644
--- 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) {
diff --git a/public/app/plugins/datasource/influxdb/mocks.ts b/public/app/plugins/datasource/influxdb/mocks.ts
index af3c0184c8d..22b7b8be019 100644
--- 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']],
+            },
+          },
+        ],
+      },
+    },
+  },
+};