Loki: Don't split queries if they use $__range variables (#69852)

* Loki: Don't split queries if they use  variables

* Update public/app/plugins/datasource/loki/querySplitting.ts

Co-authored-by: Matias Chomicki <matyax@gmail.com>

---------

Co-authored-by: Matias Chomicki <matyax@gmail.com>

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

@@ -17,7 +17,7 @@ import { LoadingState } from '@grafana/schema';
 import { LokiDatasource } from './datasource';
 import { splitTimeRange as splitLogsTimeRange } from './logsTimeSplitting';
 import { splitTimeRange as splitMetricTimeRange } from './metricTimeSplitting';
-import { isLogsQuery, isQueryWithDistinct } from './queryUtils';
+import { isLogsQuery, isQueryWithDistinct, isQueryWithRangeVariable } from './queryUtils';
 import { combineResponses } from './responseUtils';
 import { trackGroupedQueries } from './tracking';
 import { LokiGroupedRequest, LokiQuery, LokiQueryType } from './types';
@@ -213,13 +213,19 @@ function getNextRequestPointers(requests: LokiGroupedRequest[], requestGroup: nu
   };
 }
 
-function querySupporstSplitting(query: LokiQuery) {
-  return query.queryType !== LokiQueryType.Instant && !isQueryWithDistinct(query.expr);
+function querySupportsSplitting(query: LokiQuery) {
+  return (
+    query.queryType !== LokiQueryType.Instant &&
+    !isQueryWithDistinct(query.expr) &&
+    // Queries with $__range variable should not be split because then the interpolated $__range variable is incorrect
+    // because it is interpolated on the backend with the split timeRange
+    !isQueryWithRangeVariable(query.expr)
+  );
 }
 
 export function runSplitQuery(datasource: LokiDatasource, request: DataQueryRequest<LokiQuery>) {
   const queries = request.targets.filter((query) => !query.hide);
-  const [nonSplittingQueries, normalQueries] = partition(queries, (query) => !querySupporstSplitting(query));
+  const [nonSplittingQueries, normalQueries] = partition(queries, (query) => !querySupportsSplitting(query));
   const [logQueries, metricQueries] = partition(normalQueries, (query) => isLogsQuery(query.expr));
 
   request.queryGroupId = uuidv4();

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

@@ -10,6 +10,7 @@ import {
   obfuscate,
   requestSupportsSplitting,
   isQueryWithDistinct,
+  isQueryWithRangeVariable,
 } from './queryUtils';
 import { LokiQuery, LokiQueryType } from './types';
 
@@ -294,6 +295,28 @@ describe('isQueryWithDistinct', () => {
   });
 });
 
+describe('isQueryWithRangeVariableDuration', () => {
+  it('identifies queries using $__range variable', () => {
+    expect(isQueryWithRangeVariable('rate({job="grafana"}[$__range])')).toBe(true);
+  });
+
+  it('identifies queries using $__range_s variable', () => {
+    expect(isQueryWithRangeVariable('rate({job="grafana"}[$__range_s])')).toBe(true);
+  });
+
+  it('identifies queries using $__range_ms variable', () => {
+    expect(isQueryWithRangeVariable('rate({job="grafana"}[$__range_ms])')).toBe(true);
+  });
+
+  it('does not return false positives', () => {
+    expect(isQueryWithRangeVariable('rate({job="grafana"} | logfmt | value="$__range" [5m])')).toBe(false);
+    expect(isQueryWithRangeVariable('rate({job="grafana"} | logfmt | value="[$__range]" [5m])')).toBe(false);
+    expect(isQueryWithRangeVariable('rate({job="grafana"} [$range])')).toBe(false);
+    expect(isQueryWithRangeVariable('rate({job="grafana"} [$_range])')).toBe(false);
+    expect(isQueryWithRangeVariable('rate({job="grafana"} [$_range_ms])')).toBe(false);
+  });
+});
+
 describe('getParserFromQuery', () => {
   it('returns no parser', () => {
     expect(getParserFromQuery('{job="grafana"}')).toBeUndefined();

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

@@ -18,6 +18,7 @@ import {
   Matcher,
   Identifier,
   Distinct,
+  Range,
 } from '@grafana/lezer-logql';
 import { DataQuery } from '@grafana/schema';
 
@@ -303,6 +304,22 @@ export function isQueryWithDistinct(query: string): boolean {
   return hasDistinct;
 }
 
+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;
+}
+
 export function getStreamSelectorsFromQuery(query: string): string[] {
   const labelMatcherPositions = getStreamSelectorPositions(query);