Loki Query Utils: Extract and refactor common functionality (#70185)

* Query utils: refactor isQueryPipelineErrorFiltering and getLogQueryFromMetricsQuery

* Query utils: refactor isQueryWithRangeVariable and getHighlighterExpressionsFromQuery

* Get parser: return log expression even if no pipelineExpr is present

* Update tests

public/app/plugins/datasource/loki/components/monaco-query-field/monaco-completion-provider/situation.test.ts

@@ -95,7 +95,7 @@ describe('situation', () => {
       type: 'AFTER_SELECTOR',
       afterPipe: true,
       hasSpace: false,
-      logQuery: '{place="luna"}| logfmt |',
+      logQuery: '{place="luna"} | logfmt |',
     });
   });
 
@@ -179,7 +179,7 @@ describe('situation', () => {
       'quantile_over_time(0.99, {cluster="ops-tools1",container="ingress-nginx"} | json | __error__ = "" | unwrap ^',
       {
         type: 'AFTER_UNWRAP',
-        logQuery: '{cluster="ops-tools1",container="ingress-nginx"}| json | __error__ = ""',
+        logQuery: '{cluster="ops-tools1",container="ingress-nginx"} | json | __error__ = ""',
       }
     );
 

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

@@ -11,6 +11,8 @@ import {
   requestSupportsSplitting,
   isQueryWithDistinct,
   isQueryWithRangeVariable,
+  isQueryPipelineErrorFiltering,
+  getLogQueryFromMetricsQuery,
   getNormalizedLokiQuery,
 } from './queryUtils';
 import { LokiQuery, LokiQueryType } from './types';
@@ -393,3 +395,24 @@ describe('requestSupportsSplitting', () => {
     expect(requestSupportsSplitting(requests)).toBe(true);
   });
 });
+
+describe('isQueryPipelineErrorFiltering', () => {
+  it('identifies pipeline error filters', () => {
+    expect(isQueryPipelineErrorFiltering('{job="grafana"} | logfmt | __error__=""')).toBe(true);
+    expect(isQueryPipelineErrorFiltering('{job="grafana"} | logfmt | error=""')).toBe(false);
+  });
+});
+
+describe('getLogQueryFromMetricsQuery', () => {
+  it('returns the log query from a metric query', () => {
+    expect(getLogQueryFromMetricsQuery('count_over_time({job="grafana"} | logfmt | label="value" [1m])')).toBe(
+      '{job="grafana"} | logfmt | label="value"'
+    );
+    expect(getLogQueryFromMetricsQuery('count_over_time({job="grafana"} [1m])')).toBe('{job="grafana"}');
+    expect(
+      getLogQueryFromMetricsQuery(
+        'sum(quantile_over_time(0.5, {label="$var"} | logfmt | __error__=`` | unwrap latency | __error__=`` [$__interval]))'
+      )
+    ).toBe('{label="$var"} | logfmt | __error__=``');
+  });
+});

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

