Add AdHoc Filters

package.json

@@ -14,14 +14,15 @@
   "license": "Apache-2.0",
   "devDependencies": {
     "@grafana/data": "latest",
-    "@grafana/toolkit": "latest",
     "@grafana/runtime": "latest",
+    "@grafana/toolkit": "latest",
     "@grafana/ui": "latest",
     "@types/lodash": "latest"
   },
   "engines": {
     "node": ">=14"
   },
   "dependencies": {
+    "pgsql-ast-parser": "^11.1.0"
   }
 }

src/adHocFilter.ts

@@ -0,0 +1,51 @@
+import { getTable } from "ast";
+
+export default class AdHocFilter {
+  private _targetTable = '';
+
+  setTargetTable(table: string) {
+    console.log('adHocFilter: setTargetTable called')
+    this._targetTable = table;
+  }
+
+  setTargetTableFromQuery(query: string) {
+    console.log('adHocFilter: setTargetTableFromQuery called');
+    this._targetTable = getTable(query);
+    if (this._targetTable === '') {
+      console.error('Failed to get table from adhoc query.');
+      throw new Error('Failed to get table from adhoc query.');
+    }
+  }
+
+  apply(sql: string, adHocFilters: AdHocVariableFilter[]): string {
+    console.log('adHocFilter: apply called')
+    if (sql === '' || !adHocFilters || adHocFilters.length === 0) {
+      console.log('return \''+sql+'\' from adHocFilter empty check 1');
+      return sql;
+    }
+    const filter = adHocFilters[0];
+    if (filter.key.includes('.')) {
+      this._targetTable = filter.key.split('.')[0];
+    }
+    if (this._targetTable === '' || !sql.match(new RegExp(`.*\\b${this._targetTable}\\b.*`, 'gi'))) {
+      console.log('return \''+sql+'\' from adHocFilter emtpy check 2');
+      return sql;
+    }
+    let filters = adHocFilters.map((f, i) => {
+      const key = f.key.includes('.') ? f.key.split('.')[1] : f.key;
+      const value = isNaN(Number(f.value)) ? `\\'${f.value}\\'` : Number(f.value);
+      const condition = i !== adHocFilters.length - 1 ? (f.condition ? f.condition : 'AND') : '';
+      return ` ${key} ${f.operator} ${value} ${condition}`;
+    }).join('');
+    sql = sql.replace(';', '');
+    console.log('return \''+`${sql} settings additional_table_filters={'${this._targetTable}' : '${filters}'}`+'\' from adHocFilter');
+    return `${sql} settings additional_table_filters={'${this._targetTable}' : '${filters}'}`;
+  }
+}
+
+export type AdHocVariableFilter = {
+  key: string;
+  operator: string;
+  value: string;
+  condition: string;
+};

src/ast.ts

