Loki: Fix timerange for query stats request (#72193)

* call query stats with whole query * call stats when query type changes * fix stats request timerange * feedback: extract functions from component * feedback: remove "or string" types * feedback: use ds timerange picker end * partial: fix broken test * changes * update test comment * rename variable * add comment explain message prop * make getStatsTimeRange a datasource method * update tests for getStatsTimeRange * make getStats a datasource method * update failing tests * update failing tests * remove import * Update public/app/plugins/datasource/loki/datasource.test.ts Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com> --------- Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>
2023-09-20 16:47:21 +01:00
parent 24a1f12826
commit 9def0d2305
14 changed files with 266 additions and 144 deletions
@@ -21,7 +21,7 @@ import { buildVisualQueryFromString } from '../querybuilder/parsing';
 import { changeEditorMode, getQueryWithDefaults } from '../querybuilder/state';
 import { LokiQuery, QueryStats } from '../types';

-import { getStats, shouldUpdateStats } from './stats';
+import { shouldUpdateStats } from './stats';
 import { LokiQueryEditorProps } from './types';

 export const testIds = {
@@ -45,7 +45,8 @@ export const LokiQueryEditor = React.memo<LokiQueryEditorProps>((props) => {
  if (config.featureToggles.lokiPredefinedOperations && !query.expr && predefinedOperations) {
    query.expr = `{} ${predefinedOperations}`;
  }
-  const previousQuery = usePrevious(query.expr);
+  const previousQueryExpr = usePrevious(query.expr);
+  const previousQueryType = usePrevious(query.queryType);

  // This should be filled in from the defaults by now.
  const editorMode = query.editorMode!;
@@ -96,15 +97,22 @@ export const LokiQueryEditor = React.memo<LokiQueryEditorProps>((props) => {
  };

  useEffect(() => {
-    const update = shouldUpdateStats(query.expr, previousQuery, timerange, previousTimerange);
+    const update = shouldUpdateStats(
+      query.expr,
+      previousQueryExpr,
+      timerange,
+      previousTimerange,
+      query.queryType,
+      previousQueryType
+    );
    if (update) {
      const makeAsyncRequest = async () => {
-        const stats = await getStats(datasource, query.expr);
+        const stats = await datasource.getStats(query);
        setQueryStats(stats);
      };
      makeAsyncRequest();
    }
-  }, [datasource, timerange, previousTimerange, query, previousQuery, setQueryStats]);
+  }, [datasource, timerange, previousTimerange, query, previousQueryExpr, previousQueryType, setQueryStats]);

  return (
    <>
@@ -180,13 +188,7 @@ export const LokiQueryEditor = React.memo<LokiQueryEditorProps>((props) => {
      <Space v={0.5} />
      <EditorRows>
        {editorMode === QueryEditorMode.Code && (
-          <LokiQueryCodeEditor
-            {...props}
-            query={query}
-            onChange={onChangeInternal}
-            showExplain={explain}
-            setQueryStats={setQueryStats}
-          />
+          <LokiQueryCodeEditor {...props} query={query} onChange={onChangeInternal} showExplain={explain} />
        )}
        {editorMode === QueryEditorMode.Builder && (
          <LokiQueryBuilderContainer
@@ -4,25 +4,15 @@ import React from 'react';

 import { CoreApp } from '@grafana/data';

-import { LokiDatasource } from '../datasource';
+import { createLokiDatasource } from '../mocks';

 import { testIds as regularTestIds } from './LokiQueryEditor';
 import { LokiQueryEditorByApp } from './LokiQueryEditorByApp';
 import { testIds as alertingTestIds } from './LokiQueryEditorForAlerting';

 function setup(app: CoreApp): RenderResult {
-  const dataSource = {
-    languageProvider: {
-      start: () => Promise.resolve([]),
-      getSyntax: () => {},
-      getLabelKeys: () => [],
-      metrics: [],
-    },
-    getQueryHints: () => [],
-    getDataSamples: () => [],
-    maxLines: 20,
-    getTimeRange: jest.fn(),
-  } as unknown as LokiDatasource;
+  const dataSource = createLokiDatasource();
+  dataSource.metadataRequest = jest.fn();

  return render(
    <LokiQueryEditorByApp
@@ -12,7 +12,6 @@ export interface LokiQueryFieldProps extends QueryEditorProps<LokiDatasource, Lo
  ExtraFieldElement?: ReactNode;
  placeholder?: string;
  'data-testid'?: string;
-  onQueryType?: (query: string) => void;
 }

 interface LokiQueryFieldState {
@@ -66,7 +65,7 @@ export class LokiQueryField extends React.PureComponent<LokiQueryFieldProps, Lok
  };

  render() {
-    const { ExtraFieldElement, query, datasource, history, onRunQuery, onQueryType } = this.props;
+    const { ExtraFieldElement, query, datasource, history, onRunQuery } = this.props;
    const placeholder = this.props.placeholder ?? 'Enter a Loki query (run with Shift+Enter)';

    return (
@@ -83,7 +82,6 @@ export class LokiQueryField extends React.PureComponent<LokiQueryFieldProps, Lok
              onRunQuery={onRunQuery}
              initialValue={query.expr ?? ''}
              placeholder={placeholder}
-              onQueryType={onQueryType}
            />
          </div>
        </div>
@@ -6,7 +6,12 @@ import { createLokiDatasource } from '../../mocks';
 import MonacoQueryField from './MonacoQueryField';
 import { Props } from './MonacoQueryFieldProps';

-function renderComponent({ initialValue = '', onRunQuery = jest.fn(), onBlur = jest.fn() }: Partial<Props> = {}) {
+function renderComponent({
+  initialValue = '',
+  onRunQuery = jest.fn(),
+  onBlur = jest.fn(),
+  onChange = jest.fn(),
+}: Partial<Props> = {}) {
  const datasource = createLokiDatasource();

  render(
@@ -16,6 +21,7 @@ function renderComponent({ initialValue = '', onRunQuery = jest.fn(), onBlur = j
      history={[]}
      onRunQuery={onRunQuery}
      onBlur={onBlur}
+      onChange={onChange}
      placeholder="Enter a Loki query (run with Shift+Enter)"
    />
  );
@@ -91,15 +91,7 @@ const getStyles = (theme: GrafanaTheme2, placeholder: string) => {
  };
 };

-const MonacoQueryField = ({
-  history,
-  onBlur,
-  onRunQuery,
-  initialValue,
-  datasource,
-  placeholder,
-  onQueryType,
-}: Props) => {
+const MonacoQueryField = ({ history, onBlur, onRunQuery, initialValue, datasource, placeholder, onChange }: Props) => {
  const id = uuidv4();
  // we need only one instance of `overrideServices` during the lifetime of the react component
  const overrideServicesRef = useRef(getOverrideServices());
@@ -151,11 +143,7 @@ const MonacoQueryField = ({
  };

  const onTypeDebounced = debounce(async (query: string) => {
-    if (!onQueryType) {
-      return;
-    }
-
-    onQueryType(query);
+    onChange(query);
  }, 1000);

  return (
@@ -14,5 +14,5 @@ export type Props = {
  onBlur: (value: string) => void;
  placeholder: string;
  datasource: LokiDatasource;
-  onQueryType?: (query: string) => void;
+  onChange: (query: string) => void;
 };
@@ -23,5 +23,5 @@ export const MonacoQueryFieldWrapper = (props: Props) => {
    onChange(value);
  };

-  return <MonacoQueryFieldLazy onRunQuery={handleRunQuery} onBlur={handleBlur} {...rest} />;
+  return <MonacoQueryFieldLazy onRunQuery={handleRunQuery} onBlur={handleBlur} onChange={onChange} {...rest} />;
 };
@@ -1,15 +1,18 @@
 import { dateTime, getDefaultTimeRange } from '@grafana/data';

-import { createLokiDatasource } from '../mocks';
+import { LokiQueryType } from '../types';

-import { getStats, shouldUpdateStats } from './stats';
+import { shouldUpdateStats } from './stats';

 describe('shouldUpdateStats', () => {
  const timerange = getDefaultTimeRange();
+  let queryType = LokiQueryType.Range;
+  let prevQueryType = LokiQueryType.Range;
+
  it('should return true if the query has changed', () => {
    const query = '{job="grafana"}';
    const prevQuery = '{job="not-grafana"}';
-    expect(shouldUpdateStats(query, prevQuery, timerange, timerange)).toBe(true);
+    expect(shouldUpdateStats(query, prevQuery, timerange, timerange, queryType, prevQueryType)).toBe(true);
  });

  it('should return true if the timerange has changed', () => {
@@ -17,25 +20,25 @@ describe('shouldUpdateStats', () => {
    const prevQuery = '{job="grafana"}';
    timerange.raw.from = 'now-14h';
    const prevTimerange = getDefaultTimeRange();
-    expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange)).toBe(true);
+    expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange, queryType, prevQueryType)).toBe(true);
  });

  it('should return true if the previous query was undefined', () => {
    const query = '{job="grafana"}';
    const prevQuery = undefined;
-    expect(shouldUpdateStats(query, prevQuery, timerange, timerange)).toBe(true);
+    expect(shouldUpdateStats(query, prevQuery, timerange, timerange, queryType, prevQueryType)).toBe(true);
  });

  it('should return true if the query really changed, otherwise false', () => {
    const prevQuery = '{job="grafana"}';
    const query = `${prevQuery} `;
-    expect(shouldUpdateStats(query, prevQuery, timerange, timerange)).toBe(false);
+    expect(shouldUpdateStats(query, prevQuery, timerange, timerange, queryType, prevQueryType)).toBe(false);
  });

  it('should return false if the query and timerange have not changed', () => {
    const query = '{job="grafana"}';
    const prevQuery = '{job="grafana"}';
-    expect(shouldUpdateStats(query, prevQuery, timerange, timerange)).toBe(false);
+    expect(shouldUpdateStats(query, prevQuery, timerange, timerange, queryType, prevQueryType)).toBe(false);
  });

  it('should return false if the query and timerange with absolute and relative mixed have not changed', () => {
@@ -46,57 +49,13 @@ describe('shouldUpdateStats', () => {

    const prevTimerange = getDefaultTimeRange();
    prevTimerange.raw.from = now;
-    expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange)).toBe(false);
+    expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange, queryType, prevQueryType)).toBe(false);
  });
-});

-describe('makeStatsRequest', () => {
-  const datasource = createLokiDatasource();
-
-  it('should return null if there is no query', () => {
-    const query = '';
-    expect(getStats(datasource, query)).resolves.toBe(null);
-  });
-
-  it('should return null if the query is invalid', () => {
-    const query = '{job="grafana",';
-    expect(getStats(datasource, query)).resolves.toBe(null);
-  });
-
-  it('should return null if the response has no data', () => {
-    const query = '{job="grafana"}';
-    datasource.getQueryStats = jest.fn().mockResolvedValue({ streams: 0, chunks: 0, bytes: 0, entries: 0 });
-    expect(getStats(datasource, query)).resolves.toBe(null);
-  });
-
-  it('should return the stats if the response has data', () => {
+  it('should return true if the query type has changed', () => {
    const query = '{job="grafana"}';
-
-    datasource.getQueryStats = jest
-      .fn()
-      .mockResolvedValue({ streams: 1, chunks: 12611, bytes: 12913664, entries: 78344 });
-    expect(getStats(datasource, query)).resolves.toEqual({
-      streams: 1,
-      chunks: 12611,
-      bytes: 12913664,
-      entries: 78344,
-    });
-  });
-
-  it('should support queries with variables', () => {
-    const query = 'count_over_time({job="grafana"}[$__interval])';
-
-    datasource.interpolateString = jest
-      .fn()
-      .mockImplementationOnce((value: string) => value.replace('$__interval', '1h'));
-    datasource.getQueryStats = jest
-      .fn()
-      .mockResolvedValue({ streams: 1, chunks: 12611, bytes: 12913664, entries: 78344 });
-    expect(getStats(datasource, query)).resolves.toEqual({
-      streams: 1,
-      chunks: 12611,
-      bytes: 12913664,
-      entries: 78344,
-    });
+    const prevQuery = '{job="grafana"}';
+    prevQueryType = LokiQueryType.Instant;
+    expect(shouldUpdateStats(query, prevQuery, timerange, timerange, queryType, prevQueryType)).toBe(true);
  });
 });
@@ -1,21 +1,6 @@
 import { DateTime, isDateTime, TimeRange } from '@grafana/data';

-import { LokiDatasource } from '../datasource';
-import { QueryStats } from '../types';
-
-export async function getStats(datasource: LokiDatasource, query: string): Promise<QueryStats | null> {
-  if (!query) {
-    return null;
-  }
-
-  const response = await datasource.getQueryStats(query);
-
-  if (!response) {
-    return null;
-  }
-
-  return Object.values(response).every((v) => v === 0) ? null : response;
-}
+import { LokiQueryType } from '../types';

 /**
 * This function compares two time values. If the first is absolute, it compares them using `DateTime.isSame`.
@@ -37,9 +22,11 @@ export function shouldUpdateStats(
  query: string,
  prevQuery: string | undefined,
  timerange: TimeRange,
-  prevTimerange: TimeRange | undefined
+  prevTimerange: TimeRange | undefined,
+  queryType: LokiQueryType | undefined,
+  prevQueryType: LokiQueryType | undefined
 ): boolean {
-  if (prevQuery === undefined || query.trim() !== prevQuery.trim()) {
+  if (prevQuery === undefined || query.trim() !== prevQuery.trim() || queryType !== prevQueryType) {
    return true;
  }

@@ -1288,21 +1288,26 @@ describe('LokiDatasource', () => {

  describe('getQueryStats', () => {
    let ds: LokiDatasource;
+    let query: LokiQuery;
    beforeEach(() => {
      ds = createLokiDatasource(templateSrvStub);
      ds.statsMetadataRequest = jest.fn().mockResolvedValue({ streams: 1, chunks: 1, bytes: 1, entries: 1 });
      ds.interpolateString = jest.fn().mockImplementation((value: string) => value.replace('$__interval', '1m'));
+
+      query = { refId: 'A', expr: '', queryType: LokiQueryType.Range };
    });

    it('uses statsMetadataRequest', async () => {
-      const result = await ds.getQueryStats('{foo="bar"}');
+      query.expr = '{foo="bar"}';
+      const result = await ds.getQueryStats(query);

      expect(ds.statsMetadataRequest).toHaveBeenCalled();
      expect(result).toEqual({ streams: 1, chunks: 1, bytes: 1, entries: 1 });
    });

    it('supports queries with template variables', async () => {
-      const result = await ds.getQueryStats('rate({instance="server\\1"}[$__interval])');
+      query.expr = 'rate({instance="server\\1"}[$__interval])';
+      const result = await ds.getQueryStats(query);

      expect(result).toEqual({
        streams: 1,
@@ -1313,14 +1318,16 @@ describe('LokiDatasource', () => {
    });

    it('does not call stats if the query is invalid', async () => {
-      const result = await ds.getQueryStats('rate({label="value"}');
+      query.expr = 'rate({label="value"}';
+      const result = await ds.getQueryStats(query);

      expect(ds.statsMetadataRequest).not.toHaveBeenCalled();
      expect(result).toBe(undefined);
    });

    it('combines the stats of each label matcher', async () => {
-      const result = await ds.getQueryStats('count_over_time({foo="bar"}[1m]) + count_over_time({test="test"}[1m])');
+      query.expr = 'count_over_time({foo="bar"}[1m]) + count_over_time({test="test"}[1m])';
+      const result = await ds.getQueryStats(query);

      expect(ds.statsMetadataRequest).toHaveBeenCalled();
      expect(result).toEqual({ streams: 2, chunks: 2, bytes: 2, entries: 2 });
@@ -1424,6 +1431,122 @@ describe('applyTemplateVariables', () => {
      );
    });
  });
+
+  describe('getStatsTimeRange', () => {
+    let query: LokiQuery;
+    let datasource: LokiDatasource;
+
+    beforeEach(() => {
+      query = { refId: 'A', expr: '', queryType: LokiQueryType.Range };
+      datasource = createLokiDatasource();
+
+      datasource.getTimeRangeParams = jest.fn().mockReturnValue({
+        start: 1672552800000000000, // 01 Jan 2023 06:00:00 GMT
+        end: 1672639200000000000, //   02 Jan 2023 06:00:00 GMT
+      });
+    });
+
+    it('should return the ds picker timerange for a logs query with range type', () => {
+      // log queries with range type should request the ds picker timerange
+      // in this case (1 day)
+      query.expr = '{job="grafana"}';
+
+      expect(datasource.getStatsTimeRange(query, 0)).toEqual({
+        start: 1672552800000000000, // 01 Jan 2023 06:00:00 GMT
+        end: 1672639200000000000, //   02 Jan 2023 06:00:00 GMT
+      });
+    });
+
+    it('should return nothing for a logs query with instant type', () => {
+      // log queries with instant type should be invalid.
+      query.queryType = LokiQueryType.Instant;
+      query.expr = '{job="grafana"}';
+
+      expect(datasource.getStatsTimeRange(query, 0)).toEqual({
+        start: undefined,
+        end: undefined,
+      });
+    });
+
+    it('should return the ds picker timerange', () => {
+      // metric queries with range type should request ds picker timerange
+      // in this case (1 day)
+      query.expr = 'rate({job="grafana"}[5m])';
+
+      expect(datasource.getStatsTimeRange(query, 0)).toEqual({
+        start: 1672552800000000000, // 01 Jan 2023 06:00:00 GMT
+        end: 1672639200000000000, //   02 Jan 2023 06:00:00 GMT
+      });
+    });
+
+    it('should return the range duration for an instant metric query', () => {
+      // metric queries with instant type should request range duration
+      // in this case (5 minutes)
+      query.queryType = LokiQueryType.Instant;
+      query.expr = 'rate({job="grafana"}[5m])';
+
+      expect(datasource.getStatsTimeRange(query, 0)).toEqual({
+        start: 1672638900000000000, // 02 Jan 2023 05:55:00 GMT
+        end: 1672639200000000000, //   02 Jan 2023 06:00:00 GMT
+      });
+    });
+  });
+});
+
+describe('makeStatsRequest', () => {
+  const datasource = createLokiDatasource();
+  let query: LokiQuery;
+
+  beforeEach(() => {
+    query = { refId: 'A', expr: '', queryType: LokiQueryType.Range };
+  });
+
+  it('should return null if there is no query', () => {
+    query.expr = '';
+    expect(datasource.getStats(query)).resolves.toBe(null);
+  });
+
+  it('should return null if the query is invalid', () => {
+    query.expr = '{job="grafana",';
+    expect(datasource.getStats(query)).resolves.toBe(null);
+  });
+
+  it('should return null if the response has no data', () => {
+    query.expr = '{job="grafana"}';
+    datasource.getQueryStats = jest.fn().mockResolvedValue({ streams: 0, chunks: 0, bytes: 0, entries: 0 });
+    expect(datasource.getStats(query)).resolves.toBe(null);
+  });
+
+  it('should return the stats if the response has data', () => {
+    query.expr = '{job="grafana"}';
+
+    datasource.getQueryStats = jest
+      .fn()
+      .mockResolvedValue({ streams: 1, chunks: 12611, bytes: 12913664, entries: 78344 });
+    expect(datasource.getStats(query)).resolves.toEqual({
+      streams: 1,
+      chunks: 12611,
+      bytes: 12913664,
+      entries: 78344,
+    });
+  });
+
+  it('should support queries with variables', () => {
+    query.expr = 'count_over_time({job="grafana"}[$__interval])';
+
+    datasource.interpolateString = jest
+      .fn()
+      .mockImplementationOnce((value: string) => value.replace('$__interval', '1h'));
+    datasource.getQueryStats = jest
+      .fn()
+      .mockResolvedValue({ streams: 1, chunks: 12611, bytes: 12913664, entries: 78344 });
+    expect(datasource.getStats(query)).resolves.toEqual({
+      streams: 1,
+      chunks: 12611,
+      bytes: 12913664,
+      entries: 78344,
+    });
+  });
 });

 describe('getTimeRange*()', () => {
@@ -38,6 +38,8 @@ import {
  QueryFilterOptions,
  renderLegendFormat,
 } from '@grafana/data';
+import { intervalToMs } from '@grafana/data/src/datetime/rangeutil';
+import { Duration } from '@grafana/lezer-logql';
 import { BackendSrvRequest, config, DataSourceWithBackend, FetchError } from '@grafana/runtime';
 import { DataQuery } from '@grafana/schema';
 import { convertToWebSocketUrl } from 'app/core/utils/explore';
@@ -77,6 +79,7 @@ import { runSplitQuery } from './querySplitting';
 import {
  getLogQueryFromMetricsQuery,
  getLokiQueryFromDataQuery,
+  getNodesFromQuery,
  getNormalizedLokiQuery,
  getStreamSelectorsFromQuery,
  isLogsQuery,
@@ -454,22 +457,30 @@ export class LokiDatasource
    return await this.getResource(url, params, options);
  }

-  async getQueryStats(query: string): Promise<QueryStats | undefined> {
+  async getQueryStats(query: LokiQuery): Promise<QueryStats | undefined> {
    // if query is invalid, clear stats, and don't request
-    if (isQueryWithError(this.interpolateString(query, placeHolderScopedVars))) {
+    if (isQueryWithError(this.interpolateString(query.expr, placeHolderScopedVars))) {
      return undefined;
    }

-    const { start, end } = this.getTimeRangeParams();
-    const labelMatchers = getStreamSelectorsFromQuery(query);
-
+    const labelMatchers = getStreamSelectorsFromQuery(query.expr);
    let statsForAll: QueryStats = { streams: 0, chunks: 0, bytes: 0, entries: 0 };

-    for (const labelMatcher of labelMatchers) {
+    for (const idx in labelMatchers) {
+      const { start, end } = this.getStatsTimeRange(query, Number(idx));
+
+      if (start === undefined || end === undefined) {
+        return { streams: 0, chunks: 0, bytes: 0, entries: 0, message: 'Query size estimate not available.' };
+      }
+
      try {
        const data = await this.statsMetadataRequest(
          'index/stats',
-          { query: labelMatcher, start, end },
+          {
+            query: labelMatchers[idx],
+            start: start,
+            end: end,
+          },
          { showErrorAlert: false }
        );

@@ -487,6 +498,60 @@ export class LokiDatasource
    return statsForAll;
  }

+  getStatsTimeRange(query: LokiQuery, idx: number): { start: number | undefined; end: number | undefined } {
+    let start: number, end: number;
+    const NS_IN_MS = 1000000;
+    const durationNodes = getNodesFromQuery(query.expr, [Duration]);
+    const durations = durationNodes.map((d) => query.expr.substring(d.from, d.to));
+
+    if (isLogsQuery(query.expr)) {
+      // logs query with instant type can not be estimated
+      if (query.queryType === LokiQueryType.Instant) {
+        return { start: undefined, end: undefined };
+      }
+      // logs query with range type
+      return this.getTimeRangeParams();
+    }
+
+    if (query.queryType === LokiQueryType.Instant) {
+      // metric query with instant type
+
+      if (!!durations[idx]) {
+        // if query has a duration e.g. [1m]
+        end = this.getTimeRangeParams().end;
+        start = end - intervalToMs(durations[idx]) * NS_IN_MS;
+        return { start, end };
+      } else {
+        // if query has no duration e.g. [$__interval]
+
+        if (/(\$__auto|\$__range)/.test(query.expr)) {
+          // if $__auto or $__range is used, we can estimate the time range using the selected range
+          return this.getTimeRangeParams();
+        }
+
+        // otherwise we cant estimate the time range
+        return { start: undefined, end: undefined };
+      }
+    }
+
+    // metric query with range type
+    return this.getTimeRangeParams();
+  }
+
+  async getStats(query: LokiQuery): Promise<QueryStats | null> {
+    if (!query) {
+      return null;
+    }
+
+    const response = await this.getQueryStats(query);
+
+    if (!response) {
+      return null;
+    }
+
+    return Object.values(response).every((v) => v === 0) ? null : response;
+  }
+
  async metricFindQuery(query: LokiVariableQuery | string) {
    if (!query) {
      return Promise.resolve([]);
@@ -8,16 +8,13 @@ import { getModKey } from 'app/core/utils/browser';

 import { testIds } from '../../components/LokiQueryEditor';
 import { LokiQueryField } from '../../components/LokiQueryField';
-import { getStats } from '../../components/stats';
 import { LokiQueryEditorProps } from '../../components/types';
 import { formatLogqlQuery } from '../../queryUtils';
-import { QueryStats } from '../../types';

 import { LokiQueryBuilderExplained } from './LokiQueryBuilderExplained';

 type Props = LokiQueryEditorProps & {
  showExplain: boolean;
-  setQueryStats: React.Dispatch<React.SetStateAction<QueryStats | null>>;
 };

 export function LokiQueryCodeEditor({
@@ -30,7 +27,6 @@ export function LokiQueryCodeEditor({
  app,
  showExplain,
  history,
-  setQueryStats,
 }: Props) {
  const styles = useStyles2(getStyles);

@@ -49,10 +45,6 @@ export function LokiQueryCodeEditor({
        data={data}
        app={app}
        data-testid={testIds.editor}
-        onQueryType={async (query: string) => {
-          const stats = await getStats(datasource, query);
-          setQueryStats(stats);
-        }}
        ExtraFieldElement={
          <>
            {lokiFormatQuery && (
@@ -152,6 +152,8 @@ export interface QueryStats {
  chunks: number;
  bytes: number;
  entries: number;
+  // The error message displayed in the UI when we cant estimate the size of the query.
+  message?: string;
 }

 export interface ContextFilter {
@@ -19,11 +19,6 @@ export function QueryOptionGroup({ title, children, collapsedInfo, queryStats }:
  const [isOpen, toggleOpen] = useToggle(false);
  const styles = useStyles2(getStyles);

-  const convertUnits = (): string => {
-    const { text, suffix } = getValueFormat('bytes')(queryStats!.bytes, 1);
-    return text + suffix;
-  };
-
  return (
    <div className={styles.wrapper}>
      <Collapse
@@ -46,12 +41,14 @@ export function QueryOptionGroup({ title, children, collapsedInfo, queryStats }:
      >
        <div className={styles.body}>{children}</div>
      </Collapse>
+
      {queryStats && config.featureToggles.lokiQuerySplitting && (
        <Tooltip content="Note: the query will be split into multiple parts and executed in sequence. Query limits will only apply each individual part.">
          <Icon tabIndex={0} name="info-circle" className={styles.tooltip} size="sm" />
        </Tooltip>
      )}
-      {queryStats && <p className={styles.stats}>This query will process approximately {convertUnits()}.</p>}
+
+      {queryStats && <p className={styles.stats}>{generateQueryStats(queryStats)}</p>}
    </div>
  );
 }
@@ -103,3 +100,16 @@ const getStyles = (theme: GrafanaTheme2) => {
    }),
  };
 };
+
+const generateQueryStats = (queryStats: QueryStats) => {
+  if (queryStats.message) {
+    return queryStats.message;
+  }
+
+  return `This query will process approximately ${convertUnits(queryStats)}.`;
+};
+
+const convertUnits = (queryStats: QueryStats): string => {
+  const { text, suffix } = getValueFormat('bytes')(queryStats.bytes, 1);
+  return text + suffix;
+};