Loki: Fix adding of ad hoc filters for queries with parser and line_format expressions (#42590)

* Loki: Fix ad hoc filtewrs for queries with line_format

* Fix tests

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

@@ -0,0 +1,121 @@
+import { addLabelToQuery, addLabelToSelector } from './add_label_to_query';
+
+describe('addLabelToQuery()', () => {
+  it('should add label to simple query', () => {
+    expect(() => {
+      addLabelToQuery('foo', '', '');
+    }).toThrow();
+    expect(addLabelToQuery('{}', 'bar', 'baz')).toBe('{bar="baz"}');
+    expect(addLabelToQuery('{x="yy"}', 'bar', 'baz')).toBe('{bar="baz",x="yy"}');
+  });
+
+  it('should add custom operator', () => {
+    expect(addLabelToQuery('{}', 'bar', 'baz', '!=')).toBe('{bar!="baz"}');
+    expect(addLabelToQuery('{x="yy"}', 'bar', 'baz', '!=')).toBe('{bar!="baz",x="yy"}');
+  });
+
+  it('should not modify ranges', () => {
+    expect(addLabelToQuery('rate({}[1m])', 'foo', 'bar')).toBe('rate({foo="bar"}[1m])');
+  });
+
+  it('should detect in-order function use', () => {
+    expect(addLabelToQuery('sum by (xx) ({})', 'bar', 'baz')).toBe('sum by (xx) ({bar="baz"})');
+  });
+
+  it('should convert number Infinity to +Inf', () => {
+    expect(addLabelToQuery('sum(rate({}[5m])) by (le)', 'le', Infinity)).toBe('sum(rate({le="+Inf"}[5m])) by (le)');
+  });
+
+  it('should handle selectors with punctuation', () => {
+    expect(addLabelToQuery('{instance="my-host.com:9100"}', 'bar', 'baz')).toBe(
+      '{bar="baz",instance="my-host.com:9100"}'
+    );
+    expect(addLabelToQuery('{list="a,b,c"}', 'bar', 'baz')).toBe('{bar="baz",list="a,b,c"}');
+  });
+
+  it('should work on arithmetical expressions', () => {
+    expect(addLabelToQuery('{} + {}', 'bar', 'baz')).toBe('{bar="baz"} + {bar="baz"}');
+    expect(addLabelToQuery('avg({}) + sum({})', 'bar', 'baz')).toBe('avg({bar="baz"}) + sum({bar="baz"})');
+    expect(addLabelToQuery('{x="yy"} * {y="zz",a="bb"} * {}', 'bar', 'baz')).toBe(
+      '{bar="baz",x="yy"} * {a="bb",bar="baz",y="zz"} * {bar="baz"}'
+    );
+  });
+
+  it('should not add duplicate labels to a query', () => {
+    expect(addLabelToQuery(addLabelToQuery('{x="yy"}', 'bar', 'baz', '!='), 'bar', 'baz', '!=')).toBe(
+      '{bar!="baz",x="yy"}'
+    );
+    expect(addLabelToQuery(addLabelToQuery('rate({}[1m])', 'foo', 'bar'), 'foo', 'bar')).toBe('rate({foo="bar"}[1m])');
+    expect(addLabelToQuery(addLabelToQuery('{list="a,b,c"}', 'bar', 'baz'), 'bar', 'baz')).toBe(
+      '{bar="baz",list="a,b,c"}'
+    );
+    expect(addLabelToQuery(addLabelToQuery('avg({}) + sum({})', 'bar', 'baz'), 'bar', 'baz')).toBe(
+      'avg({bar="baz"}) + sum({bar="baz"})'
+    );
+  });
+
+  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"'
+    );
+  });
+
+  it('should add labels to metrics with logical operators', () => {
+    expect(addLabelToQuery('{} or {}', 'bar', 'baz')).toBe('{bar="baz"} or {bar="baz"}');
+    expect(addLabelToQuery('{} and {}', 'bar', 'baz')).toBe('{bar="baz"} and {bar="baz"}');
+  });
+
+  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)'
+    );
+  });
+
+  it('should not add ad-hoc filter to range', () => {
+    expect(addLabelToQuery('avg(rate(({job="foo"} > 0)[3h:])) by (label)', 'bar', 'baz')).toBe(
+      'avg(rate(({bar="baz",job="foo"} > 0)[3h:])) by (label)'
+    );
+  });
+  it('should not add ad-hoc filter to labels in label list provided with the group modifier', () => {
+    expect(
+      addLabelToQuery(
+        'max by (id, name, type) ({type=~"foo|bar|baz-test"}) * on(id) group_right(id, type, name) sum by (id) ({}) * 1000',
+        'bar',
+        'baz'
+      )
+    ).toBe(
+      'max by (id, name, type) ({bar="baz",type=~"foo|bar|baz-test"}) * on(id) group_right(id, type, name) sum by (id) ({bar="baz"}) * 1000'
+    );
+  });
+  it('should not add ad-hoc filter to labels in label list provided with the group modifier', () => {
+    expect(addLabelToQuery('rate({}[${__range_s}s])', 'bar', 'baz')).toBe('rate({bar="baz"}[${__range_s}s])');
+  });
+  it('should not add ad-hoc filter to labels to math operations', () => {
+    expect(addLabelToQuery('count({job!="foo"} < (5*1024*1024*1024) or vector(0)) - 1', 'bar', 'baz')).toBe(
+      'count({bar="baz",job!="foo"} < (5*1024*1024*1024) or vector(0)) - 1'
+    );
+  });
+  it('should not add adhoc filter to line_format expressions', () => {
+    expect(addLabelToQuery('{foo="bar"} | logfmt | line_format {{.status}}', 'bar', 'baz')).toBe(
+      '{bar="baz",foo="bar"} | logfmt | line_format {{.status}}'
+    );
+  });
+});
+
+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"}');
+  });
+});

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

