Prometheus: Fix support of AdHoc filters for specific queries (#46547)

* Use parser to deal with ad hoc filters

* Add comment

* Fix variable naming

* Fix tests

public/app/plugins/datasource/prometheus/add_label_to_query.test.ts

@@ -1,4 +1,4 @@
-import { addLabelToQuery, addLabelToSelector } from './add_label_to_query';
+import { addLabelToQuery } from './add_label_to_query';
 
 describe('addLabelToQuery()', () => {
   it('should add label to simple query', () => {
@@ -7,13 +7,13 @@ describe('addLabelToQuery()', () => {
     }).toThrow();
     expect(addLabelToQuery('foo', 'bar', 'baz')).toBe('foo{bar="baz"}');
     expect(addLabelToQuery('foo{}', 'bar', 'baz')).toBe('foo{bar="baz"}');
-    expect(addLabelToQuery('foo{x="yy"}', 'bar', 'baz')).toBe('foo{bar="baz",x="yy"}');
+    expect(addLabelToQuery('foo{x="yy"}', 'bar', 'baz')).toBe('foo{x="yy", bar="baz"}');
     expect(addLabelToQuery('metric > 0.001', 'foo', 'bar')).toBe('metric{foo="bar"} > 0.001');
   });
 
   it('should add custom operator', () => {
     expect(addLabelToQuery('foo{}', 'bar', 'baz', '!=')).toBe('foo{bar!="baz"}');
-    expect(addLabelToQuery('foo{x="yy"}', 'bar', 'baz', '!=')).toBe('foo{bar!="baz",x="yy"}');
+    expect(addLabelToQuery('foo{x="yy"}', 'bar', 'baz', '!=')).toBe('foo{x="yy", bar!="baz"}');
   });
 
   it('should not modify ranges', () => {
@@ -32,33 +32,33 @@ describe('addLabelToQuery()', () => {
 
   it('should handle selectors with punctuation', () => {
     expect(addLabelToQuery('foo{instance="my-host.com:9100"}', 'bar', 'baz')).toBe(
-      'foo{bar="baz",instance="my-host.com:9100"}'
+      'foo{instance="my-host.com:9100", bar="baz"}'
     );
     expect(addLabelToQuery('foo:metric:rate1m', 'bar', 'baz')).toBe('foo:metric:rate1m{bar="baz"}');
     expect(addLabelToQuery('avg(foo:metric:rate1m{a="b"})', 'bar', 'baz')).toBe(
-      'avg(foo:metric:rate1m{a="b",bar="baz"})'
+      'avg(foo:metric:rate1m{a="b", bar="baz"})'
     );
-    expect(addLabelToQuery('foo{list="a,b,c"}', 'bar', 'baz')).toBe('foo{bar="baz",list="a,b,c"}');
+    expect(addLabelToQuery('foo{list="a,b,c"}', 'bar', 'baz')).toBe('foo{list="a,b,c", bar="baz"}');
   });
 
   it('should work on arithmetical expressions', () => {
     expect(addLabelToQuery('foo + foo', 'bar', 'baz')).toBe('foo{bar="baz"} + foo{bar="baz"}');
-    expect(addLabelToQuery('foo{x="yy"} + metric', 'bar', 'baz')).toBe('foo{bar="baz",x="yy"} + metric{bar="baz"}');
+    expect(addLabelToQuery('foo{x="yy"} + metric', 'bar', 'baz')).toBe('foo{x="yy", bar="baz"} + metric{bar="baz"}');
     expect(addLabelToQuery('avg(foo) + sum(xx_yy)', 'bar', 'baz')).toBe('avg(foo{bar="baz"}) + sum(xx_yy{bar="baz"})');
     expect(addLabelToQuery('foo{x="yy"} * metric{y="zz",a="bb"} * metric2', 'bar', 'baz')).toBe(
-      'foo{bar="baz",x="yy"} * metric{a="bb",bar="baz",y="zz"} * metric2{bar="baz"}'
+      'foo{x="yy", bar="baz"} * metric{y="zz", a="bb", bar="baz"} * metric2{bar="baz"}'
     );
   });
 
   it('should not add duplicate labels to a query', () => {
     expect(addLabelToQuery(addLabelToQuery('foo{x="yy"}', 'bar', 'baz', '!='), 'bar', 'baz', '!=')).toBe(
-      'foo{bar!="baz",x="yy"}'
+      'foo{x="yy", bar!="baz"}'
     );
     expect(addLabelToQuery(addLabelToQuery('rate(metric[1m])', 'foo', 'bar'), 'foo', 'bar')).toBe(
       'rate(metric{foo="bar"}[1m])'
     );
     expect(addLabelToQuery(addLabelToQuery('foo{list="a,b,c"}', 'bar', 'baz'), 'bar', 'baz')).toBe(
-      'foo{bar="baz",list="a,b,c"}'
+      'foo{list="a,b,c", bar="baz"}'
     );
     expect(addLabelToQuery(addLabelToQuery('avg(foo) + sum(xx_yy)', 'bar', 'baz'), 'bar', 'baz')).toBe(
       'avg(foo{bar="baz"}) + sum(xx_yy{bar="baz"})'
@@ -66,14 +66,8 @@ describe('addLabelToQuery()', () => {
   });
 
   it('should not remove filters', () => {
-    expect(addLabelToQuery('{x="y"} |="yy"', 'bar', 'baz')).toBe('{bar="baz",x="y"} |="yy"');
-    expect(addLabelToQuery('{x="y"} |="yy" !~"xx"', 'bar', 'baz')).toBe('{bar="baz",x="y"} |="yy" !~"xx"');
-  });
-
-  it('should add label to query properly with Loki datasource', () => {
-    expect(addLabelToQuery('{job="grafana"} |= "foo-bar"', 'filename', 'test.txt', undefined, true)).toBe(
-      '{filename="test.txt",job="grafana"} |= "foo-bar"'
-    );
+    expect(addLabelToQuery('{x="y"} |="yy"', 'bar', 'baz')).toBe('{x="y", bar="baz"} |="yy"');
+    expect(addLabelToQuery('{x="y"} |="yy" !~"xx"', 'bar', 'baz')).toBe('{x="y", bar="baz"} |="yy" !~"xx"');
   });
 
   it('should add labels to metrics with logical operators', () => {
@@ -83,13 +77,13 @@ describe('addLabelToQuery()', () => {
 
   it('should not add ad-hoc filter to template variables', () => {
     expect(addLabelToQuery('sum(rate({job="foo"}[2m])) by (value $variable)', 'bar', 'baz')).toBe(
-      'sum(rate({bar="baz",job="foo"}[2m])) by (value $variable)'
+      'sum(rate({job="foo", bar="baz"}[2m])) by (value $variable)'
     );
   });
 
   it('should not add ad-hoc filter to range', () => {
     expect(addLabelToQuery('avg(rate((my_metric{job="foo"} > 0)[3h:])) by (label)', 'bar', 'baz')).toBe(
-      'avg(rate((my_metric{bar="baz",job="foo"} > 0)[3h:])) by (label)'
+      'avg(rate((my_metric{job="foo", bar="baz"} > 0)[3h:])) by (label)'
     );
   });
   it('should not add ad-hoc filter to labels in label list provided with the group modifier', () => {
@@ -100,7 +94,7 @@ describe('addLabelToQuery()', () => {
         'baz'
       )
     ).toBe(
-      'max by (id, name, type) (my_metric{bar="baz",type=~"foo|bar|baz-test"}) * on(id) group_right(id, type, name) sum by (id) (my_metric{bar="baz"}) * 1000'
+      'max by (id, name, type) (my_metric{type=~"foo|bar|baz-test", bar="baz"}) * on(id) group_right(id, type, name) sum by (id) (my_metric{bar="baz"}) * 1000'
     );
   });
   it('should not add ad-hoc filter to labels in label list provided with the group modifier', () => {
@@ -110,20 +104,11 @@ describe('addLabelToQuery()', () => {
   });
   it('should not add ad-hoc filter to labels to math operations', () => {
     expect(addLabelToQuery('count(my_metric{job!="foo"} < (5*1024*1024*1024) or vector(0)) - 1', 'bar', 'baz')).toBe(
-      'count(my_metric{bar="baz",job!="foo"} < (5*1024*1024*1024) or vector(0)) - 1'
+      'count(my_metric{job!="foo", bar="baz"} < (5*1024*1024*1024) or vector(0)) - 1'
     );
   });
-});
 
-describe('addLabelToSelector()', () => {
-  test('should add a label to an empty selector', () => {
-    expect(addLabelToSelector('{}', 'foo', 'bar')).toBe('{foo="bar"}');
-    expect(addLabelToSelector('', 'foo', 'bar')).toBe('{foo="bar"}');
-  });
-  test('should add a label to a selector', () => {
-    expect(addLabelToSelector('{foo="bar"}', 'baz', '42')).toBe('{baz="42",foo="bar"}');
-  });
-  test('should add a label to a selector with custom operator', () => {
-    expect(addLabelToSelector('{}', 'baz', '42', '!=')).toBe('{baz!="42"}');
+  it('should not add ad-hoc filter bool operator', () => {
+    expect(addLabelToQuery('ALERTS < bool 1', 'bar', 'baz')).toBe('ALERTS{bar="baz"} < bool 1');
   });
 });

public/app/plugins/datasource/prometheus/add_label_to_query.ts

@@ -1,129 +1,99 @@
-import { chain, isEqual } from 'lodash';
-import { OPERATORS, LOGICAL_OPERATORS, PROM_KEYWORDS } from './promql';
+import { parser } from 'lezer-promql';
+import { buildVisualQueryFromString } from './querybuilder/parsing';
+import { PromQueryModeller } from './querybuilder/PromQueryModeller';
+import { QueryBuilderLabelFilter } from './querybuilder/shared/types';
+import { PromVisualQuery } from './querybuilder/types';
 
-const builtInWords = [...PROM_KEYWORDS, ...OPERATORS, ...LOGICAL_OPERATORS];
-
-// We want to extract all possible metrics and also keywords
-const metricsAndKeywordsRegexp = /([A-Za-z:][\w:]*)\b(?![\]{=!",])/g;
-// Safari currently doesn't support negative lookbehind. When it does, we should refactor this.
-// We are creating 2 matching groups. (\$) is for the Grafana's variables such as ${__rate_s}. We want to ignore
-// ${__rate_s} and not add variable to it.
-const selectorRegexp = /(\$)?{([^{]*)}/g;
-
-export function addLabelToQuery(
-  query: string,
-  key: string,
-  value: string | number,
-  operator?: string,
-  hasNoMetrics?: boolean
-): string {
+/**
+ * Adds label filter to existing query. Useful for query modification for example for ad hoc filters.
+ *
+ * It uses PromQL parser to find instances of metric and labels, alters them and then splices them back into the query.
+ * Ideally we could use the parse -> change -> render is a simple 3 steps but right now building the visual query
+ * object does not support all possible queries.
+ *
+ * So instead this just operates on substrings of the query with labels and operates just on those. This makes this
+ * more robust and can alter even invalid queries, and preserves in general the query structure and whitespace.
+ * @param query
+ * @param key
+ * @param value
+ * @param operator
+ */
+export function addLabelToQuery(query: string, key: string, value: string | number, operator = '='): string {
   if (!key || !value) {
     throw new Error('Need label to add to query.');
   }
 
+  const vectorSelectorPositions = getVectorSelectorPositions(query);
+  if (!vectorSelectorPositions.length) {
+    return query;
+  }
+
+  const filter = toLabelFilter(key, value, operator);
+  return addFilter(query, vectorSelectorPositions, filter);
+}
+
+type VectorSelectorPosition = { from: number; to: number; query: PromVisualQuery };
+
+/**
+ * Parse the string and get all VectorSelector positions in the query together with parsed representation of the vector
+ * selector.
+ * @param query
+ */
+function getVectorSelectorPositions(query: string): VectorSelectorPosition[] {
+  const tree = parser.parse(query);
+  const positions: VectorSelectorPosition[] = [];
+  tree.iterate({
+    enter: (type, from, to, get): false | void => {
+      if (type.name === 'VectorSelector') {
+        const visQuery = buildVisualQueryFromString(query.substring(from, to));
+        positions.push({ query: visQuery.query, from, to });
+        return false;
+      }
+    },
+  });
+  return positions;
+}
+
+function toLabelFilter(key: string, value: string | number, operator: string): QueryBuilderLabelFilter {
   // We need to make sure that we convert the value back to string because it may be a number
   const transformedValue = value === Infinity ? '+Inf' : value.toString();
+  return { label: key, op: operator, value: transformedValue };
+}
 
-  // Add empty selectors to bare metric names
-  let previousWord: string;
+function addFilter(
+  query: string,
+  vectorSelectorPositions: VectorSelectorPosition[],
+  filter: QueryBuilderLabelFilter
+): string {
+  const modeller = new PromQueryModeller();
+  let newQuery = '';
+  let prev = 0;
 
-  query = query.replace(metricsAndKeywordsRegexp, (match, word, offset) => {
-    const isMetric = isWordMetric(query, word, offset, previousWord, hasNoMetrics);
-    previousWord = word;
+  for (let i = 0; i < vectorSelectorPositions.length; i++) {
+    // This is basically just doing splice on a string for each matched vector selector.
 
-    return isMetric ? `${word}{}` : word;
-  });
+    const match = vectorSelectorPositions[i];
+    const isLast = i === vectorSelectorPositions.length - 1;
 
-  // Adding label to existing selectors
-  let match = selectorRegexp.exec(query);
-  const parts = [];
-  let lastIndex = 0;
-  let suffix = '';
+    const start = query.substring(prev, match.from);
+    const end = isLast ? query.substring(match.to) : '';
 
-  while (match) {
-    const prefix = query.slice(lastIndex, match.index);
-    lastIndex = match.index + match[2].length + 2;
-    suffix = query.slice(match.index + match[0].length);
-    // If we matched 1st group, we know it is Grafana's variable and we don't want to add labels
-    if (match[1]) {
-      parts.push(prefix);
-      parts.push(match[0]);
-    } else {
-      // If we didn't match first group, we are inside selector and we want to add labels
-      const selector = match[2];
-      const selectorWithLabel = addLabelToSelector(selector, key, transformedValue, operator);
-      parts.push(prefix, selectorWithLabel);
+    if (!labelExists(match.query.labels, filter)) {
+      // We don't want to add duplicate labels.
+      match.query.labels.push(filter);
     }
-
-    match = selectorRegexp.exec(query);
+    const newLabels = modeller.renderQuery(match.query);
+    newQuery += start + newLabels + end;
+    prev = match.to;
   }
-
-  parts.push(suffix);
-  return parts.join('');
+  return newQuery;
 }
 
-const labelRegexp = /(\w+)\s*(=|!=|=~|!~)\s*("[^"]*")/g;
-
-export function addLabelToSelector(selector: string, labelKey: string, labelValue: string, labelOperator?: string) {
-  const parsedLabels = [];
-
-  // Split selector into labels
-  if (selector) {
-    let match = labelRegexp.exec(selector);
-    while (match) {
-      parsedLabels.push({ key: match[1], operator: match[2], value: match[3] });
-      match = labelRegexp.exec(selector);
-    }
-  }
-
-  // Add new label
-  const operatorForLabelKey = labelOperator || '=';
-  parsedLabels.push({ key: labelKey, operator: operatorForLabelKey, value: `"${labelValue}"` });
-
-  // Sort labels by key and put them together
-  const formatted = chain(parsedLabels)
-    .uniqWith(isEqual)
-    .compact()
-    .sortBy('key')
-    .map(({ key, operator, value }) => `${key}${operator}${value}`)
-    .value()
-    .join(',');
-
-  return `{${formatted}}`;
+/**
+ * Check if label exists in the list of labels but ignore the operator.
+ * @param labels
+ * @param filter
+ */
+function labelExists(labels: QueryBuilderLabelFilter[], filter: QueryBuilderLabelFilter) {
+  return labels.find((label) => label.label === filter.label && label.value === filter.value);
 }
-
-function isPositionInsideChars(text: string, position: number, openChar: string, closeChar: string) {
-  const nextSelectorStart = text.slice(position).indexOf(openChar);
-  const nextSelectorEnd = text.slice(position).indexOf(closeChar);
-  return nextSelectorEnd > -1 && (nextSelectorStart === -1 || nextSelectorStart > nextSelectorEnd);
-}
-
-function isWordMetric(query: string, word: string, offset: number, previousWord: string, hasNoMetrics?: boolean) {
-  const insideSelector = isPositionInsideChars(query, offset, '{', '}');
-  // Handle "sum by (key) (metric)"
-  const previousWordIsKeyWord = previousWord && OPERATORS.indexOf(previousWord) > -1;
-  // Check for colon as as "word boundary" symbol
-  const isColonBounded = word.endsWith(':');
-  // Check for words that start with " which means that they are not metrics
-  const startsWithQuote = query[offset - 1] === '"';
-  // Check for template variables
-  const isTemplateVariable = query[offset - 1] === '$';
-  // Check for time units
-  const isTimeUnit = ['s', 'm', 'h', 'd', 'w'].includes(word) && Boolean(Number(query[offset - 1]));
-
-  if (
-    !hasNoMetrics &&
-    !insideSelector &&
-    !isColonBounded &&
-    !previousWordIsKeyWord &&
-    !startsWithQuote &&
-    !isTemplateVariable &&
-    !isTimeUnit &&
-    builtInWords.indexOf(word) === -1
-  ) {
-    return true;
-  }
-  return false;
-}
-
-export default addLabelToQuery;

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

@@ -247,7 +247,7 @@ describe('PrometheusDatasource', () => {
         },
       ]);
       const result = ds.createQuery(target as any, { interval: '15s' } as any, 0, 0);
-      expect(result).toMatchObject({ expr: 'metric{job="foo",k1="v1",k2!="v2"} - metric{k1="v1",k2!="v2"}' });
+      expect(result).toMatchObject({ expr: 'metric{job="foo", k1="v1", k2!="v2"} - metric{k1="v1", k2!="v2"}' });
     });
 
     it('should add escaping if needed to regex filter expressions', () => {
@@ -265,7 +265,7 @@ describe('PrometheusDatasource', () => {
       ]);
       const result = ds.createQuery(target as any, { interval: '15s' } as any, 0, 0);
       expect(result).toMatchObject({
-        expr: `metric{job="foo",k1=~"v.*",k2=~"v\\\\'.*"} - metric{k1=~"v.*",k2=~"v\\\\'.*"}`,
+        expr: `metric{job="foo", k1=~"v.*", k2=~"v\\\\'.*"} - metric{k1=~"v.*", k2=~"v\\\\'.*"}`,
       });
     });
   });
@@ -637,7 +637,7 @@ describe('PrometheusDatasource', () => {
       };
 
       const result = ds.applyTemplateVariables(query, {});
-      expect(result).toMatchObject({ expr: 'test{job="bar",k1="v1",k2!="v2"}' });
+      expect(result).toMatchObject({ expr: 'test{job="bar", k1="v1", k2!="v2"}' });
     });
   });
 
@@ -2172,7 +2172,7 @@ describe('modifyQuery', () => {
         const result = ds.modifyQuery(query, action);
 
         expect(result.refId).toEqual('A');
-        expect(result.expr).toEqual('go_goroutines{cluster="us-cluster",pod="pod-123"}');
+        expect(result.expr).toEqual('go_goroutines{cluster="us-cluster", pod="pod-123"}');
       });
     });
   });
@@ -2202,7 +2202,7 @@ describe('modifyQuery', () => {
         const result = ds.modifyQuery(query, action);
 
         expect(result.refId).toEqual('A');
-        expect(result.expr).toEqual('go_goroutines{cluster="us-cluster",pod!="pod-123"}');
+        expect(result.expr).toEqual('go_goroutines{cluster="us-cluster", pod!="pod-123"}');
       });
     });
   });

public/app/plugins/datasource/prometheus/datasource.ts

@@ -34,7 +34,7 @@ import {
 import { safeStringifyValue } from 'app/core/utils/explore';
 import { getTimeSrv, TimeSrv } from 'app/features/dashboard/services/TimeSrv';
 import { getTemplateSrv, TemplateSrv } from 'app/features/templating/template_srv';
-import addLabelToQuery from './add_label_to_query';
+import { addLabelToQuery } from './add_label_to_query';
 import PrometheusLanguageProvider from './language_provider';
 import { expandRecordingRules } from './language_utils';
 import { getInitHints, getQueryHints } from './query_hints';
@@ -949,15 +949,15 @@ export class PrometheusDatasource
 
   enhanceExprWithAdHocFilters(expr: string) {
     const adhocFilters = this.templateSrv.getAdhocFilters(this.name);
-    let finalQuery = expr;
-    finalQuery = adhocFilters.reduce((acc: string, filter: { key?: any; operator?: any; value?: any }) => {
+
+    const finalQuery = adhocFilters.reduce((acc: string, filter: { key?: any; operator?: any; value?: any }) => {
       const { key, operator } = filter;
       let { value } = filter;
       if (operator === '=~' || operator === '!~') {
         value = prometheusRegularEscape(value);
       }
       return addLabelToQuery(acc, key, value, operator);
-    }, finalQuery);
+    }, expr);
     return finalQuery;
   }
 

public/app/plugins/datasource/prometheus/language_utils.test.ts

@@ -163,7 +163,7 @@ describe('expandRecordingRules()', () => {
       expandRecordingRules('metricA{label1="value1",label2="value,2"}', {
         metricA: 'rate(fooA[])',
       })
-    ).toBe('rate(fooA{label1="value1",label2="value,2"}[])');
+    ).toBe('rate(fooA{label1="value1", label2="value,2"}[])');
     expect(
       expandRecordingRules('metricA{label1="value1"} / metricB{label2="value2"}', {
         metricA: 'rate(fooA[])',
@@ -175,7 +175,7 @@ describe('expandRecordingRules()', () => {
         metricA: 'rate(fooA[])',
         metricB: 'rate(fooB[])',
       })
-    ).toBe('rate(fooA{label1="value1",label2="value2"}[])/ rate(fooB{label3="value3"}[])');
+    ).toBe('rate(fooA{label1="value1", label2="value2"}[])/ rate(fooB{label3="value3"}[])');
   });
 });
 

public/app/plugins/datasource/prometheus/language_utils.ts

@@ -152,7 +152,7 @@ function addLabelsToExpression(expr: string, invalidLabelsRegexp: RegExp) {
     return '';
   });
 
-  // Loop trough all of the label objects and add them to query.
+  // Loop through all label objects and add them to query.
   // As a starting point we have valid query without the labels.
   let result = exprBeforeRegexMatch;
   arrayOfLabelObjects.filter(Boolean).forEach((obj) => {