Prometheus query builder: Replace select with AsyncSelect to support high cardinality prometheus instances (#57300)

* Replace current select with AsyncSelect component to facilitate autocomplete via prometheus server instead of client application
Co-authored-by: Kyle Brandt <kyle@grafana.com>

public/app/plugins/datasource/prometheus/querybuilder/components/MetricSelect.test.tsx

@@ -2,18 +2,58 @@ import { render, screen, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
 
+import { DataSourceInstanceSettings, MetricFindValue } from '@grafana/data/src';
+
+import { PrometheusDatasource } from '../../datasource';
+import { PromOptions } from '../../types';
+
 import { MetricSelect } from './MetricSelect';
 
+const instanceSettings = {
+  url: 'proxied',
+  id: 1,
+  directUrl: 'direct',
+  user: 'test',
+  password: 'mupp',
+  jsonData: { httpMethod: 'GET' },
+} as unknown as DataSourceInstanceSettings<PromOptions>;
+
+const dataSourceMock = new PrometheusDatasource(instanceSettings);
+const mockValues = [{ label: 'random_metric' }, { label: 'unique_metric' }, { label: 'more_unique_metric' }];
+
+// Mock metricFindQuery which will call backend API
+//@ts-ignore
+dataSourceMock.metricFindQuery = jest.fn((query: string) => {
+  // Use the label values regex to get the values inside the label_values function call
+  const labelValuesRegex = /^label_values\((?:(.+),\s*)?([a-zA-Z_][a-zA-Z0-9_]*)\)\s*$/;
+  const queryValueArray = query.match(labelValuesRegex) as RegExpMatchArray;
+  const queryValueRaw = queryValueArray[1] as string;
+
+  // Remove the wrapping regex
+  const queryValue = queryValueRaw.substring(queryValueRaw.indexOf('".*') + 3, queryValueRaw.indexOf('.*"'));
+
+  // Run the regex that we'd pass into prometheus API against the strings in the test
+  return Promise.resolve(
+    mockValues
+      .filter((value) => value.label.match(queryValue))
+      .map((result) => {
+        return {
+          text: result.label,
+        };
+      }) as MetricFindValue[]
+  );
+});
+
 const props = {
+  labelsFilters: [],
+  datasource: dataSourceMock,
   query: {
     metric: '',
     labels: [],
     operations: [],
   },
   onChange: jest.fn(),
-  onGetMetrics: jest
+  onGetMetrics: jest.fn().mockResolvedValue(mockValues),
-    .fn()
-    .mockResolvedValue([{ label: 'random_metric' }, { label: 'unique_metric' }, { label: 'more_unique_metric' }]),
 };
 
 describe('MetricSelect', () => {

public/app/plugins/datasource/prometheus/querybuilder/components/MetricSelect.tsx

@@ -1,11 +1,14 @@
 import { css } from '@emotion/css';
+import debounce from 'debounce-promise';
 import React, { useCallback, useState } from 'react';
 import Highlighter from 'react-highlight-words';
 
 import { SelectableValue, toOption, GrafanaTheme2 } from '@grafana/data';
 import { EditorField, EditorFieldGroup } from '@grafana/experimental';
-import { Select, FormatOptionLabelMeta, useStyles2 } from '@grafana/ui';
+import { AsyncSelect, FormatOptionLabelMeta, useStyles2 } from '@grafana/ui';
 
+import { PrometheusDatasource } from '../../datasource';
+import { QueryBuilderLabelFilter } from '../shared/types';
 import { PromVisualQuery } from '../types';
 
 // We are matching words split with space
@@ -15,9 +18,13 @@ export interface Props {
   query: PromVisualQuery;
   onChange: (query: PromVisualQuery) => void;
   onGetMetrics: () => Promise<SelectableValue[]>;
+  datasource: PrometheusDatasource;
+  labelsFilters: QueryBuilderLabelFilter[];
 }
 
-export function MetricSelect({ query, onChange, onGetMetrics }: Props) {
+const MAX_NUMBER_OF_RESULTS = 1000;
+
+export function MetricSelect({ datasource, query, onChange, onGetMetrics, labelsFilters }: Props) {
   const styles = useStyles2(getStyles);
   const [state, setState] = useState<{
     metrics?: Array<SelectableValue<any>>;
@@ -57,25 +64,87 @@ export function MetricSelect({ query, onChange, onGetMetrics }: Props) {
     [styles.highlight]
   );
 
+  const formatLabelFilters = (labelsFilters: QueryBuilderLabelFilter[]): string[] => {
+    return labelsFilters.map((label) => {
+      return `,${label.label}="${label.value}"`;
+    });
+  };
+
+  /**
+   * Transform queryString and any currently set label filters into label_values() string
+   */
+  const queryAndFilterToLabelValuesString = (
+    queryString: string,
+    labelsFilters: QueryBuilderLabelFilter[] | undefined
+  ): string => {
+    return `label_values({__name__=~".*${queryString}"${
+      labelsFilters ? formatLabelFilters(labelsFilters).join() : ''
+    }},__name__)`;
+  };
+
+  /**
+   * There aren't any spaces in the metric names, so let's introduce a wildcard into the regex for each space to better facilitate a fuzzy search
+   */
+  const regexifyLabelValuesQueryString = (query: string) => {
+    const queryArray = query.split(' ');
+    return queryArray.map((query) => `${query}.*`).join('');
+  };
+
+  /**
+   * Reformat the query string and label filters to return all valid results for current query editor state
+   */
+  const formatKeyValueStringsForLabelValuesQuery = (
+    query: string,
+    labelsFilters?: QueryBuilderLabelFilter[]
+  ): string => {
+    const queryString = regexifyLabelValuesQueryString(query);
+
+    return queryAndFilterToLabelValuesString(queryString, labelsFilters);
+  };
+
+  /**
+   * Gets label_values response from prometheus API for current autocomplete query string and any existing labels filters
+   */
+  const getMetricLabels = (query: string) => {
+    // Since some customers can have millions of metrics, whenever the user changes the autocomplete text we want to call the backend and request all metrics that match the current query string
+    const results = datasource.metricFindQuery(formatKeyValueStringsForLabelValuesQuery(query, labelsFilters));
+    return results.then((results) => {
+      if (results.length > MAX_NUMBER_OF_RESULTS) {
+        results.splice(0, results.length - MAX_NUMBER_OF_RESULTS);
+      }
+      return results.map((result) => {
+        return {
+          label: result.text,
+          value: result.text,
+        };
+      });
+    });
+  };
+
+  const debouncedSearch = debounce((query: string) => getMetricLabels(query), 300);
+
   return (
     <EditorFieldGroup>
       <EditorField label="Metric">
-        <Select
+        <AsyncSelect
           inputId="prometheus-metric-select"
           className={styles.select}
           value={query.metric ? toOption(query.metric) : undefined}
           placeholder="Select metric"
-          virtualized
           allowCustomValue
           formatOptionLabel={formatOptionLabel}
           filterOption={customFilterOption}
           onOpenMenu={async () => {
             setState({ isLoading: true });
             const metrics = await onGetMetrics();
+            if (metrics.length > MAX_NUMBER_OF_RESULTS) {
+              metrics.splice(0, metrics.length - MAX_NUMBER_OF_RESULTS);
+            }
             setState({ metrics, isLoading: undefined });
           }}
+          loadOptions={debouncedSearch}
           isLoading={state.isLoading}
-          options={state.metrics}
+          defaultOptions={state.metrics}
           onChange={({ value }) => {
             if (value) {
               onChange({ ...query, metric: value });

public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilder.tsx

@@ -99,7 +99,13 @@ export const PromQueryBuilder = React.memo<Props>((props) => {
   return (
     <>
       <EditorRow>
-        <MetricSelect query={query} onChange={onChange} onGetMetrics={onGetMetrics} />
+        <MetricSelect
+          query={query}
+          onChange={onChange}
+          onGetMetrics={onGetMetrics}
+          datasource={datasource}
+          labelsFilters={query.labels}
+        />
         <LabelFilters
           labelsFilters={query.labels}
           onChange={onChangeLabels}