@@ -0,0 +1,138 @@
+import { chain, isEqual } from 'lodash';
+import { LOKI_KEYWORDS } from './syntax';
+import { PROM_KEYWORDS, OPERATORS, LOGICAL_OPERATORS } from 'app/plugins/datasource/prometheus/promql';
+
+const builtInWords = [...PROM_KEYWORDS, ...OPERATORS, ...LOGICAL_OPERATORS, ...LOKI_KEYWORDS];
+
+// We want to extract all possible metrics and also keywords
+const metricsAndKeywordsRegexp = /([A-Za-z:][\w:]*)\b(?![\]{=!",])/g;
+
+export function addLabelToQuery(
+  query: string,
+  key: string,
+  value: string | number,
+  operator?: string,
+  hasNoMetrics?: boolean
+): string {
+  if (!key || !value) {
+    throw new Error('Need label to add to query.');
+  }
+
+  // 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();
+
+  // Add empty selectors to bare metric names
+  let previousWord: string;
+
+  query = query.replace(metricsAndKeywordsRegexp, (match, word, offset) => {
+    const isMetric = isWordMetric(query, word, offset, previousWord, hasNoMetrics);
+    previousWord = word;
+
+    return isMetric ? `${word}{}` : word;
+  });
+
+  //This is a RegExp for stream selector - e.g. {job="grafana"}
+  const selectorRegexp = /(\$)?{([^{]*)}/g;
+  const parts = [];
+  let lastIndex = 0;
+  let suffix = '';
+
+  let match = selectorRegexp.exec(query);
+  /* 
+    There are 2 possible false positive scenarios: 
+    
+    1. We match Grafana's variables with ${ syntax - such as${__rate_s}. To filter these out we could use negative lookbehind,
+    but Safari browser currently doesn't support it. Therefore we need to hack this by creating 2 matching groups. 
+    (\$) is for the Grafana's variables and if we match it, we know this is not a stream selector and we don't want to add label.
+
+    2. Log queries can include {{.label}} syntax when line_format is used. We need to filter these out by checking
+    if match starts with "{."
+  */
+  while (match) {
+    const prefix = query.slice(lastIndex, match.index);
+    lastIndex = match.index + match[2].length + 2;
+    suffix = query.slice(match.index + match[0].length);
+
+    // Filtering our false positives
+    if (match[0].startsWith('{.') || 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);
+    }
+
+    match = selectorRegexp.exec(query);
+  }
+
+  parts.push(suffix);
+  return parts.join('');
+}
+
+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}}`;
+}
+
+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/loki/datasource.test.ts

@@ -872,7 +872,7 @@ describe('LokiDatasource', () => {
       });
       describe('and query has parser', () => {
         it('then the correct label should be added for logs query', () => {
-          assertAdHocFilters('{bar="baz"} | logfmt', '{bar="baz"} | logfmt | job="grafana"', ds);
+          assertAdHocFilters('{bar="baz"} | logfmt', '{bar="baz",job="grafana"} | logfmt', ds);
         });
         it('then the correct label should be added for metrics query', () => {
           assertAdHocFilters('rate({bar="baz"} | logfmt [5m])', 'rate({bar="baz",job="grafana"} | logfmt [5m])', ds);
@@ -907,7 +907,7 @@ describe('LokiDatasource', () => {
       });
       describe('and query has parser', () => {
         it('then the correct label should be added for logs query', () => {
-          assertAdHocFilters('{bar="baz"} | logfmt', '{bar="baz"} | logfmt | job!="grafana"', ds);
+          assertAdHocFilters('{bar="baz"} | logfmt', '{bar="baz",job!="grafana"} | logfmt', ds);
         });
         it('then the correct label should be added for metrics query', () => {
           assertAdHocFilters('rate({bar="baz"} | logfmt [5m])', 'rate({bar="baz",job!="grafana"} | logfmt [5m])', ds);

public/app/plugins/datasource/loki/datasource.ts

@@ -33,7 +33,7 @@ import {
 } from '@grafana/data';
 import { BackendSrvRequest, FetchError, getBackendSrv } from '@grafana/runtime';
 import { getTemplateSrv, TemplateSrv } from 'app/features/templating/template_srv';
-import { addLabelToQuery } from 'app/plugins/datasource/prometheus/add_label_to_query';
+import { addLabelToQuery } from './add_label_to_query';
 import { getTimeSrv, TimeSrv } from 'app/features/dashboard/services/TimeSrv';
 import { convertToWebSocketUrl } from 'app/core/utils/explore';
 import {
@@ -707,14 +707,21 @@ export class LokiDatasource
         value = lokiRegularEscape(value);
       }
 
-      return this.addLabelToQuery(acc, key, value, operator);
+      return this.addLabelToQuery(acc, key, value, operator, true);
     }, expr);
 
     return expr;
   }
 
-  addLabelToQuery(queryExpr: string, key: string, value: string | number, operator: string) {
-    if (queryHasPipeParser(queryExpr) && !isMetricsQuery(queryExpr)) {
+  addLabelToQuery(
+    queryExpr: string,
+    key: string,
+    value: string | number,
+    operator: string,
+    // Override to make sure that we use label as actual label and not parsed label
+    notParsedLabelOverride?: boolean
+  ) {
+    if (queryHasPipeParser(queryExpr) && !isMetricsQuery(queryExpr) && !notParsedLabelOverride) {
       // If query has parser, we treat all labels as parsed and use | key="value" syntax
       return addParsedLabelToQuery(queryExpr, key, value, operator);
     } else {

public/app/plugins/datasource/loki/syntax.ts

@@ -168,6 +168,7 @@ export const RANGE_VEC_FUNCTIONS = [
 ];
 
 export const FUNCTIONS = [...AGGREGATION_OPERATORS, ...RANGE_VEC_FUNCTIONS];
+export const LOKI_KEYWORDS = [...FUNCTIONS, ...PIPE_OPERATORS, ...PIPE_PARSERS].map((keyword) => keyword.label);
 
 const tokenizer: Grammar = {
   comment: {

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

@@ -1,21 +1,7 @@
 import { chain, isEqual } from 'lodash';
+import { OPERATORS, LOGICAL_OPERATORS, PROM_KEYWORDS } from './promql';
 
-const keywords = 'by|without|on|ignoring|group_left|group_right|bool';
-const logicalOperators = 'or|and|unless';
-
-// Duplicate from mode-prometheus.js, which can't be used in tests due to global ace not being loaded.
-const builtInWords = [
-  keywords,
-  logicalOperators,
-  'count|count_values|min|max|avg|sum|stddev|stdvar|bottomk|topk|quantile',
-  'true|false|null|__name__|job',
-  'abs|absent|ceil|changes|clamp_max|clamp_min|count_scalar|day_of_month|day_of_week|days_in_month|delta|deriv',
-  'drop_common_labels|exp|floor|histogram_quantile|holt_winters|hour|idelta|increase|irate|label_replace|ln|log2',
-  'log10|minute|month|predict_linear|rate|resets|round|scalar|sort|sort_desc|sqrt|time|vector|year|avg_over_time',
-  'min_over_time|max_over_time|sum_over_time|count_over_time|quantile_over_time|stddev_over_time|stdvar_over_time',
-]
-  .join('|')
-  .split('|');
+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;
@@ -115,7 +101,7 @@ function isPositionInsideChars(text: string, position: number, openChar: string,
 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 && keywords.split('|').indexOf(previousWord) > -1;
+  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

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

@@ -16,6 +16,7 @@ export const RATE_RANGES: CompletionItem[] = [
 ];
 
 export const OPERATORS = ['by', 'group_left', 'group_right', 'ignoring', 'on', 'offset', 'without'];
+export const LOGICAL_OPERATORS = ['or', 'and', 'unless'];
 
 const AGGREGATION_OPERATORS: CompletionItem[] = [
   {
@@ -427,6 +428,8 @@ export const FUNCTIONS = [
   },
 ];
 
+export const PROM_KEYWORDS = FUNCTIONS.map((keyword) => keyword.label);
+
 const tokenizer: Grammar = {
   comment: {
     pattern: /#.*/,