Skip to content

Commit 6e17206

Browse files
Dhruv-JDhruv-J
committed
Add AdHoc Filters
This PR attemps to add AdHoc filter functionality to the snowflake datasource plugin for grafana. The AST is inspired from the AST used with ClickHouse datasource plugin for grafana. Fixes issue michelin#24 Signed-off-by: Dhruv-J <[email protected]>
1 parent 300a6a3 commit 6e17206

File tree

6 files changed

+1621
-1438
lines changed

6 files changed

+1621
-1438
lines changed

‎.DS_Store

4 KB
Binary file not shown.

‎package.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,15 @@
1414
"license": "Apache-2.0",
1515
"devDependencies": {
1616
"@grafana/data": "latest",
17-
"@grafana/toolkit": "latest",
1817
"@grafana/runtime": "latest",
18+
"@grafana/toolkit": "latest",
1919
"@grafana/ui": "latest",
2020
"@types/lodash": "latest"
2121
},
2222
"engines": {
2323
"node": ">=14"
2424
},
2525
"dependencies": {
26+
"pgsql-ast-parser": "^11.1.0"
2627
}
2728
}

‎src/adHocFilter.ts

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
import { getTable } from "ast";
2+
3+
export default class AdHocFilter {
4+
private _targetTable = '';
5+
6+
setTargetTable(table: string) {
7+
this._targetTable = table;
8+
}
9+
10+
setTargetTableFromQuery(query: string) {
11+
this._targetTable = getTable(query);
12+
if (this._targetTable === '') {
13+
console.error('Failed to get table from adhoc query.');
14+
throw new Error('Failed to get table from adhoc query.');
15+
}
16+
}
17+
18+
apply(sql: string, adHocFilters: AdHocVariableFilter[]): string {
19+
if (this._targetTable === '' || sql === '' || !adHocFilters || adHocFilters.length === 0) {
20+
return sql;
21+
}
22+
const filter = adHocFilters[0];
23+
if (filter.key.includes('.')) {
24+
this._targetTable = filter.key.split('.')[0];
25+
}
26+
if (this._targetTable === '' || !sql.match(new RegExp(`.*\\b${this._targetTable}\\b.*`, 'gi'))) {
27+
return sql;
28+
}
29+
let filters = adHocFilters.map((f, i) => {
30+
const key = f.key.includes('.') ? f.key.split('.')[1] : f.key;
31+
const value = isNaN(Number(f.value)) ? `\\'${f.value}\\'` : Number(f.value);
32+
const condition = i !== adHocFilters.length - 1 ? (f.condition ? f.condition : 'AND') : '';
33+
return ` ${key} ${f.operator} ${value} ${condition}`;
34+
}).join('');
35+
sql = sql.replace(';', '');
36+
return `${sql} settings additional_table_filters={'${this._targetTable}' : '${filters}'}`;
37+
}
38+
}
39+
40+
export type AdHocVariableFilter = {
41+
key: string;
42+
operator: string;
43+
value: string;
44+
condition: string;
45+
};

