Loki: Fix producing correct log volume query for query with comments (#53254)

* Loki: Remove comments from log volume query

* Update, add more tests

* Update based on suggestions

* Update

* Update tests to use it.each

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

@@ -44,12 +44,18 @@ import { getTemplateSrv, TemplateSrv } from 'app/features/templating/template_sr
 import { serializeParams } from '../../../core/utils/fetch';
 import { renderLegendFormat } from '../prometheus/legend';
 
-import { addLabelFormatToQuery, addLabelToQuery, addNoPipelineErrorToQuery, addParserToQuery } from './addToQuery';
 import { transformBackendResult } from './backendResultTransformer';
 import { LokiAnnotationsQueryEditor } from './components/AnnotationsQueryEditor';
 import LanguageProvider from './language_provider';
 import { escapeLabelValueInSelector } from './language_utils';
 import { LiveStreams, LokiLiveTarget } from './live_streams';
+import {
+  addLabelFormatToQuery,
+  addLabelToQuery,
+  addNoPipelineErrorToQuery,
+  addParserToQuery,
+  removeCommentsFromQuery,
+} from './modifyQuery';
 import { getQueryHints } from './queryHints';
 import { getNormalizedLokiQuery, isLogsQuery, isValidQuery } from './query_utils';
 import { sortDataFrameByTime } from './sortDataFrame';
@@ -119,11 +125,12 @@ export class LokiDatasource
 
     const logsVolumeRequest = cloneDeep(request);
     logsVolumeRequest.targets = logsVolumeRequest.targets.filter(isQuerySuitable).map((target) => {
+      const query = removeCommentsFromQuery(target.expr);
       return {
         ...target,
         instant: false,
         volumeQuery: true,
-        expr: `sum by (level) (count_over_time(${target.expr}[$__interval]))`,
+        expr: `sum by (level) (count_over_time(${query}[$__interval]))`,
       };
     });
 

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

@@ -1,4 +1,10 @@
-import { addLabelFormatToQuery, addLabelToQuery, addNoPipelineErrorToQuery, addParserToQuery } from './addToQuery';
+import {
+  addLabelFormatToQuery,
+  addLabelToQuery,
+  addNoPipelineErrorToQuery,
+  addParserToQuery,
+  removeCommentsFromQuery,
+} from './modifyQuery';
 
 describe('addLabelToQuery()', () => {
   it('should add label to simple query', () => {
@@ -233,3 +239,41 @@ describe('addLabelFormatToQuery', () => {
     );
   });
 });
+
+describe('removeCommentsFromQuery', () => {
+  it.each`
+    query                                                                             | expectedResult
+    ${'{job="grafana"}#hello'}                                                        | ${'{job="grafana"}'}
+    ${'{job="grafana"} | logfmt #hello'}                                              | ${'{job="grafana"} | logfmt '}
+    ${'{job="grafana", bar="baz"} |="test" | logfmt | label_format level=lvl #hello'} | ${'{job="grafana", bar="baz"} |="test" | logfmt | label_format level=lvl '}
+  `('strips comments in log query:  {$query}', ({ query, expectedResult }) => {
+    expect(removeCommentsFromQuery(query)).toBe(expectedResult);
+  });
+
+  it.each`
+    query                                                                      | expectedResult
+    ${'{job="grafana"}'}                                                       | ${'{job="grafana"}'}
+    ${'{job="grafana"} | logfmt'}                                              | ${'{job="grafana"} | logfmt'}
+    ${'{job="grafana", bar="baz"} |="test" | logfmt | label_format level=lvl'} | ${'{job="grafana", bar="baz"} |="test" | logfmt | label_format level=lvl'}
+  `('returns original query if no comments in log query:  {$query}', ({ query, expectedResult }) => {
+    expect(removeCommentsFromQuery(query)).toBe(expectedResult);
+  });
+
+  it.each`
+    query                                                       | expectedResult
+    ${'count_over_time({job="grafana"}[10m])#hello'}            | ${'count_over_time({job="grafana"}[10m])'}
+    ${'count_over_time({job="grafana"} | logfmt[10m])#hello'}   | ${'count_over_time({job="grafana"} | logfmt[10m])'}
+    ${'rate({job="grafana"} | logfmt | foo="bar" [10m])#hello'} | ${'rate({job="grafana"} | logfmt | foo="bar" [10m])'}
+  `('strips comments in metrics query:  {$query}', ({ query, expectedResult }) => {
+    expect(removeCommentsFromQuery(query)).toBe(expectedResult);
+  });
+
+  it.each`
+    query                                                     | expectedResult
+    ${'count_over_time({job="grafana"}[10m])#hello'}          | ${'count_over_time({job="grafana"}[10m])'}
+    ${'count_over_time({job="grafana"} | logfmt[10m])#hello'} | ${'count_over_time({job="grafana"} | logfmt[10m])'}
+    ${'rate({job="grafana"} | logfmt | foo="bar" [10m])'}     | ${'rate({job="grafana"} | logfmt | foo="bar" [10m])'}
+  `('returns original query if no comments in metrics query:  {$query}', ({ query, expectedResult }) => {
+    expect(removeCommentsFromQuery(query)).toBe(expectedResult);
+  });
+});

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

@@ -1,6 +1,6 @@
 import { sortBy } from 'lodash';
 
-import { parser } from '@grafana/lezer-logql';
+import { LineComment, parser } from '@grafana/lezer-logql';
 
 import { QueryBuilderLabelFilter } from '../prometheus/querybuilder/shared/types';
 
@@ -90,6 +90,30 @@ export function addLabelFormatToQuery(query: string, labelFormat: { originalLabe
   return addLabelFormat(query, logQueryPositions, labelFormat);
 }
 
+/**
+ * Removes all comments from query.
+ * It uses  LogQL parser to find all LineComments and removes them.
+ */
+export function removeCommentsFromQuery(query: string): string {
+  const lineCommentPositions = getLineCommentPositions(query);
+
+  if (!lineCommentPositions.length) {
+    return query;
+  }
+
+  let newQuery = '';
+  let prev = 0;
+
+  for (let lineCommentPosition of lineCommentPositions) {
+    const beforeComment = query.substring(prev, lineCommentPosition.from);
+    const afterComment = query.substring(lineCommentPosition.to);
+
+    newQuery += beforeComment + afterComment;
+    prev = lineCommentPosition.to;
+  }
+  return newQuery;
+}
+
 /**
  * Parse the string and get all Selector positions in the query together with parsed representation of the
  * selector.
@@ -331,6 +355,20 @@ function addLabelFormat(
   return newQuery;
 }
 
+function getLineCommentPositions(query: string): Position[] {
+  const tree = parser.parse(query);
+  const positions: Position[] = [];
+  tree.iterate({
+    enter: (type, from, to, get): false | void => {
+      if (type.id === LineComment) {
+        positions.push({ from, to });
+        return false;
+      }
+    },
+  });
+  return positions;
+}
+
 /**
  * Check if label exists in the list of labels but ignore the operator.
  * @param labels