Loki: Remove prom utils in import/export from abstract query (#80312)

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

@@ -320,10 +320,30 @@ describe('fetchLabels', () => {
 describe('Query imports', () => {
   const datasource = setup({});
 
-  it('returns empty queries', async () => {
-    const instance = new LanguageProvider(datasource);
-    const result = await instance.importFromAbstractQuery({ refId: 'bar', labelMatchers: [] });
-    expect(result).toEqual({ refId: 'bar', expr: '', queryType: LokiQueryType.Range });
+  describe('importing from abstract query', () => {
+    it('returns empty queries', async () => {
+      const instance = new LanguageProvider(datasource);
+      const result = await instance.importFromAbstractQuery({ refId: 'bar', labelMatchers: [] });
+      expect(result).toEqual({ refId: 'bar', expr: '', queryType: LokiQueryType.Range });
+    });
+
+    it('returns valid query', () => {
+      const instance = new LanguageProvider(datasource);
+      const result = instance.importFromAbstractQuery({
+        refId: 'bar',
+        labelMatchers: [
+          { name: 'label1', operator: AbstractLabelOperator.Equal, value: 'value1' },
+          { name: 'label2', operator: AbstractLabelOperator.NotEqual, value: 'value2' },
+          { name: 'label3', operator: AbstractLabelOperator.EqualRegEx, value: 'value3' },
+          { name: 'label4', operator: AbstractLabelOperator.NotEqualRegEx, value: 'value4' },
+        ],
+      });
+      expect(result).toEqual({
+        refId: 'bar',
+        expr: '{label1="value1", label2!="value2", label3=~"value3", label4!~"value4"}',
+        queryType: LokiQueryType.Range,
+      });
+    });
   });
 
   describe('exporting to abstract query', () => {
@@ -345,6 +365,42 @@ describe('Query imports', () => {
         ],
       });
     });
+
+    it('exports labels in metric query', async () => {
+      const instance = new LanguageProvider(datasource);
+      const abstractQuery = instance.exportToAbstractQuery({
+        refId: 'bar',
+        expr: 'rate({label1="value1", label2!="value2"}[5m])',
+        instant: true,
+        range: false,
+      });
+      expect(abstractQuery).toMatchObject({
+        refId: 'bar',
+        labelMatchers: [
+          { name: 'label1', operator: AbstractLabelOperator.Equal, value: 'value1' },
+          { name: 'label2', operator: AbstractLabelOperator.NotEqual, value: 'value2' },
+        ],
+      });
+    });
+
+    it('exports labels in query with multiple stream selectors', async () => {
+      const instance = new LanguageProvider(datasource);
+      const abstractQuery = instance.exportToAbstractQuery({
+        refId: 'bar',
+        expr: 'rate({label1="value1", label2!="value2"}[5m]) + rate({label3=~"value3", label4!~"value4"}[5m])',
+        instant: true,
+        range: false,
+      });
+      expect(abstractQuery).toMatchObject({
+        refId: 'bar',
+        labelMatchers: [
+          { name: 'label1', operator: AbstractLabelOperator.Equal, value: 'value1' },
+          { name: 'label2', operator: AbstractLabelOperator.NotEqual, value: 'value2' },
+          { name: 'label3', operator: AbstractLabelOperator.EqualRegEx, value: 'value3' },
+          { name: 'label4', operator: AbstractLabelOperator.NotEqualRegEx, value: 'value4' },
+        ],
+      });
+    });
   });
 
   describe('getParserAndLabelKeys()', () => {

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

@@ -1,17 +1,18 @@
+import { flatten } from 'lodash';
 import { LRUCache } from 'lru-cache';
-import Prism from 'prismjs';
 
 import { LanguageProvider, AbstractQuery, KeyValue, getDefaultTimeRange, TimeRange } from '@grafana/data';
 import { config } from '@grafana/runtime';
-import { extractLabelMatchers, processLabels, toPromLikeExpr } from 'app/plugins/datasource/prometheus/language_utils';
 
 import { DEFAULT_MAX_LINES_SAMPLE, LokiDatasource } from './datasource';
+import { abstractQueryToExpr, mapAbstractOperatorsToOp, processLabels } from './languageUtils';
+import { getStreamSelectorsFromQuery } from './queryUtils';
+import { buildVisualQueryFromString } from './querybuilder/parsing';
 import {
   extractLabelKeysFromDataFrame,
   extractLogParserFromDataFrame,
   extractUnwrapLabelKeysFromDataFrame,
 } from './responseUtils';
-import syntax from './syntax';
 import { ParserAndLabelKeysResult, LokiQuery, LokiQueryType, LabelType } from './types';
 
 const NS_IN_MS = 1000000;
@@ -88,20 +89,32 @@ export default class LokiLanguageProvider extends LanguageProvider {
   importFromAbstractQuery(labelBasedQuery: AbstractQuery): LokiQuery {
     return {
       refId: labelBasedQuery.refId,
-      expr: toPromLikeExpr(labelBasedQuery),
+      expr: abstractQueryToExpr(labelBasedQuery),
       queryType: LokiQueryType.Range,
     };
   }
 
   exportToAbstractQuery(query: LokiQuery): AbstractQuery {
-    const lokiQuery = query.expr;
-    if (!lokiQuery || lokiQuery.length === 0) {
+    if (!query.expr || query.expr.length === 0) {
       return { refId: query.refId, labelMatchers: [] };
     }
-    const tokens = Prism.tokenize(lokiQuery, syntax);
+    const streamSelectors = getStreamSelectorsFromQuery(query.expr);
+
+    const labelMatchers = streamSelectors.map((streamSelector) => {
+      const visualQuery = buildVisualQueryFromString(streamSelector).query;
+      const matchers = visualQuery.labels.map((label) => {
+        return {
+          name: label.label,
+          value: label.value,
+          operator: mapAbstractOperatorsToOp[label.op],
+        };
+      });
+      return matchers;
+    });
+
     return {
       refId: query.refId,
-      labelMatchers: extractLabelMatchers(tokens),
+      labelMatchers: flatten(labelMatchers),
     };
   }
 

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

@@ -1,9 +1,11 @@
-import { toDataFrame, FieldType } from '@grafana/data';
+import { toDataFrame, FieldType, AbstractQuery, AbstractLabelOperator } from '@grafana/data';
 
 import {
+  abstractQueryToExpr,
   escapeLabelValueInExactSelector,
   getLabelTypeFromFrame,
   isBytesString,
+  processLabels,
   unescapeLabelValue,
 } from './languageUtils';
 import { LabelType } from './types';
@@ -90,3 +92,38 @@ describe('getLabelTypeFromFrame', () => {
     expect(getLabelTypeFromFrame('job', frameWithoutTypes, 0)).toBe(null);
   });
 });
+
+describe('abstractQueryToExpr', () => {
+  it('export abstract query to expr', () => {
+    const abstractQuery: AbstractQuery = {
+      refId: 'bar',
+      labelMatchers: [
+        { name: 'label1', operator: AbstractLabelOperator.Equal, value: 'value1' },
+        { name: 'label2', operator: AbstractLabelOperator.NotEqual, value: 'value2' },
+        { name: 'label3', operator: AbstractLabelOperator.EqualRegEx, value: 'value3' },
+        { name: 'label4', operator: AbstractLabelOperator.NotEqualRegEx, value: 'value4' },
+      ],
+    };
+
+    expect(abstractQueryToExpr(abstractQuery)).toBe(
+      '{label1="value1", label2!="value2", label3=~"value3", label4!~"value4"}'
+    );
+  });
+});
+
+describe('processLabels', () => {
+  it('export abstract query to expr', () => {
+    const labels: Array<{ [key: string]: string }> = [
+      { label1: 'value1' },
+      { label2: 'value2' },
+      { label3: 'value3' },
+      { label1: 'value1' },
+      { label1: 'value1b' },
+    ];
+
+    expect(processLabels(labels)).toEqual({
+      keys: ['label1', 'label2', 'label3'],
+      values: { label1: ['value1', 'value1b'], label2: ['value2'], label3: ['value3'] },
+    });
+  });
+});

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

@@ -1,4 +1,6 @@
-import { DataFrame, TimeRange } from '@grafana/data';
+import { invert } from 'lodash';
+
+import { AbstractLabelMatcher, AbstractLabelOperator, AbstractQuery, DataFrame, TimeRange } from '@grafana/data';
 
 import { LabelType } from './types';
 
@@ -111,3 +113,56 @@ export function getLabelTypeFromFrame(labelKey: string, frame?: DataFrame, index
       return null;
   }
 }
+
+export const mapOpToAbstractOp: Record<AbstractLabelOperator, string> = {
+  [AbstractLabelOperator.Equal]: '=',
+  [AbstractLabelOperator.NotEqual]: '!=',
+  [AbstractLabelOperator.EqualRegEx]: '=~',
+  [AbstractLabelOperator.NotEqualRegEx]: '!~',
+};
+
+// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+export const mapAbstractOperatorsToOp = invert(mapOpToAbstractOp) as Record<string, AbstractLabelOperator>;
+
+export function abstractQueryToExpr(labelBasedQuery: AbstractQuery): string {
+  const expr = labelBasedQuery.labelMatchers
+    .map((selector: AbstractLabelMatcher) => {
+      const operator = mapOpToAbstractOp[selector.operator];
+      if (operator) {
+        return `${selector.name}${operator}"${selector.value}"`;
+      } else {
+        return '';
+      }
+    })
+    .filter((e: string) => e !== '')
+    .join(', ');
+
+  return expr ? `{${expr}}` : '';
+}
+
+export function processLabels(labels: Array<{ [key: string]: string }>) {
+  const valueSet: { [key: string]: Set<string> } = {};
+  labels.forEach((label) => {
+    Object.keys(label).forEach((key) => {
+      if (!valueSet[key]) {
+        valueSet[key] = new Set();
+      }
+      if (!valueSet[key].has(label[key])) {
+        valueSet[key].add(label[key]);
+      }
+    });
+  });
+
+  const valueArray: { [key: string]: string[] } = {};
+  limitSuggestions(Object.keys(valueSet)).forEach((key) => {
+    valueArray[key] = limitSuggestions(Array.from(valueSet[key]));
+  });
+
+  return { values: valueArray, keys: Object.keys(valueArray) };
+}
+
+// Max number of items (metrics, labels, values) that we display as suggestions. Prevents from running out of memory.
+export const SUGGESTIONS_LIMIT = 10000;
+export function limitSuggestions(items: string[]) {
+  return items.slice(0, SUGGESTIONS_LIMIT);
+}