@@ -0,0 +1,94 @@
+import { parseFirst, Statement, SelectFromStatement, astMapper, toSql } from 'pgsql-ast-parser';
+
+export function sqlToStatement(sql: string): Statement {
+  console.log('ast: sqlToStatement called')
+  const replaceFuncs: Array<{
+    startIndex: number;
+    name: string;
+    replacementName: string;
+  }> = [];
+  const re = /(\$__|\$|default)/gi;
+  let regExpArray: RegExpExecArray | null;
+  while ((regExpArray = re.exec(sql)) !== null) {
+    replaceFuncs.push({ startIndex: regExpArray.index, name: regExpArray[0], replacementName: '' });
+  }
+
+  //need to process in reverse so starting positions aren't effected by replacing other things
+  for (let i = replaceFuncs.length - 1; i >= 0; i--) {
+    const si = replaceFuncs[i].startIndex;
+    const replacementName = 'f' + (Math.random() + 1).toString(36).substring(7);
+    replaceFuncs[i].replacementName = replacementName;
+    sql = sql.substring(0, si) + replacementName + sql.substring(si + replaceFuncs[i].name.length);
+  }
+
+  let ast: Statement;
+  try {
+    ast = parseFirst(sql);
+    console.log('parsed sql: '+sql+' into: '+ast)
+  } catch (err) {
+    console.error(`Failed to parse SQL statement into an AST: ${err}`);
+    return {} as Statement;
+  }
+
+  const mapper = astMapper((map) => ({
+    tableRef: (t) => {
+      const rfs = replaceFuncs.find((x) => x.replacementName === t.schema);
+      if (rfs) {
+        return { ...t, schema: t.schema?.replace(rfs.replacementName, rfs.name) };
+      }
+      const rft = replaceFuncs.find((x) => x.replacementName === t.name);
+      if (rft) {
+        return { ...t, name: t.name.replace(rft.replacementName, rft.name) };
+      }
+      return map.super().tableRef(t);
+    },
+    ref: (r) => {
+      const rf = replaceFuncs.find((x) => r.name.startsWith(x.replacementName));
+      if (rf) {
+        const d = r.name.replace(rf.replacementName, rf.name);
+        return { ...r, name: d };
+      }
+      return map.super().ref(r);
+    },
+    call: (c) => {
+      const rf = replaceFuncs.find((x) => c.function.name.startsWith(x.replacementName));
+      if (rf) {
+        return { ...c, function: { ...c.function, name: c.function.name.replace(rf.replacementName, rf.name) } };
+      }
+      return map.super().call(c);
+    },
+  }));
+  return mapper.statement(ast)!;
+}
+
+export function getTable(sql: string): string {
+  const stm = sqlToStatement(sql);
+  console.log('attempting to parse for statement: '+stm);
+  if (stm.type !== 'select' || !stm.from?.length || stm.from?.length <= 0) {
+    return '';
+  }
+  switch (stm.from![0].type) {
+    case 'table': {
+      const table = stm.from![0];
+      const tableName = `${table.name.schema ? `${table.name.schema}.` : ''}${table.name.name}`;
+      // table names are case-sensitive and pgsql parser removes casing,
+      // so we need to get the casing from the raw sql
+      const s = new RegExp(`\\b${tableName}\\b`, 'gi').exec(sql);
+      return s ? s[0] : tableName;
+    }
+    case 'statement': {
+      const table = stm.from![0];
+      return getTable(toSql.statement(table.statement));
+    }
+  }
+  return '';
+}
+
+export function getFields(sql: string): string[] {
+  console.log('ast: getFields called')
+  const stm = sqlToStatement(sql) as SelectFromStatement;
+  if (stm.type !== 'select' || !stm.columns?.length || stm.columns?.length <= 0) {
+    return [];
+  }
+  return stm.columns.map((x) => `${x.expr} as ${x.alias?.name}`);
+}

src/datasource.ts

@@ -1,25 +1,114 @@
-import { DataSourceWithBackend, getTemplateSrv } from '@grafana/runtime';
-import { DataQueryRequest, DataFrame, MetricFindValue, DataSourceInstanceSettings, ScopedVars } from '@grafana/data';
+import { DataSourceWithBackend, getTemplateSrv, TemplateSrv } from '@grafana/runtime';
+import { DataQueryRequest, DataQueryResponse, TypedVariableModel, DataFrame, MetricFindValue, DataSourceInstanceSettings, ScopedVars, DataFrameView, vectorator } from '@grafana/data';
 import { SnowflakeQuery, SnowflakeOptions } from './types';
 import { switchMap, map } from 'rxjs/operators';
 import { firstValueFrom } from 'rxjs';
