Loki: Add option to define chunk duration per query (#64834)

* add query option to configure chunk ranges

* remove `isValidDuration` check

* Update public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilderOptions.tsx

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

* change to `chunkDuration`
added tests

* no need to call `toString`

---------

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

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

@@ -190,4 +190,74 @@ describe('runPartitionedQueries()', () => {
       });
     });
   });
+
+  describe('Splitting targets based on chunkDuration', () => {
+    const range1h = {
+      from: dateTime('2023-02-08T05:00:00.000Z'),
+      to: dateTime('2023-02-08T06:00:00.000Z'),
+      raw: {
+        from: dateTime('2023-02-08T05:00:00.000Z'),
+        to: dateTime('2023-02-08T06:00:00.000Z'),
+      },
+    };
+    beforeEach(() => {
+      jest.spyOn(datasource, 'runQuery').mockReturnValue(of({ data: [], refId: 'A' }));
+    });
+    test('with 30m chunkDuration runs 2 queries', async () => {
+      const request = getQueryOptions<LokiQuery>({
+        targets: [{ expr: '{a="b"}', refId: 'A', chunkDuration: '30m' }],
+        range: range1h,
+      });
+      await expect(runPartitionedQueries(datasource, request)).toEmitValuesWith(() => {
+        expect(datasource.runQuery).toHaveBeenCalledTimes(2);
+      });
+    });
+    test('with 1h chunkDuration runs 1 queries', async () => {
+      const request = getQueryOptions<LokiQuery>({
+        targets: [{ expr: '{a="b"}', refId: 'A', chunkDuration: '1h' }],
+        range: range1h,
+      });
+      await expect(runPartitionedQueries(datasource, request)).toEmitValuesWith(() => {
+        expect(datasource.runQuery).toHaveBeenCalledTimes(1);
+      });
+    });
+    test('with 1h chunkDuration and 2 targets runs 1 queries', async () => {
+      const request = getQueryOptions<LokiQuery>({
+        targets: [
+          { expr: '{a="b"}', refId: 'A', chunkDuration: '1h' },
+          { expr: '{a="b"}', refId: 'B', chunkDuration: '1h' },
+        ],
+        range: range1h,
+      });
+      await expect(runPartitionedQueries(datasource, request)).toEmitValuesWith(() => {
+        expect(datasource.runQuery).toHaveBeenCalledTimes(1);
+      });
+    });
+    test('with 1h/30m chunkDuration and 2 targets runs 3 queries', async () => {
+      const request = getQueryOptions<LokiQuery>({
+        targets: [
+          { expr: '{a="b"}', refId: 'A', chunkDuration: '1h' },
+          { expr: '{a="b"}', refId: 'B', chunkDuration: '30m' },
+        ],
+        range: range1h,
+      });
+      await expect(runPartitionedQueries(datasource, request)).toEmitValuesWith(() => {
+        // 2 x 30m + 1 x 1h
+        expect(datasource.runQuery).toHaveBeenCalledTimes(3);
+      });
+    });
+    test('with 1h/30m chunkDuration and 1 log and 2 metric target runs 3 queries', async () => {
+      const request = getQueryOptions<LokiQuery>({
+        targets: [
+          { expr: '{a="b"}', refId: 'A', chunkDuration: '1h' },
+          { expr: 'count_over_time({c="d"}[1m])', refId: 'C', chunkDuration: '30m' },
+        ],
+        range: range1h,
+      });
+      await expect(runPartitionedQueries(datasource, request)).toEmitValuesWith(() => {
+        // 2 x 30m + 1 x 1h
+        expect(datasource.runQuery).toHaveBeenCalledTimes(3);
+      });
+    });
+  });
 });

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

@@ -1,7 +1,14 @@
-import { partition } from 'lodash';
+import { groupBy, partition } from 'lodash';
-import { Subscriber, Observable, Subscription } from 'rxjs';
+import { Observable, Subscriber, Subscription } from 'rxjs';
 
-import { DataQueryRequest, DataQueryResponse, dateTime, TimeRange } from '@grafana/data';
+import {
+  DataQueryRequest,
+  DataQueryResponse,
+  dateTime,
+  durationToMilliseconds,
+  parseDuration,
+  TimeRange,
+} from '@grafana/data';
 import { LoadingState } from '@grafana/schema';
 
 import { LokiDatasource } from './datasource';
@@ -10,24 +17,12 @@ import { getRangeChunks as getMetricRangeChunks } from './metricTimeSplit';
 import { combineResponses, isLogsQuery } from './queryUtils';
 import { LokiQuery, LokiQueryType } from './types';
 