‎src/ast.ts

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
import { parseFirst, Statement, SelectFromStatement, astMapper, toSql } from 'pgsql-ast-parser';
2+
3+
export function sqlToStatement(sql: string): Statement {
4+
const replaceFuncs: Array<{
5+
startIndex: number;
6+
name: string;
7+
replacementName: string;
8+
}> = [];
9+
const re = /(\$__|\$|default)/gi;
10+
let regExpArray: RegExpExecArray | null;
11+
while ((regExpArray = re.exec(sql)) !== null) {
12+
replaceFuncs.push({ startIndex: regExpArray.index, name: regExpArray[0], replacementName: '' });
13+
}
14+
15+
//need to process in reverse so starting positions aren't effected by replacing other things
16+
for (let i = replaceFuncs.length - 1; i >= 0; i--) {
17+
const si = replaceFuncs[i].startIndex;
18+
const replacementName = 'f' + (Math.random() + 1).toString(36).substring(7);
19+
replaceFuncs[i].replacementName = replacementName;
20+
sql = sql.substring(0, si) + replacementName + sql.substring(si + replaceFuncs[i].name.length);
21+
}
22+
23+
let ast: Statement;
24+
try {
25+
ast = parseFirst(sql);
26+
} catch (err) {
27+
console.error(`Failed to parse SQL statement into an AST: ${err}`);
28+
return {} as Statement;
29+
}
30+
31+
const mapper = astMapper((map) => ({
32+
tableRef: (t) => {
33+
const rfs = replaceFuncs.find((x) => x.replacementName === t.schema);
34+
if (rfs) {
35+
return { ...t, schema: t.schema?.replace(rfs.replacementName, rfs.name) };
36+
}
37+
const rft = replaceFuncs.find((x) => x.replacementName === t.name);
38+
if (rft) {
39+
return { ...t, name: t.name.replace(rft.replacementName, rft.name) };
40+
}
41+
return map.super().tableRef(t);
42+
},
43+
ref: (r) => {
44+
const rf = replaceFuncs.find((x) => r.name.startsWith(x.replacementName));
45+
if (rf) {
46+
const d = r.name.replace(rf.replacementName, rf.name);
47+
return { ...r, name: d };
48+
}
49+
return map.super().ref(r);
50+
},
51+
call: (c) => {
52+
const rf = replaceFuncs.find((x) => c.function.name.startsWith(x.replacementName));
53+
if (rf) {
54+
return { ...c, function: { ...c.function, name: c.function.name.replace(rf.replacementName, rf.name) } };
55+
}
56+
return map.super().call(c);
57+
},
58+
}));
59+
return mapper.statement(ast)!;
60+
}
61+
62+
export function getTable(sql: string): string {
63+
const stm = sqlToStatement(sql);
64+
if (stm.type !== 'select' || !stm.from?.length || stm.from?.length <= 0) {
65+
return '';
66+
}
67+
switch (stm.from![0].type) {
68+
case 'table': {
69+
const table = stm.from![0];
70+
const tableName = `${table.name.schema ? `${table.name.schema}.` : ''}${table.name.name}`;
71+
// table names are case-sensitive and pgsql parser removes casing,
72+
// so we need to get the casing from the raw sql
73+
const s = new RegExp(`\\b${tableName}\\b`, 'gi').exec(sql);
74+
return s ? s[0] : tableName;
75+
}
76+
case 'statement': {
77+
const table = stm.from![0];
78+
return getTable(toSql.statement(table.statement));
79+
}
80+
}
81+
return '';
82+
}
83+
84+
export function getFields(sql: string): string[] {
85+
const stm = sqlToStatement(sql) as SelectFromStatement;
86+
if (stm.type !== 'select' || !stm.columns?.length || stm.columns?.length <= 0) {
87+
return [];
88+
}
89+
return stm.columns.map((x) => `${x.expr} as ${x.alias?.name}`);
90+
}

‎src/datasource.ts

Lines changed: 128 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,111 @@
1-
import { DataSourceWithBackend, getTemplateSrv } from '@grafana/runtime';
2-
import { DataQueryRequest, DataFrame, MetricFindValue, DataSourceInstanceSettings, ScopedVars } from '@grafana/data';
1+
import { DataSourceWithBackend, getTemplateSrv, TemplateSrv } from '@grafana/runtime';
2+
import { DataQueryRequest, DataQueryResponse, TypedVariableModel, DataFrame, MetricFindValue, DataSourceInstanceSettings, ScopedVars, ArrayDataFrame, vectorator } from '@grafana/data';
33
import { SnowflakeQuery, SnowflakeOptions } from './types';
44
import { switchMap, map } from 'rxjs/operators';
55
import { firstValueFrom } from 'rxjs';
6+
import AdHocFilter from './adHocFilter';
67