+import AdHocFilter from './adHocFilter';
 
 export class DataSource extends DataSourceWithBackend<SnowflakeQuery, SnowflakeOptions> {
+  templateSrv: TemplateSrv;
+  adHocFilter: AdHocFilter;
+  skipAdHocFilter = false;
+  adHocFiltersStatus = AdHocFilterStatus.none;
+
   constructor(instanceSettings: DataSourceInstanceSettings<SnowflakeOptions>) {
     super(instanceSettings);
     this.annotations = {};
+    this.templateSrv = getTemplateSrv();
+    this.adHocFilter = new AdHocFilter();
   }
 
   applyTemplateVariables(query: SnowflakeQuery, scopedVars: ScopedVars): SnowflakeQuery {
-    query.queryText = getTemplateSrv().replace(query.queryText, scopedVars);
-    return query;
+    console.log('applyTemplateVariables with query: '+query.queryText);
+    let rawQuery = query.queryText || '';
+    const templateSrv = getTemplateSrv();
+    if(!this.skipAdHocFilter) {
+      const adHocFilters = (templateSrv as any)?.getAdhocFilters(this.name);
+      if (this.adHocFiltersStatus === AdHocFilterStatus.disabled && adHocFilters.length > 0) {
+        throw new Error(`unable to appply ad hoc filters`);
+      }
+      rawQuery = this.adHocFilter.apply(rawQuery, adHocFilters);
+    }
+    this.skipAdHocFilter = false;
+    rawQuery = this.applyConditionalAll(rawQuery, getTemplateSrv().getVariables());
+    return {
+      ...query,
+      queryText: getTemplateSrv().replace(rawQuery, scopedVars) || '',
+    };
+  }
+
+  applyConditionalAll(rawQuery: string, templateVars: TypedVariableModel[]): string {
+    console.log('applyConditionalAll with query: '+rawQuery);
+    if (!rawQuery) {
+      return rawQuery;
+    }
+    const macro = '$__conditionalAll(';
+    let macroIndex = rawQuery.lastIndexOf(macro);
+
+    while (macroIndex !== -1) {
+      const params = this.getMacroArgs(rawQuery, macroIndex + macro.length - 1);
+      if (params.length !== 2) {
+        return rawQuery;
+      }
+      const templateVar = params[1].trim();
+      const key = templateVars.find( (x) => x.name === templateVar.substring(1, templateVar.length)) as any;
+      let phrase = params[0];
+      let value = key?.current.value.toString();
+      if (value === '' || value === '$__all') {
+        phrase = '1=1';
+      }
+      rawQuery = rawQuery.replace(`${macro}${params[0]},${params[1]})`, phrase);
+      macroIndex = rawQuery.lastIndexOf(macro);
+    }
+    return rawQuery;
+  }
+
+  private getMacroArgs(query: string, argsIndex: number): string[] {
+    const args = [] as string[];
+    const re = /\(|\)|,/g;
+    let bracketCount = 0;
+    let lastArgEndIndex = 1;
+    let regExpArray: RegExpExecArray | null;
+    const argsSubstr = query.substring(argsIndex, query.length);
+    while ((regExpArray = re.exec(argsSubstr)) !== null) {
+      const foundNode = regExpArray[0];
+      if (foundNode === '(') {
+        bracketCount++;
+      } else if (foundNode === ')') {
+        bracketCount--;
+      }
+      if (foundNode === ',' && bracketCount === 1) {
+        args.push(argsSubstr.substring(lastArgEndIndex, re.lastIndex - 1));
+        lastArgEndIndex = re.lastIndex;
+      }
+      if (bracketCount === 0) {
+        args.push(argsSubstr.substring(lastArgEndIndex, re.lastIndex - 1));
+        console.log('args: '+args)
+        return args;
+      }
+    }
+    return [];
   }
 
   filterQuery(query: SnowflakeQuery): boolean {
+    console.log('filterQuery called with: '+query.queryText+' and result: '+(query.queryText !== '' && !query.hide))
     return query.queryText !== '' && !query.hide;
   }
 
+  runQuery(request: Partial<SnowflakeQuery>): Promise<DataFrame> {
+    console.log('runQuery called')
+    return new Promise( (resolve) => {
+      const req = {
+        targets: [{ ...request, refId: String(Math.random()) }]
+      } as DataQueryRequest<SnowflakeQuery>;
+      this.query(req).subscribe((res: DataQueryResponse) => {
+        resolve(res.data[0] || { fields: [] });
+      });
+    });
+  }
+
   async metricFindQuery(queryText: string): Promise<MetricFindValue[]> {
+    console.log('metricFindQuery with query: '+queryText);
     if (!queryText) {
       return Promise.resolve([]);
     }
@@ -51,4 +140,107 @@ export class DataSource extends DataSourceWithBackend<SnowflakeQuery, SnowflakeO
         )
       ));
   }
