Loki: Exclude queries using DISTINCT from query splitting (#69377)

* Query utils: add function to identify distinct queries

* Chore: return false when the expected node is found

* Query splitting: exclude distinct queries from splitting

* Group queries excluded from splitting

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

@@ -182,6 +182,19 @@ describe('runSplitQuery()', () => {
         expect(datasource.runQuery).toHaveBeenCalledTimes(1);
       });
     });
+    test('Groups queries using distinct', async () => {
+      const request = getQueryOptions<LokiQuery>({
+        targets: [
+          { expr: '{a="b"} | distinct field', refId: 'A' },
+          { expr: 'count_over_time({c="d"} | distinct something [1m])', refId: 'B' },
+        ],
+        range,
+      });
+      await expect(runSplitQuery(datasource, request)).toEmitValuesWith(() => {
+        // Queries using distinct are omitted from splitting
+        expect(datasource.runQuery).toHaveBeenCalledTimes(1);
+      });
+    });
     test('Respects maxLines of logs queries', async () => {
       const { logFrameA } = getMockFrames();
       const request = getQueryOptions<LokiQuery>({
@@ -199,17 +212,18 @@ describe('runSplitQuery()', () => {
         expect(datasource.runQuery).toHaveBeenCalledTimes(4);
       });
     });
-    test('Groups multiple queries into logs, queries, and instant', async () => {
+    test('Groups multiple queries into logs, queries, instant, and distinct', async () => {
       const request = getQueryOptions<LokiQuery>({
         targets: [
           { expr: 'count_over_time({a="b"}[1m])', refId: 'A', queryType: LokiQueryType.Instant },
           { expr: '{c="d"}', refId: 'B' },
           { expr: 'count_over_time({c="d"}[1m])', refId: 'C' },
+          { expr: 'count_over_time({c="d"} | distinct id [1m])', refId: 'D' },
         ],
         range,
       });
       await expect(runSplitQuery(datasource, request)).toEmitValuesWith(() => {
-        // 3 days, 3 chunks, 3x Logs + 3x Metric + 1x Instant, 7 requests.
+        // 3 days, 3 chunks, 3x Logs + 3x Metric + (1x Instant | Distinct), 7 requests.
         expect(datasource.runQuery).toHaveBeenCalledTimes(7);
       });
     });

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

@@ -15,7 +15,7 @@ import { LoadingState } from '@grafana/schema';
 import { LokiDatasource } from './datasource';
 import { splitTimeRange as splitLogsTimeRange } from './logsTimeSplitting';
 import { splitTimeRange as splitMetricTimeRange } from './metricTimeSplitting';
-import { isLogsQuery } from './queryUtils';
+import { isLogsQuery, isQueryWithDistinct } from './queryUtils';
 import { combineResponses } from './responseUtils';
 import { trackGroupedQueries } from './tracking';
 import { LokiGroupedRequest, LokiQuery, LokiQueryType } from './types';
@@ -171,9 +171,13 @@ function getNextRequestPointers(requests: LokiGroupedRequest[], requestGroup: nu
   };
 }
 
+function querySupporstSplitting(query: LokiQuery) {
+  return query.queryType !== LokiQueryType.Instant && !isQueryWithDistinct(query.expr);
+}
+
 export function runSplitQuery(datasource: LokiDatasource, request: DataQueryRequest<LokiQuery>) {
   const queries = request.targets.filter((query) => !query.hide);
-  const [instantQueries, normalQueries] = partition(queries, (query) => query.queryType === LokiQueryType.Instant);
+  const [nonSplittingQueries, normalQueries] = partition(queries, (query) => !querySupporstSplitting(query));
   const [logQueries, metricQueries] = partition(normalQueries, (query) => isLogsQuery(query.expr));
 
   request.queryGroupId = uuidv4();
@@ -218,9 +222,9 @@ export function runSplitQuery(datasource: LokiDatasource, request: DataQueryRequ
     }
   }
 
-  if (instantQueries.length) {
+  if (nonSplittingQueries.length) {
     requests.push({
-      request: { ...request, targets: instantQueries },
+      request: { ...request, targets: nonSplittingQueries },
       partition: [request.range],
     });
   }

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

@@ -9,6 +9,7 @@ import {
   getParserFromQuery,
   obfuscate,
   requestSupportsSplitting,
+  isQueryWithDistinct,
 } from './queryUtils';
 import { LokiQuery, LokiQueryType } from './types';
 
@@ -281,6 +282,18 @@ describe('isQueryWithLabelFormat', () => {
   });
 });
 
+describe('isQueryWithDistinct', () => {
+  it('identifies queries using distinct', () => {
+    expect(isQueryWithDistinct('{job="grafana"} | distinct id')).toBe(true);
+    expect(isQueryWithDistinct('count_over_time({job="grafana"} | distinct id [1m])')).toBe(true);
+  });
+
+  it('does not return false positives', () => {
+    expect(isQueryWithDistinct('{label="distinct"} | logfmt')).toBe(false);
+    expect(isQueryWithDistinct('count_over_time({job="distinct"} | json [1m])')).toBe(false);
+  });
+});
+
 describe('getParserFromQuery', () => {
   it('returns no parser', () => {
     expect(getParserFromQuery('{job="grafana"}')).toBeUndefined();

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

@@ -17,6 +17,7 @@ import {
   MetricExpr,
   Matcher,
   Identifier,
+  Distinct,
 } from '@grafana/lezer-logql';
 import { DataQuery } from '@grafana/schema';
 
@@ -218,6 +219,7 @@ export function isQueryWithLabelFormat(query: string): boolean {
     enter: ({ type }): false | void => {
       if (type.id === LabelFormatExpr) {
         queryWithLabelFormat = true;
+        return false;
       }
     },
   });
@@ -260,10 +262,10 @@ export function isQueryWithLabelFilter(query: string): boolean {
   let hasLabelFilter = false;
 
   tree.iterate({
-    enter: ({ type, node }): false | void => {
+    enter: ({ type }): false | void => {
       if (type.id === LabelFilter) {
         hasLabelFilter = true;
-        return;
+        return false;
       }
     },
   });
@@ -279,7 +281,7 @@ export function isQueryWithLineFilter(query: string): boolean {
     enter: ({ type }): false | void => {
       if (type.id === LineFilter) {
         queryWithLineFilter = true;
-        return;
+        return false;
       }
     },
   });
@@ -287,6 +289,20 @@ export function isQueryWithLineFilter(query: string): boolean {
   return queryWithLineFilter;
 }
 
+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;
+}
+
 export function getStreamSelectorsFromQuery(query: string): string[] {
   const labelMatcherPositions = getStreamSelectorPositions(query);