@@ -38,15 +38,7 @@ export function formatQuery(selector: string | undefined): string {
 export function getHighlighterExpressionsFromQuery(input: string): string[] {
   const results = [];
 
-  const tree = parser.parse(input);
-  const filters: SyntaxNode[] = [];
-  tree.iterate({
-    enter: ({ type, node }): void => {
-      if (type.id === LineFilter) {
-        filters.push(node);
-      }
-    },
-  });
+  const filters = getNodesFromQuery(input, [LineFilter]);
 
   for (let filter of filters) {
     const pipeExact = filter.getChild(Filter)?.getChild(PipeExact);
@@ -139,92 +131,73 @@ export function parseToNodeNamesArray(query: string): string[] {
   return queryParts;
 }
 
-export function isValidQuery(query: string): boolean {
-  let isValid = true;
+export function isQueryWithNode(query: string, nodeType: number): boolean {
+  let isQueryWithNode = false;
   const tree = parser.parse(query);
   tree.iterate({
     enter: ({ type }): false | void => {
-      if (type.id === ErrorId) {
-        isValid = false;
+      if (type.id === nodeType) {
+        isQueryWithNode = true;
+        return false;
       }
     },
   });
-  return isValid;
+  return isQueryWithNode;
 }
 
-export function isLogsQuery(query: string): boolean {
-  let isLogsQuery = true;
+export function getNodesFromQuery(query: string, nodeTypes: number[]): SyntaxNode[] {
+  const nodes: SyntaxNode[] = [];
   const tree = parser.parse(query);
   tree.iterate({
-    enter: ({ type }): false | void => {
-      if (type.id === MetricExpr) {
-        isLogsQuery = false;
+    enter: (node): false | void => {
+      if (nodeTypes.includes(node.type.id)) {
+        nodes.push(node.node);
       }
     },
   });
-  return isLogsQuery;
+  return nodes;
+}
+
+export function getNodeFromQuery(query: string, nodeType: number): SyntaxNode | undefined {
+  const nodes = getNodesFromQuery(query, [nodeType]);
+  return nodes.length > 0 ? nodes[0] : undefined;
+}
+
+export function isValidQuery(query: string): boolean {
+  return !isQueryWithNode(query, ErrorId);
+}
+
+export function isLogsQuery(query: string): boolean {
+  return !isQueryWithNode(query, MetricExpr);
 }
 
 export function isQueryWithParser(query: string): { queryWithParser: boolean; parserCount: number } {
-  let parserCount = 0;
-  const tree = parser.parse(query);
-  tree.iterate({
-    enter: ({ type }): false | void => {
-      if (type.id === LabelParser || type.id === JsonExpressionParser) {
-        parserCount++;
-      }
-    },
-  });
+  const nodes = getNodesFromQuery(query, [LabelParser, JsonExpressionParser]);
+  const parserCount = nodes.length;
   return { queryWithParser: parserCount > 0, parserCount };
 }
 
 export function getParserFromQuery(query: string): string | undefined {
-  const tree = parser.parse(query);
-  let logParser: string | undefined = undefined;
-  tree.iterate({
-    enter: (node: SyntaxNode): false | void => {
-      if (node.type.id === LabelParser || node.type.id === JsonExpressionParser) {
-        logParser = query.substring(node.from, node.to).trim();
-        return false;
-      }
-    },
-  });
-
-  return logParser;
+  const parsers = getNodesFromQuery(query, [LabelParser, JsonExpressionParser]);
+  return parsers.length > 0 ? query.substring(parsers[0].from, parsers[0].to).trim() : undefined;
 }
 
 export function isQueryPipelineErrorFiltering(query: string): boolean {
-  let isQueryPipelineErrorFiltering = false;
-  const tree = parser.parse(query);
-  tree.iterate({
-    enter: ({ type, node }): false | void => {
-      if (type.id === LabelFilter) {
-        const label = node.getChild(Matcher)?.getChild(Identifier);
-        if (label) {
-          const labelName = query.substring(label.from, label.to);
-          if (labelName === '__error__') {
-            isQueryPipelineErrorFiltering = true;
-          }
-        }
+  const labels = getNodesFromQuery(query, [LabelFilter]);
+  for (const node of labels) {
+    const label = node.getChild(Matcher)?.getChild(Identifier);
+    if (label) {
+      const labelName = query.substring(label.from, label.to);
+      if (labelName === '__error__') {
+        return true;
       }
-    },
-  });
-
-  return isQueryPipelineErrorFiltering;
+    }
+  }
+  return false;
 }
 
 export function isQueryWithLabelFormat(query: string): boolean {
-  let queryWithLabelFormat = false;
-  const tree = parser.parse(query);
-  tree.iterate({
-    enter: ({ type }): false | void => {
-      if (type.id === LabelFormatExpr) {
-        queryWithLabelFormat = true;
-        return false;
-      }
-    },
-  });
-  return queryWithLabelFormat;
+  return isQueryWithNode(query, LabelFormatExpr);
 }
 
 export function getLogQueryFromMetricsQuery(query: string): string {
@@ -232,92 +205,39 @@ export function getLogQueryFromMetricsQuery(query: string): string {
     return query;
   }
 
-  const tree = parser.parse(query);
-
   // Log query in metrics query composes of Selector & PipelineExpr
-  let selector = '';
-  tree.iterate({
-    enter: ({ type, from, to }): false | void => {
-      if (type.id === Selector) {
-        selector = query.substring(from, to);
-        return false;
-      }
-    },
-  });
+  const selectorNode = getNodeFromQuery(query, Selector);
+  if (!selectorNode) {
+    return query;
+  }
+  const selector = query.substring(selectorNode.from, selectorNode.to);
 
-  let pipelineExpr = '';
-  tree.iterate({
-    enter: ({ type, from, to }): false | void => {
-      if (type.id === PipelineExpr) {
-        pipelineExpr = query.substring(from, to);
-        return false;
-      }
-    },
-  });
+  const pipelineExprNode = getNodeFromQuery(query, PipelineExpr);
+  const pipelineExpr = pipelineExprNode ? query.substring(pipelineExprNode.from, pipelineExprNode.to) : '';
 
-  return selector + pipelineExpr;
+  return `${selector} ${pipelineExpr}`.trim();
 }
 
 export function isQueryWithLabelFilter(query: string): boolean {
-  const tree = parser.parse(query);
-  let hasLabelFilter = false;
-
-  tree.iterate({
-    enter: ({ type }): false | void => {
-      if (type.id === LabelFilter) {
-        hasLabelFilter = true;
-        return false;
-      }
-    },
-  });
-
-  return hasLabelFilter;
+  return isQueryWithNode(query, LabelFilter);
 }
 
 export function isQueryWithLineFilter(query: string): boolean {
-  const tree = parser.parse(query);
-  let queryWithLineFilter = false;
-
-  tree.iterate({
-    enter: ({ type }): false | void => {
-      if (type.id === LineFilter) {
-        queryWithLineFilter = true;
-        return false;
-      }
-    },
-  });
-
-  return queryWithLineFilter;
+  return isQueryWithNode(query, LineFilter);
 }
 
 export function isQueryWithDistinct(query: string): boolean {
-  let hasDistinct = false;
-  const tree = parser.parse(query);
-  tree.iterate({
-    enter: ({ type }): false | void => {
-      if (type.id === Distinct) {
-        hasDistinct = true;
-        return false;
-      }
-    },
-  });
-  return hasDistinct;
+  return isQueryWithNode(query, Distinct);
 }
 
 export function isQueryWithRangeVariable(query: string): boolean {
-  let hasRangeVariableDuration = false;
-  const tree = parser.parse(query);
-  tree.iterate({
-    enter: ({ type, from, to }): false | void => {
-      if (type.id === Range) {
-        if (query.substring(from, to).match(/\[\$__range(_s|_ms)?/)) {
-          hasRangeVariableDuration = true;
-          return false;
-        }
-      }
-    },
-  });
-  return hasRangeVariableDuration;
+  const rangeNodes = getNodesFromQuery(query, [Range]);
+  for (const node of rangeNodes) {
+    if (query.substring(node.from, node.to).match(/\[\$__range(_s|_ms)?/)) {
+      return true;
+    }
+  }
+  return false;
 }
 
 export function getStreamSelectorsFromQuery(query: string): string[] {