Loki: Show query size approximation (#62109)

* feat: make api request to /loki/api/v1/index/stats * fix: add /index/stats to callResource valid urls * feat: make call to getQueryStats when the query changes * feat: render user tooltip displaying the estimated value for processed data * fix: add new props to component tests * test: add tests for query size estimation * fix: disable error message on request failure * refactor: add suggestions from code review * refactor: only pass required query string
2024-11-26 02:40:26 -06:00 · 2023-01-27 17:32:53 +00:00 · 2023-01-27 17:32:53 +00:00 · b2c8126e6e
commit b2c8126e6e
parent 5ad7cca9d4
10 changed files with 243 additions and 91 deletions
--- a/pkg/tsdb/loki/loki.go
+++ b/pkg/tsdb/loki/loki.go
@ -112,7 +112,8 @@ func callResource(ctx context.Context, req *backend.CallResourceRequest, sender
 	}
 	if (!strings.HasPrefix(url, "labels?")) &&
 		(!strings.HasPrefix(url, "label/")) && // the `/label/$label_name/values` form
-		(!strings.HasPrefix(url, "series?")) {
+		(!strings.HasPrefix(url, "series?")) &&
+		(!strings.HasPrefix(url, "index/stats?")) {
 		return fmt.Errorf("invalid resource URL: %s", url)
 	}
 	lokiURL := fmt.Sprintf("/loki/api/v1/%s", url)
--- a/public/app/plugins/datasource/loki/components/LokiQueryEditor.tsx
+++ b/public/app/plugins/datasource/loki/components/LokiQueryEditor.tsx
@ -171,6 +171,7 @@ export const LokiQueryEditor = React.memo<LokiQueryEditorProps>((props) => {
          onRunQuery={onRunQuery}
          app={app}
          maxLines={datasource.maxLines}
+          datasource={datasource}
        />
      </EditorRows>
    </>
--- a/public/app/plugins/datasource/loki/datasource.ts
+++ b/public/app/plugins/datasource/loki/datasource.ts
@ -34,7 +34,7 @@ import {
  TimeRange,
  toUtc,
 } from '@grafana/data';
-import { config, DataSourceWithBackend, FetchError } from '@grafana/runtime';
+import { BackendSrvRequest, config, DataSourceWithBackend, FetchError } from '@grafana/runtime';
 import { DataQuery } from '@grafana/schema';
 import { queryLogsSample, queryLogsVolume } from 'app/core/logsModel';
 import { convertToWebSocketUrl } from 'app/core/utils/explore';
@ -71,6 +71,7 @@ import { getQueryHints } from './queryHints';
 import {
  getLogQueryFromMetricsQuery,
  getNormalizedLokiQuery,
+  getStreamSelectorsFromQuery,
  getParserFromQuery,
  isLogsQuery,
  isValidQuery,
@ -86,6 +87,7 @@ import {
  LokiQueryType,
  LokiVariableQuery,
  LokiVariableQueryType,
+  QueryStats,
  SupportingQueryType,
 } from './types';
 import { LokiVariableSupport } from './variables';
@ -401,15 +403,43 @@ export class LokiDatasource
    return queries.map((query) => this.languageProvider.exportToAbstractQuery(query));
  }

-  async metadataRequest(url: string, params?: Record<string, string | number>) {
+  async metadataRequest(url: string, params?: Record<string, string | number>, options?: Partial<BackendSrvRequest>) {
    // url must not start with a `/`, otherwise the AJAX-request
    // going from the browser will contain `//`, which can cause problems.
    if (url.startsWith('/')) {
      throw new Error(`invalid metadata request url: ${url}`);
    }

-    const res = await this.getResource(url, params);
-    return res.data || [];
+    const res = await this.getResource(url, params, options);
+    return res.data ?? (res || []);
+  }
+
+  async getQueryStats(query: LokiQuery): Promise<QueryStats> {
+    const { start, end } = this.getTimeRangeParams();
+    const labelMatchers = getStreamSelectorsFromQuery(query.expr);
+
+    let statsForAll: QueryStats = { streams: 0, chunks: 0, bytes: 0, entries: 0 };
+
+    for (const labelMatcher of labelMatchers) {
+      try {
+        const data = await this.metadataRequest(
+          'index/stats',
+          { query: labelMatcher, start, end },
+          { showErrorAlert: false }
+        );
+
+        statsForAll = {
+          streams: statsForAll.streams + data.streams,
+          chunks: statsForAll.chunks + data.chunks,
+          bytes: statsForAll.bytes + data.bytes,
+          entries: statsForAll.entries + data.entries,
+        };
+      } catch (e) {
+        break;
+      }
+    }
+
+    return statsForAll;
  }

  async metricFindQuery(query: LokiVariableQuery | string) {
--- a/public/app/plugins/datasource/loki/modifyQuery.ts
+++ b/public/app/plugins/datasource/loki/modifyQuery.ts
@ -139,7 +139,7 @@ export function removeCommentsFromQuery(query: string): string {
 * selector.
 * @param query
 */
-function getStreamSelectorPositions(query: string): Position[] {
+export function getStreamSelectorPositions(query: string): Position[] {
  const tree = parser.parse(query);
  const positions: Position[] = [];
  tree.iterate({
--- a/public/app/plugins/datasource/loki/queryUtils.ts
+++ b/public/app/plugins/datasource/loki/queryUtils.ts
@ -21,6 +21,7 @@ import {

 import { ErrorId } from '../prometheus/querybuilder/shared/parsingUtils';

+import { getStreamSelectorPositions } from './modifyQuery';
 import { LokiQuery, LokiQueryType } from './types';

 export function formatQuery(selector: string | undefined): string {
@ -284,3 +285,13 @@ export function isQueryWithLineFilter(query: string): boolean {

  return queryWithLineFilter;
 }
+
+export function getStreamSelectorsFromQuery(query: string): string[] {
+  const labelMatcherPositions = getStreamSelectorPositions(query);
+
+  const labelMatchers = labelMatcherPositions.map((labelMatcher) => {
+    return query.slice(labelMatcher.from, labelMatcher.to);
+  });
+
+  return labelMatchers;
+}
--- a/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilderOptions.test.tsx
+++ b/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilderOptions.test.tsx
@ -2,6 +2,7 @@ import { fireEvent, render, screen } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';

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

 import { LokiQueryBuilderOptions } from './LokiQueryBuilderOptions';
@ -47,6 +48,7 @@ function setup(queryOverrides: Partial<LokiQuery> = {}) {
    onRunQuery: jest.fn(),
    onChange: jest.fn(),
    maxLines: 20,
+    datasource: createLokiDatasource(),
  };

  const { container } = render(<LokiQueryBuilderOptions {...props} />);
--- a/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilderOptions.tsx
+++ b/public/app/plugins/datasource/loki/querybuilder/components/LokiQueryBuilderOptions.tsx
@ -1,4 +1,5 @@
-import React from 'react';
+import React, { useEffect, useState } from 'react';
+import { usePrevious } from 'react-use';

 import { CoreApp, SelectableValue } from '@grafana/data';
 import { EditorField, EditorRow } from '@grafana/experimental';
@ -7,8 +8,9 @@ import { RadioButtonGroup, Select, AutoSizeInput } from '@grafana/ui';
 import { QueryOptionGroup } from 'app/plugins/datasource/prometheus/querybuilder/shared/QueryOptionGroup';

 import { preprocessMaxLines, queryTypeOptions, RESOLUTION_OPTIONS } from '../../components/LokiOptionFields';
+import { LokiDatasource } from '../../datasource';
 import { isLogsQuery } from '../../queryUtils';
-import { LokiQuery, LokiQueryType } from '../../types';
+import { LokiQuery, LokiQueryType, QueryStats } from '../../types';

 export interface Props {
  query: LokiQuery;
@ -16,83 +18,112 @@ export interface Props {
  onRunQuery: () => void;
  maxLines: number;
  app?: CoreApp;
+  datasource: LokiDatasource;
 }

-export const LokiQueryBuilderOptions = React.memo<Props>(({ app, query, onChange, onRunQuery, maxLines }) => {
-  const onQueryTypeChange = (value: LokiQueryType) => {
-    onChange({ ...query, queryType: value });
-    onRunQuery();
-  };
+export const LokiQueryBuilderOptions = React.memo<Props>(
+  ({ app, query, onChange, onRunQuery, maxLines, datasource }) => {
+    const [queryStats, setQueryStats] = useState<QueryStats>();
+    const prevQuery = usePrevious(query);

-  const onResolutionChange = (option: SelectableValue<number>) => {
-    reportInteraction('grafana_loki_resolution_clicked', {
-      app,
-      resolution: option.value,
-    });
-    onChange({ ...query, resolution: option.value });
-    onRunQuery();
-  };
-
-  const onLegendFormatChanged = (evt: React.FormEvent<HTMLInputElement>) => {
-    onChange({ ...query, legendFormat: evt.currentTarget.value });
-    onRunQuery();
-  };
-
-  function onMaxLinesChange(e: React.SyntheticEvent<HTMLInputElement>) {
-    const newMaxLines = preprocessMaxLines(e.currentTarget.value);
-    if (query.maxLines !== newMaxLines) {
-      onChange({ ...query, maxLines: newMaxLines });
+    const onQueryTypeChange = (value: LokiQueryType) => {
+      onChange({ ...query, queryType: value });
      onRunQuery();
+    };
+
+    const onResolutionChange = (option: SelectableValue<number>) => {
+      reportInteraction('grafana_loki_resolution_clicked', {
+        app,
+        resolution: option.value,
+      });
+      onChange({ ...query, resolution: option.value });
+      onRunQuery();
+    };
+
+    const onLegendFormatChanged = (evt: React.FormEvent<HTMLInputElement>) => {
+      onChange({ ...query, legendFormat: evt.currentTarget.value });
+      onRunQuery();
+    };
+
+    function onMaxLinesChange(e: React.SyntheticEvent<HTMLInputElement>) {
+      const newMaxLines = preprocessMaxLines(e.currentTarget.value);
+      if (query.maxLines !== newMaxLines) {
+        onChange({ ...query, maxLines: newMaxLines });
+        onRunQuery();
+      }
    }
-  }

-  let queryType = query.queryType ?? (query.instant ? LokiQueryType.Instant : LokiQueryType.Range);
-  let showMaxLines = isLogsQuery(query.expr);
+    useEffect(() => {
+      if (query.expr === prevQuery?.expr) {
+        return;
+      }

-  return (
-    <EditorRow>
-      <QueryOptionGroup title="Options" collapsedInfo={getCollapsedInfo(query, queryType, showMaxLines, maxLines)}>
-        <EditorField
-          label="Legend"
-          tooltip="Series name override or template. Ex. {{hostname}} will be replaced with label value for hostname."
+      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);
+
+    return (
+      <EditorRow>
+        <QueryOptionGroup
+          title="Options"
+          collapsedInfo={getCollapsedInfo(query, queryType, showMaxLines, maxLines)}
+          queryStats={queryStats}
        >
-          <AutoSizeInput
-            placeholder="{{label}}"
-            id="loki-query-editor-legend-format"
-            type="string"
-            minWidth={14}
-            defaultValue={query.legendFormat}
-            onCommitChange={onLegendFormatChanged}
-          />
-        </EditorField>
-        <EditorField label="Type">
-          <RadioButtonGroup options={queryTypeOptions} value={queryType} onChange={onQueryTypeChange} />
-        </EditorField>
-        {showMaxLines && (
-          <EditorField label="Line limit" tooltip="Upper limit for number of log lines returned by query.">
+          <EditorField
+            label="Legend"
+            tooltip="Series name override or template. Ex. {{hostname}} will be replaced with label value for hostname."
+          >
            <AutoSizeInput
-              className="width-4"
-              placeholder={maxLines.toString()}
-              type="number"
-              min={0}
-              defaultValue={query.maxLines?.toString() ?? ''}
-              onCommitChange={onMaxLinesChange}
+              placeholder="{{label}}"
+              id="loki-query-editor-legend-format"
+              type="string"
+              minWidth={14}
+              defaultValue={query.legendFormat}
+              onCommitChange={onLegendFormatChanged}
            />
          </EditorField>
-        )}
-        <EditorField label="Resolution">
-          <Select
-            isSearchable={false}
-            onChange={onResolutionChange}
-            options={RESOLUTION_OPTIONS}
-            value={query.resolution || 1}
-            aria-label="Select resolution"
-          />
-        </EditorField>
-      </QueryOptionGroup>
-    </EditorRow>
-  );
-});
+          <EditorField label="Type">
+            <RadioButtonGroup options={queryTypeOptions} value={queryType} onChange={onQueryTypeChange} />
+          </EditorField>
+          {showMaxLines && (
+            <EditorField label="Line limit" tooltip="Upper limit for number of log lines returned by query.">
+              <AutoSizeInput
+                className="width-4"
+                placeholder={maxLines.toString()}
+                type="number"
+                min={0}
+                defaultValue={query.maxLines?.toString() ?? ''}
+                onCommitChange={onMaxLinesChange}
+              />
+            </EditorField>
+          )}
+          <EditorField label="Resolution">
+            <Select
+              isSearchable={false}
+              onChange={onResolutionChange}
+              options={RESOLUTION_OPTIONS}
+              value={query.resolution || 1}
+              aria-label="Select resolution"
+            />
+          </EditorField>
+        </QueryOptionGroup>
+      </EditorRow>
+    );
+  }
+);

 function getCollapsedInfo(
  query: LokiQuery,
--- a/public/app/plugins/datasource/loki/types.ts
+++ b/public/app/plugins/datasource/loki/types.ts
@ -154,6 +154,13 @@ export interface LokiVariableQuery extends DataQuery {
  stream?: string;
 }

+export interface QueryStats {
+  streams: number;
+  chunks: number;
+  bytes: number;
+  entries: number;
+}
+
 export enum SupportingQueryType {
  LogsVolume = 'logsVolume',
  LogsSample = 'logsSample',
--- a/public/app/plugins/datasource/prometheus/querybuilder/shared/QueryOptionGroup.test.tsx
+++ b/public/app/plugins/datasource/prometheus/querybuilder/shared/QueryOptionGroup.test.tsx
@ -0,0 +1,48 @@
+import { render, screen } from '@testing-library/react';
+import React from 'react';
+
+import { QueryOptionGroup, Props } from './QueryOptionGroup';
+
+describe('Query size approximation', () => {
+  const _1KiB = 1024; // size of 1 KiB in bytes
+  const _1GiB = 1073741824; // ...
+  const _1PiB = 1125899906842624;
+
+  it('renders the correct data value given 1 KiB', async () => {
+    const props = createProps(_1KiB);
+    render(<QueryOptionGroup {...props} />);
+    expect(screen.getByText(/This query will process approximately 1.0 KiB/)).toBeInTheDocument();
+  });
+
+  it('renders the correct data value given 1 GiB', async () => {
+    const props = createProps(_1GiB);
+    render(<QueryOptionGroup {...props} />);
+    expect(screen.getByText(/This query will process approximately 1.0 GiB/)).toBeInTheDocument();
+  });
+
+  it('renders the correct data value given 1 PiB', async () => {
+    const props = createProps(_1PiB);
+    render(<QueryOptionGroup {...props} />);
+    expect(screen.getByText(/This query will process approximately 1.0 PiB/)).toBeInTheDocument();
+  });
+
+  it('updates the data value on data change', async () => {
+    const props1 = createProps(_1KiB);
+    const props2 = createProps(_1PiB);
+
+    const { rerender } = render(<QueryOptionGroup {...props1} />);
+    expect(screen.getByText(/This query will process approximately 1.0 KiB/)).toBeInTheDocument();
+
+    rerender(<QueryOptionGroup {...props2} />);
+    expect(screen.getByText(/This query will process approximately 1.0 PiB/)).toBeInTheDocument();
+  });
+});
+
+function createProps(bytes?: number): Props {
+  return {
+    title: 'Options',
+    collapsedInfo: ['Type: Range', 'Line limit: 1000'],
+    children: <div></div>,
+    queryStats: { streams: 0, chunks: 0, bytes: bytes ?? 0, entries: 0 },
+  };
+}
--- a/public/app/plugins/datasource/prometheus/querybuilder/shared/QueryOptionGroup.tsx
+++ b/public/app/plugins/datasource/prometheus/querybuilder/shared/QueryOptionGroup.tsx
@ -2,42 +2,58 @@ import { css } from '@emotion/css';
 import React from 'react';
 import { useToggle } from 'react-use';

-import { GrafanaTheme2 } from '@grafana/data';
+import { getValueFormat, GrafanaTheme2 } from '@grafana/data';
 import { Stack } from '@grafana/experimental';
 import { Icon, useStyles2 } from '@grafana/ui';
+import { QueryStats } from 'app/plugins/datasource/loki/types';

 export interface Props {
  title: string;
  collapsedInfo: string[];
+  queryStats?: QueryStats;
  children: React.ReactNode;
 }

-export function QueryOptionGroup({ title, children, collapsedInfo }: Props) {
+export function QueryOptionGroup({ title, children, collapsedInfo, queryStats }: Props) {
  const [isOpen, toggleOpen] = useToggle(false);
  const styles = useStyles2(getStyles);

+  const convertUnits = (): string => {
+    const { text, suffix } = getValueFormat('bytes')(queryStats!.bytes, 1);
+    return text + suffix;
+  };
+
  return (
-    <Stack gap={0} direction="column">
-      <div className={styles.header} onClick={toggleOpen} title="Click to edit options">
-        <div className={styles.toggle}>
-          <Icon name={isOpen ? 'angle-down' : 'angle-right'} />
-        </div>
-        <h6 className={styles.title}>{title}</h6>
-        {!isOpen && (
-          <div className={styles.description}>
-            {collapsedInfo.map((x, i) => (
-              <span key={i}>{x}</span>
-            ))}
+    <div className={styles.wrapper}>
+      <Stack gap={0} direction="column">
+        <div className={styles.header} onClick={toggleOpen} title="Click to edit options">
+          <div className={styles.toggle}>
+            <Icon name={isOpen ? 'angle-down' : 'angle-right'} />
          </div>
-        )}
-      </div>
-      {isOpen && <div className={styles.body}>{children}</div>}
-    </Stack>
+          <h6 className={styles.title}>{title}</h6>
+          {!isOpen && (
+            <div className={styles.description}>
+              {collapsedInfo.map((x, i) => (
+                <span key={i}>{x}</span>
+              ))}
+            </div>
+          )}
+        </div>
+        {isOpen && <div className={styles.body}>{children}</div>}
+      </Stack>
+      {queryStats && <p className={styles.stats}>This query will process approximately {convertUnits()}.</p>}
+    </div>
  );
 }

 const getStyles = (theme: GrafanaTheme2) => {
  return {
+    wrapper: css({
+      width: '100%',
+      display: 'flex',
+      justifyContent: 'space-between',
+      alignItems: 'baseline',
+    }),
    switchLabel: css({
      color: theme.colors.text.secondary,
      cursor: 'pointer',
@ -79,5 +95,10 @@ const getStyles = (theme: GrafanaTheme2) => {
      color: theme.colors.text.secondary,
      marginRight: `${theme.spacing(1)}`,
    }),
+    stats: css({
+      margin: '0px',
+      color: theme.colors.text.secondary,
+      fontSize: theme.typography.bodySmall.fontSize,
+    }),
  };
 };