-declare global {
-  interface Window {
-    lokiChunkDuration: number;
-  }
-}
-
-/**
- * Purposely exposing it to support doing tests without needing to update the repo.
- * TODO: remove.
- * Hardcoded to 1 day.
- */
-window.lokiChunkDuration = 24 * 60 * 60 * 1000;
-
 export function partitionTimeRange(
   isLogsQuery: boolean,
   originalTimeRange: TimeRange,
   intervalMs: number,
-  resolution: number
+  resolution: number,
+  duration: number
 ): TimeRange[] {
   // the `step` value that will be finally sent to Loki is rougly the same as `intervalMs`,
   // but there are some complications.
@@ -41,8 +36,6 @@ export function partitionTimeRange(
   const safeStep = Math.ceil((end - start) / 11000);
   const step = Math.max(intervalMs * resolution, safeStep);
 
-  const duration = window.lokiChunkDuration;
-
   const ranges = isLogsQuery
     ? getLogsRangeChunks(start, end, duration)
     : getMetricRangeChunks(start, end, step, duration);
@@ -179,19 +172,41 @@ export function runPartitionedQueries(datasource: LokiDatasource, request: DataQ
   const [instantQueries, normalQueries] = partition(queries, (query) => query.queryType === LokiQueryType.Instant);
   const [logQueries, metricQueries] = partition(normalQueries, (query) => isLogsQuery(query.expr));
 
+  const oneDayMs = 24 * 60 * 60 * 1000;
+  const rangePartitionedLogQueries = groupBy(logQueries, (query) =>
+    query.chunkDuration ? durationToMilliseconds(parseDuration(query.chunkDuration)) : oneDayMs
+  );
+  const rangePartitionedMetricQueries = groupBy(metricQueries, (query) =>
+    query.chunkDuration ? durationToMilliseconds(parseDuration(query.chunkDuration)) : oneDayMs
+  );
+
   const requests: LokiGroupedRequest = [];
-  if (logQueries.length) {
+  for (const [chunkRangeMs, queries] of Object.entries(rangePartitionedLogQueries)) {
     requests.push({
-      request: { ...request, targets: logQueries },
+      request: { ...request, targets: queries },
-      partition: partitionTimeRange(true, request.range, request.intervalMs, logQueries[0].resolution ?? 1),
+      partition: partitionTimeRange(
+        true,
+        request.range,
+        request.intervalMs,
+        queries[0].resolution ?? 1,
+        Number(chunkRangeMs)
+      ),
     });
   }
-  if (metricQueries.length) {
+
+  for (const [chunkRangeMs, queries] of Object.entries(rangePartitionedMetricQueries)) {
     requests.push({
-      request: { ...request, targets: metricQueries },
+      request: { ...request, targets: queries },
-      partition: partitionTimeRange(false, request.range, request.intervalMs, metricQueries[0].resolution ?? 1),
+      partition: partitionTimeRange(
+        false,
+        request.range,
+        request.intervalMs,
+        queries[0].resolution ?? 1,
+        Number(chunkRangeMs)
+      ),
     });
   }
+
   if (instantQueries.length) {
     requests.push({
       request: { ...request, targets: instantQueries },

public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilderOptions.tsx

@@ -1,10 +1,10 @@
 import React, { useEffect, useState } from 'react';
 import { usePrevious } from 'react-use';
 
-import { CoreApp, SelectableValue } from '@grafana/data';
+import { CoreApp, isValidDuration, SelectableValue } from '@grafana/data';
 import { EditorField, EditorRow } from '@grafana/experimental';
-import { reportInteraction } from '@grafana/runtime';
+import { config, reportInteraction } from '@grafana/runtime';
-import { RadioButtonGroup, Select, AutoSizeInput } from '@grafana/ui';
+import { AutoSizeInput, RadioButtonGroup, Select } from '@grafana/ui';
 import { QueryOptionGroup } from 'app/plugins/datasource/prometheus/querybuilder/shared/QueryOptionGroup';
 
 import { preprocessMaxLines, queryTypeOptions, RESOLUTION_OPTIONS } from '../../components/LokiOptionFields';
@@ -24,6 +24,7 @@ export interface Props {
 export const LokiQueryBuilderOptions = React.memo<Props>(
   ({ app, query, onChange, onRunQuery, maxLines, datasource }) => {
     const [queryStats, setQueryStats] = useState<QueryStats>();
+    const [chunkRangeValid, setChunkRangeValid] = useState(true);
     const prevQuery = usePrevious(query);
 
     const onQueryTypeChange = (value: LokiQueryType) => {
@@ -40,6 +41,17 @@ export const LokiQueryBuilderOptions = React.memo<Props>(
       onRunQuery();
     };
 
+    const onChunkRangeChange = (evt: React.FormEvent<HTMLInputElement>) => {
+      const value = evt.currentTarget.value;
+      if (!isValidDuration(value)) {
+        setChunkRangeValid(false);
+        return;
+      }
+      setChunkRangeValid(true);
+      onChange({ ...query, chunkDuration: value });
+      onRunQuery();
+    };
+
     const onLegendFormatChanged = (evt: React.FormEvent<HTMLInputElement>) => {
       onChange({ ...query, legendFormat: evt.currentTarget.value });
       onRunQuery();
@@ -119,6 +131,21 @@ export const LokiQueryBuilderOptions = React.memo<Props>(
               aria-label="Select resolution"
             />
           </EditorField>
+          {config.featureToggles.lokiQuerySplitting && (
+            <EditorField
+              label="Chunk Duration"
+              tooltip="Defines the duration of a single query chunk when query chunking is used."
+            >
+              <AutoSizeInput
+                minWidth={14}
+                type="string"
+                min={0}
+                defaultValue={query.chunkDuration ?? '1d'}
+                onCommitChange={onChunkRangeChange}
+                invalid={!chunkRangeValid}
+              />
+            </EditorField>
+          )}
         </QueryOptionGroup>
       </EditorRow>
     );

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

@@ -35,6 +35,12 @@ export interface LokiQuery extends LokiQueryFromSchema {
   // the temporary fix (until this gets improved in the codegen), is to
   // override it here
   queryType?: LokiQueryType;
+
+  /**
+   * This is a property for the experimental query splitting feature.
+   * @experimental
+   */
+  chunkDuration?: string;
 }
 
 export interface LokiOptions extends DataSourceJsonData {