78
export class DataSource extends DataSourceWithBackend<SnowflakeQuery, SnowflakeOptions> {
9+
templateSrv: TemplateSrv;
10+
adHocFilter: AdHocFilter;
11+
skipAdHocFilter = false;
12+
adHocFiltersStatus = AdHocFilterStatus.none;
13+
814
constructor(instanceSettings: DataSourceInstanceSettings<SnowflakeOptions>) {
915
super(instanceSettings);
1016
this.annotations = {};
17+
this.templateSrv = getTemplateSrv();
18+
this.adHocFilter = new AdHocFilter();
1119
}
1220

1321
applyTemplateVariables(query: SnowflakeQuery, scopedVars: ScopedVars): SnowflakeQuery {
14-
query.queryText = getTemplateSrv().replace(query.queryText, scopedVars);
15-
return query;
22+
console.log('applyTemplateVariables');
23+
let rawQuery = query.queryText || '';
24+
const templateSrv = getTemplateSrv();
25+
if(!this.skipAdHocFilter) {
26+
const adHocFilters = (templateSrv as any)?.getAdhocFilters(this.name);
27+
if (this.adHocFiltersStatus === AdHocFilterStatus.disabled && adHocFilters.length > 0) {
28+
throw new Error(`unable to appply ad hoc filters`);
29+
}
30+
rawQuery = this.adHocFilter.apply(rawQuery, adHocFilters);
31+
}
32+
this.skipAdHocFilter = false;
33+
rawQuery = this.applyConditionalAll(rawQuery, getTemplateSrv().getVariables());
34+
return {
35+
...query,
36+
queryText: getTemplateSrv().replace(rawQuery, scopedVars) || '',
37+
};
38+
}
39+
40+
applyConditionalAll(rawQuery: string, templateVars: TypedVariableModel[]): string {
41+
console.log('applyConditionalAll');
42+
if (!rawQuery) {
43+
return rawQuery;
44+
}
45+
const macro = '$__conditionalAll(';
46+
let macroIndex = rawQuery.lastIndexOf(macro);
47+
48+
while (macroIndex !== -1) {
49+
const params = this.getMacroArgs(rawQuery, macroIndex + macro.length - 1);
50+
if (params.length !== 2) {
51+
return rawQuery;
52+
}
53+
const templateVar = params[1].trim();
54+
const key = templateVars.find( (x) => x.name === templateVar.substring(1, templateVar.length)) as any;
55+
let phrase = params[0];
56+
let value = key?.current.value.toString();
57+
if (value === '' || value === '$__all') {
58+
phrase = '1=1';
59+
}
60+
rawQuery = rawQuery.replace(`${macro}${params[0]},${params[1]})`, phrase);
61+
macroIndex = rawQuery.lastIndexOf(macro);
62+
}
63+
return rawQuery;
64+
}
65+
66+
private getMacroArgs(query: string, argsIndex: number): string[] {
67+
const args = [] as string[];
68+
const re = /\(|\)|,/g;
69+
let bracketCount = 0;
70+
let lastArgEndIndex = 1;
71+
let regExpArray: RegExpExecArray | null;
72+
const argsSubstr = query.substring(argsIndex, query.length);
73+
while ((regExpArray = re.exec(argsSubstr)) !== null) {
74+
const foundNode = regExpArray[0];
75+
if (foundNode === '(') {
76+
bracketCount++;
77+
} else if (foundNode === ')') {
78+
bracketCount--;
79+
}
80+
if (foundNode === ',' && bracketCount === 1) {
81+
args.push(argsSubstr.substring(lastArgEndIndex, re.lastIndex - 1));
82+
lastArgEndIndex = re.lastIndex;
83+
}
84+
if (bracketCount === 0) {
85+
args.push(argsSubstr.substring(lastArgEndIndex, re.lastIndex - 1));
86+
return args;
87+
}
88+
}
89+
return [];
1690
}
1791

1892
filterQuery(query: SnowflakeQuery): boolean {
1993
return query.queryText !== '' && !query.hide;
2094
}
2195

96+
runQuery(request: Partial<SnowflakeQuery>): Promise<DataFrame> {
97+
return new Promise( (resolve) => {
98+
const req = {
99+
targets: [{ ...request, refId: String(Math.random()) }]
100+
} as DataQueryRequest<SnowflakeQuery>;
101+
this.query(req).subscribe((res: DataQueryResponse) => {
102+
resolve(res.data[0] || { fields: [] });
103+
});
104+
});
105+
}
106+
22107
async metricFindQuery(queryText: string): Promise<MetricFindValue[]> {
108+
console.log('metricFindQuery');
23109
if (!queryText) {
24110
return Promise.resolve([]);
25111
}
@@ -51,4 +137,42 @@ export class DataSource extends DataSourceWithBackend<SnowflakeQuery, SnowflakeO
51137
)
52138
));
53139
}
140+
141+
async getTagKeys(options?: any): Promise<MetricFindValue[]> {
142+
const frame = await this.fetchTags();
143+
return frame.fields.map( (f) => ({ text: f.name }) );
144+
}
145+
146+
async getTagValues({ key }: any): Promise<MetricFindValue[]> {
147+
const frame = await this.fetchTags();
148+
const field = frame.fields.find( (f) => f.name === key);
149+
if (field) {
150+
return vectorator(field.values)
151+
.filter( (value) => value !== null )
152+
.map( (value) => { return { text: String(value) }; });
153+
}
154+
return [];
155+
}
156+
157+
async fetchTags(): Promise<DataFrame> {
158+
const rawQuery = this.templateSrv.replace('$snowflake_adhoc_query');
159+
if (rawQuery === '$snowflake_adhoc_query') {
160+
return new ArrayDataFrame([]);
161+
} else {
162+
this.skipAdHocFilter = true;
163+
// this.adHocFilter.setTargetTable(rawQuery)
164+
return await this.runQuery({ queryText: rawQuery });
165+
}
166+
}
167+
}
168+
169+
// enum TagType {
170+
// query,
171+
// schema,
172+
// }
173+
174+
enum AdHocFilterStatus {
175+
none = 0,
176+
enabled,
177+
disabled,
54178
}

0 commit comments

Comments
 (0)