+
+  async getTagKeys(): Promise<MetricFindValue[]> {
+    const { type, frame } = await this.fetchTags();
+    if (type === TagType.query) {
+      return frame.fields.map((f) => ({ text: f.name }));
+    }
+    const view = new DataFrameView(frame);
+    return view.map((item) => ({
+      text: `${item[2]}.${item[0]}`,
+    }));
+  }
+
+  async getTagValues({ key }: any): Promise<MetricFindValue[]> {
+    const { type } = this.getTagSource();
+    this.skipAdHocFilter = true;
+    if (type === TagType.query) {
+      return this.fetchTagValuesFromQuery(key);
+    }
+    return this.fetchTagValuesFromSchema(key);
+  }
+
+  private async fetchTagValuesFromSchema(key: string): Promise<MetricFindValue[]> {
+    const { from } = this.getTagSource();
+    const [table, col] = key.split('.');
+    const source = from?.includes('.') ? `${from.split('.')[0]}.${table}` : table;
+    const rawSql = `select distinct ${col} from ${source} limit 1000`;
+    const frame = await this.runQuery({ queryText: rawSql });
+    if (frame.fields?.length === 0) {
+      return [];
+    }
+    const field = frame.fields[0];
+    // Convert to string to avoid https://github.com/grafana/grafana/issues/12209
+    return vectorator(field.values)
+      .filter((value) => value !== null)
+      .map((value) => {
+        return { text: String(value) };
+      });
+  }
+
+  private async fetchTagValuesFromQuery(key: string): Promise<MetricFindValue[]> {
+    const { frame } = await this.fetchTags();
+    const field = frame.fields.find((f) => f.name === key);
+    if (field) {
+      // Convert to string to avoid https://github.com/grafana/grafana/issues/12209
+      return vectorator(field.values)
+        .filter((value) => value !== null)
+        .map((value) => {
+          return { text: String(value) };
+        });
+    }
+    return [];
+  }
+
+  async fetchTags(): Promise<Tags> {
+    const tagSource = this.getTagSource();
+    this.skipAdHocFilter = true;
+
+    if (tagSource.source === undefined) {
+      const sql = 'SELECT COLUMN_NAME, DATA_TYPE, TABLE_NAME FROM information_schema.columns';
+      const results = await this.runQuery({ queryText: sql });
+      return { type: TagType.schema, frame: results };
+    }
+
+    if (tagSource.type === TagType.query) {
+      this.adHocFilter.setTargetTableFromQuery(tagSource.source);
+    } else {
+      let table = tagSource.from;
+      this.adHocFilter.setTargetTable(table || '');
+    }
+
+    const results = await this.runQuery({ queryText: tagSource.source });
+    return { type: tagSource.type, frame: results }
+  }
+
+  private getTagSource() {
+    const ADHOC_VAR = '$snowflake_adhoc_query';
+    let source = getTemplateSrv().replace(ADHOC_VAR);
+    console.log('called getTagSource: '+source)
+    if (source === ADHOC_VAR) {
+      return { type: TagType.schema, source: undefined };
+    }
+    if (source.toLowerCase().startsWith('select')) {
+      return { type: TagType.query, source };
+    }
+    const sql = `SELECT COLUMN_NAME, DATA_TYPE, TABLE_NAME FROM information_schema.columns WHERE table_name='${source}'`
+    return { type: TagType.schema, source: sql, from: source }
+  }
+}
+
+enum TagType {
+  query,
+  schema,
+}
+
+enum AdHocFilterStatus {
+  none = 0,
+  enabled,
+  disabled,
+}
+
+interface Tags {
+  type?: TagType;
+  frame: DataFrame;
 }