Loki: Improve the display of loki query stats (#63623)

* fix: refresh query stats on timerange change

* partial fix: request stats update on type in code mode

* complete fix: request stats update on type in code mode

* fix: update failing tests and props

* refactor: pass only essential query string to getQueryStats

* refactor: remove unused variables

* test: fix datasource.getTimeRange is not a function

* refactor: use lodash debounce instead of setTimeout

* refactor: add suggestions from code review

* test: shouldUpdateStats and makeStatsRequest

* refactor: use more appropriate variable names

* refactor: make setQueryStats required instead of optional

* refactor: move setQueryStats into LokiQueryEditor

* fix: add missing props to LokiQueryField

* revert changes

* refactor: use inversion of control to request stats

* refactor: remove unnecessary code

public/app/plugins/datasource/loki/components/LokiQueryEditor.tsx

@@ -1,4 +1,5 @@
 import React, { SyntheticEvent, useCallback, useEffect, useState } from 'react';
+import { usePrevious } from 'react-use';
 
 import { CoreApp, LoadingState } from '@grafana/data';
 import { selectors } from '@grafana/e2e-selectors';
@@ -17,8 +18,9 @@ import { LokiQueryCodeEditor } from '../querybuilder/components/LokiQueryCodeEdi
 import { QueryPatternsModal } from '../querybuilder/components/QueryPatternsModal';
 import { buildVisualQueryFromString } from '../querybuilder/parsing';
 import { changeEditorMode, getQueryWithDefaults } from '../querybuilder/state';
-import { LokiQuery } from '../types';
+import { LokiQuery, QueryStats } from '../types';
 
+import { getStats, shouldUpdateStats } from './stats';
 import { LokiQueryEditorProps } from './types';
 
 export const testIds = {
@@ -31,9 +33,15 @@ export const LokiQueryEditor = React.memo<LokiQueryEditorProps>((props) => {
   const [queryPatternsModalOpen, setQueryPatternsModalOpen] = useState(false);
   const [dataIsStale, setDataIsStale] = useState(false);
   const [labelBrowserVisible, setLabelBrowserVisible] = useState(false);
+  const [queryStats, setQueryStats] = useState<QueryStats>();
   const { flag: explain, setFlag: setExplain } = useFlag(lokiQueryEditorExplainKey);
 
+  const timerange = datasource.getTimeRange();
+  const previousTimerange = usePrevious(timerange);
+
   const query = getQueryWithDefaults(props.query);
+  const previousQuery = usePrevious(query.expr);
+
   // This should be filled in from the defaults by now.
   const editorMode = query.editorMode!;
 
@@ -80,6 +88,17 @@ export const LokiQueryEditor = React.memo<LokiQueryEditorProps>((props) => {
     setLabelBrowserVisible((visible) => !visible);
   };
 
+  useEffect(() => {
+    const update = shouldUpdateStats(query.expr, previousQuery, timerange, previousTimerange);
+    if (update) {
+      const makeAsyncRequest = async () => {
+        const stats = await getStats(datasource, query.expr);
+        setQueryStats(stats);
+      };
+      makeAsyncRequest();
+    }
+  }, [datasource, timerange, previousTimerange, query, previousQuery, setQueryStats]);
+
   return (
     <>
       <ConfirmModal
@@ -154,7 +173,13 @@ export const LokiQueryEditor = React.memo<LokiQueryEditorProps>((props) => {
       <Space v={0.5} />
       <EditorRows>
         {editorMode === QueryEditorMode.Code && (
-          <LokiQueryCodeEditor {...props} query={query} onChange={onChangeInternal} showExplain={explain} />
+          <LokiQueryCodeEditor
+            {...props}
+            query={query}
+            onChange={onChangeInternal}
+            showExplain={explain}
+            setQueryStats={setQueryStats}
+          />
         )}
         {editorMode === QueryEditorMode.Builder && (
           <LokiQueryBuilderContainer
@@ -172,6 +197,7 @@ export const LokiQueryEditor = React.memo<LokiQueryEditorProps>((props) => {
           app={app}
           maxLines={datasource.maxLines}
           datasource={datasource}
+          queryStats={queryStats}
         />
       </EditorRows>
     </>

public/app/plugins/datasource/loki/components/LokiQueryEditorByApp.test.tsx

@@ -21,6 +21,7 @@ function setup(app: CoreApp): RenderResult {
     getQueryHints: () => [],
     getDataSamples: () => [],
     maxLines: 20,
+    getTimeRange: jest.fn(),
   } as unknown as LokiDatasource;
 
   return render(

public/app/plugins/datasource/loki/components/LokiQueryField.tsx

@@ -12,6 +12,7 @@ export interface LokiQueryFieldProps extends QueryEditorProps<LokiDatasource, Lo
   ExtraFieldElement?: ReactNode;
   placeholder?: string;
   'data-testid'?: string;
+  onQueryType?: (query: string) => void;
 }
 
 interface LokiQueryFieldState {
@@ -65,7 +66,7 @@ export class LokiQueryField extends React.PureComponent<LokiQueryFieldProps, Lok
   };
 
   render() {
-    const { ExtraFieldElement, query, app, datasource, history, onRunQuery } = this.props;
+    const { ExtraFieldElement, query, app, datasource, history, onRunQuery, onQueryType } = this.props;
     const placeholder = this.props.placeholder ?? 'Enter a Loki query (run with Shift+Enter)';
 
     return (
@@ -83,6 +84,7 @@ export class LokiQueryField extends React.PureComponent<LokiQueryFieldProps, Lok
               onRunQuery={onRunQuery}
               initialValue={query.expr ?? ''}
               placeholder={placeholder}
+              onQueryType={onQueryType}
             />
           </div>
         </div>

public/app/plugins/datasource/loki/components/monaco-query-field/MonacoQueryField.tsx

@@ -1,4 +1,5 @@
 import { css } from '@emotion/css';
+import { debounce } from 'lodash';
 import React, { useRef, useEffect } from 'react';
 import { useLatest } from 'react-use';
 import { v4 as uuidv4 } from 'uuid';
@@ -8,6 +9,8 @@ import { selectors } from '@grafana/e2e-selectors';
 import { languageConfiguration, monarchlanguage } from '@grafana/monaco-logql';
 import { useTheme2, ReactMonacoEditor, Monaco, monacoTypes, MonacoEditor } from '@grafana/ui';
 
+import { isValidQuery } from '../../queryUtils';
+
 import { Props } from './MonacoQueryFieldProps';
 import { getOverrideServices } from './getOverrideServices';
 import { getCompletionProvider, getSuggestOptions } from './monaco-completion-provider';
@@ -87,7 +90,15 @@ const getStyles = (theme: GrafanaTheme2, placeholder: string) => {
   };
 };
 
-const MonacoQueryField = ({ history, onBlur, onRunQuery, initialValue, datasource, placeholder }: Props) => {
+const MonacoQueryField = ({
+  history,
+  onBlur,
+  onRunQuery,
+  initialValue,
+  datasource,
+  placeholder,
+  onQueryType,
+}: Props) => {
   const id = uuidv4();
   // we need only one instance of `overrideServices` during the lifetime of the react component
   const overrideServicesRef = useRef(getOverrideServices());
@@ -138,6 +149,14 @@ const MonacoQueryField = ({ history, onBlur, onRunQuery, initialValue, datasourc
     editor.onDidChangeModelContent(checkDecorators);
   };
 
+  const onTypeDebounced = debounce(async (query: string) => {
+    if (!onQueryType || (isValidQuery(query) === false && query !== '')) {
+      return;
+    }
+
+    onQueryType(query);
+  }, 1000);
+
   return (
     <div
       aria-label={selectors.components.QueryField.container}
@@ -181,6 +200,8 @@ const MonacoQueryField = ({ history, onBlur, onRunQuery, initialValue, datasourc
               severity: monaco.MarkerSeverity.Error,
               ...boundary,
             }));
+
+            onTypeDebounced(query);
             monaco.editor.setModelMarkers(model, 'owner', markers);
           });
           const dataProvider = new CompletionDataProvider(langProviderRef.current, historyRef);

public/app/plugins/datasource/loki/components/monaco-query-field/MonacoQueryFieldProps.ts

@@ -14,4 +14,5 @@ export type Props = {
   onBlur: (value: string) => void;
   placeholder: string;
   datasource: LokiDatasource;
+  onQueryType?: (query: string) => void;
 };

public/app/plugins/datasource/loki/components/monaco-query-field/MonacoQueryFieldWrapper.tsx

@@ -7,6 +7,7 @@ export type Props = Omit<MonacoProps, 'onRunQuery' | 'onBlur'> & {
   onChange: (query: string) => void;
   onRunQuery: () => void;
   runQueryOnBlur: boolean;
+  onQueryType?: (query: string) => void;
 };
 
 export const MonacoQueryFieldWrapper = (props: Props) => {

public/app/plugins/datasource/loki/components/stats.test.ts

@@ -0,0 +1,59 @@
+import { TimeRange } from '@grafana/data';
+
+import { createLokiDatasource } from '../mocks';
+
+import { getStats, shouldUpdateStats } from './stats';
+
+describe('shouldUpdateStats', () => {
+  it('should return true if the query has changed', () => {
+    const query = '{job="grafana"}';
+    const prevQuery = '{job="not-grafana"}';
+    const timerange = { raw: { from: 'now-1h', to: 'now' } } as TimeRange;
+    const prevTimerange = { raw: { from: 'now-1h', to: 'now' } } as TimeRange;
+    expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange)).toBe(true);
+  });
+  it('should return true if the timerange has changed', () => {
+    const query = '{job="grafana"}';
+    const prevQuery = '{job="grafana"}';
+    const timerange = { raw: { from: 'now-1h', to: 'now' } } as TimeRange;
+    const prevTimerange = { raw: { from: 'now-2h', to: 'now' } } as TimeRange;
+    expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange)).toBe(true);
+  });
+
+  it('should return false if the query and timerange have not changed', () => {
+    const query = '{job="grafana"}';
+    const prevQuery = '{job="grafana"}';
+    const timerange = { raw: { from: 'now-1h', to: 'now' } } as TimeRange;
+    const prevTimerange = { raw: { from: 'now-1h', to: 'now' } } as TimeRange;
+    expect(shouldUpdateStats(query, prevQuery, timerange, prevTimerange)).toBe(false);
+  });
+});
+
+describe('makeStatsRequest', () => {
+  const datasource = createLokiDatasource();
+
+  it('should return undefined if there is no query', () => {
+    const query = '';
+    expect(getStats(datasource, query)).resolves.toBe(undefined); // change
+  });
+
+  it('should return undefined 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(undefined);
+  });
+
+  it('should return the stats if the response has data', () => {
+    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,
+    });
+  });
+});

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

@@ -0,0 +1,30 @@
+import { TimeRange } from '@grafana/data';
+
+import { LokiDatasource } from '../datasource';
+import { QueryStats } from '../types';
+
+export async function getStats(datasource: LokiDatasource, query: string): Promise<QueryStats | undefined> {
+  if (!query) {
+    return undefined;
+  }
+
+  const response = await datasource.getQueryStats(query);
+  return Object.values(response).every((v) => v === 0) ? undefined : response;
+}
+
+export function shouldUpdateStats(
+  query: string,
+  prevQuery: string | undefined,
+  timerange: TimeRange,
+  prevTimerange: TimeRange | undefined
+): boolean {
+  if (
+    query === prevQuery &&
+    timerange.raw.from === prevTimerange?.raw.from &&
+    timerange.raw.to === prevTimerange?.raw.to
+  ) {
+    return false;
+  }
+
+  return true;
+}

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

@@ -424,9 +424,9 @@ export class LokiDatasource
     return res.data ?? (res || []);
   }
 
-  async getQueryStats(query: LokiQuery): Promise<QueryStats> {
+  async getQueryStats(query: string): Promise<QueryStats> {
     const { start, end } = this.getTimeRangeParams();
-    const labelMatchers = getStreamSelectorsFromQuery(query.expr);
+    const labelMatchers = getStreamSelectorsFromQuery(query);
 
     let statsForAll: QueryStats = { streams: 0, chunks: 0, bytes: 0, entries: 0 };
 

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

@@ -49,6 +49,7 @@ function setup(queryOverrides: Partial<LokiQuery> = {}) {
     onChange: jest.fn(),
     maxLines: 20,
     datasource: createLokiDatasource(),
+    queryStats: { streams: 0, chunks: 0, bytes: 0, entries: 0 },
   };
 
   const { container } = render(<LokiQueryBuilderOptions {...props} />);

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

@@ -1,5 +1,4 @@
-import React, { useEffect, useState } from 'react';
-import { usePrevious } from 'react-use';
+import React, { useState } from 'react';
 
 import { CoreApp, isValidDuration, SelectableValue } from '@grafana/data';
 import { EditorField, EditorRow } from '@grafana/experimental';
@@ -19,13 +18,12 @@ export interface Props {
   maxLines: number;
   app?: CoreApp;
   datasource: LokiDatasource;
+  queryStats: QueryStats | undefined;
 }
 
 export const LokiQueryBuilderOptions = React.memo<Props>(
-  ({ app, query, onChange, onRunQuery, maxLines, datasource }) => {
-    const [queryStats, setQueryStats] = useState<QueryStats>();
+  ({ app, query, onChange, onRunQuery, maxLines, datasource, queryStats }) => {
     const [chunkRangeValid, setChunkRangeValid] = useState(true);
-    const prevQuery = usePrevious(query);
 
     const onQueryTypeChange = (value: LokiQueryType) => {
       onChange({ ...query, queryType: value });
@@ -65,25 +63,6 @@ export const LokiQueryBuilderOptions = React.memo<Props>(
       }
     }
 
-    useEffect(() => {
-      if (query.expr === prevQuery?.expr) {
-        return;
-      }
-
-      if (!query.expr) {
-        setQueryStats(undefined);
-        return;
-      }
-
-      const makeAsyncRequest = async () => {
-        const res = await datasource.getQueryStats(query);
-
-        // this filters out the case where the user has not configured loki to use tsdb, in that case all keys in the query stats will be 0
-        Object.values(res).every((v) => v === 0) ? setQueryStats(undefined) : setQueryStats(res);
-      };
-      makeAsyncRequest();
-    }, [query, prevQuery, datasource]);
-
     let queryType = query.queryType ?? (query.instant ? LokiQueryType.Instant : LokiQueryType.Range);
     let showMaxLines = isLogsQuery(query.expr);
 

public/app/plugins/datasource/loki/querybuilder/components/LokiQueryCodeEditor.test.tsx

@@ -20,6 +20,7 @@ const createDefaultProps = () => {
     onRunQuery: () => {},
     onChange: () => {},
     showExplain: false,
+    setQueryStats: () => {},
   };
 
   return props;

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

@@ -6,12 +6,15 @@ import { useStyles2 } from '@grafana/ui';
 
 import { testIds } from '../../components/LokiQueryEditor';
 import { LokiQueryField } from '../../components/LokiQueryField';
+import { getStats } from '../../components/stats';
 import { LokiQueryEditorProps } from '../../components/types';
+import { QueryStats } from '../../types';
 
 import { LokiQueryBuilderExplained } from './LokiQueryBuilderExplained';
 
 type Props = LokiQueryEditorProps & {
   showExplain: boolean;
+  setQueryStats: React.Dispatch<React.SetStateAction<QueryStats | undefined>>;
 };
 
 export function LokiQueryCodeEditor({
@@ -24,6 +27,7 @@ export function LokiQueryCodeEditor({
   app,
   showExplain,
   history,
+  setQueryStats,
 }: Props) {
   const styles = useStyles2(getStyles);
 
@@ -39,6 +43,10 @@ export function LokiQueryCodeEditor({
         data={data}
         app={app}
         data-testid={testIds.editor}
+        onQueryType={async (query: string) => {
+          const stats = await getStats(datasource, query);
+          setQueryStats(stats);
+        }}
       />
       {showExplain && <LokiQueryBuilderExplained query={query.expr} />}
     </div>