Prometheus: Improved the function selector with search (#46084)

* Change to cascader with search

* Don't change to select on backspace

* Fix error on group select

* Simplify getOperationDef

* Set cascader wrapper to fixed width

* Add placeholder

* Fix tests and type errors

* Fix props for backward compatibility

* Add comments

packages/grafana-ui/src/components/Cascader/Cascader.tsx

@@ -23,7 +23,16 @@ export interface CascaderProps {
   allowCustomValue?: boolean;
   /** A function for formatting the message for custom value creation. Only applies when allowCustomValue is set to true*/
   formatCreateLabel?: (val: string) => string;
+  /** If true all levels are shown in the input by simple concatenating the labels */
   displayAllSelectedLevels?: boolean;
+  onBlur?: () => void;
+  /** When mounted focus automatically on the input */
+  autoFocus?: boolean;
+  /** Keep the dropdown open all the time, useful in case whole cascader visibility is controlled by the parent */
+  alwaysOpen?: boolean;
+  /** Don't show what is selected in the cascader input/search. Useful when input is used just as search and the
+      cascader is hidden after selection. */
+  hideActiveLevelLabel?: boolean;
 }
 
 interface CascaderState {
@@ -117,12 +126,15 @@ export class Cascader extends React.PureComponent<CascaderProps, CascaderState>
 
   //For rc-cascader
   onChange = (value: string[], selectedOptions: CascaderOption[]) => {
+    const activeLabel = this.props.hideActiveLevelLabel
+      ? ''
+      : this.props.displayAllSelectedLevels
+      ? selectedOptions.map((option) => option.label).join(this.props.separator || DEFAULT_SEPARATOR)
+      : selectedOptions[selectedOptions.length - 1].label;
     this.setState({
       rcValue: value,
       focusCascade: true,
-      activeLabel: this.props.displayAllSelectedLevels
+      activeLabel,
-        ? selectedOptions.map((option) => option.label).join(this.props.separator || DEFAULT_SEPARATOR)
-        : selectedOptions[selectedOptions.length - 1].label,
     });
 
     this.props.onSelect(selectedOptions[selectedOptions.length - 1].value);
@@ -159,22 +171,19 @@ export class Cascader extends React.PureComponent<CascaderProps, CascaderState>
         rcValue: [],
       });
     }
+    this.props.onBlur?.();
   };
 
   onBlurCascade = () => {
     this.setState({
       focusCascade: false,
     });
+
+    this.props.onBlur?.();
   };
 
   onInputKeyDown = (e: React.KeyboardEvent<HTMLInputElement>) => {
-    if (
+    if (['ArrowDown', 'ArrowUp', 'Enter', 'ArrowLeft', 'ArrowRight', 'Backspace'].includes(e.key)) {
-      e.key === 'ArrowDown' ||
-      e.key === 'ArrowUp' ||
-      e.key === 'Enter' ||
-      e.key === 'ArrowLeft' ||
-      e.key === 'ArrowRight'
-    ) {
       return;
     }
     this.setState({
@@ -183,6 +192,14 @@ export class Cascader extends React.PureComponent<CascaderProps, CascaderState>
     });
   };
 
+  onSelectInputChange = (value: string) => {
+    if (value === '') {
+      this.setState({
+        isSearching: false,
+      });
+    }
+  };
+
   render() {
     const { allowCustomValue, formatCreateLabel, placeholder, width, changeOnSelect, options } = this.props;
     const { focusCascade, isSearching, rcValue, activeLabel } = this.state;
@@ -203,6 +220,7 @@ export class Cascader extends React.PureComponent<CascaderProps, CascaderState>
             onCreateOption={this.onCreateOption}
             formatCreateLabel={formatCreateLabel}
             width={width}
+            onInputChange={this.onSelectInputChange}
           />
         ) : (
           <RCCascader
@@ -212,9 +230,11 @@ export class Cascader extends React.PureComponent<CascaderProps, CascaderState>
             value={rcValue.value}
             fieldNames={{ label: 'label', value: 'value', children: 'items' }}
             expandIcon={null}
+            open={this.props.alwaysOpen}
           >
             <div className={disableDivFocus}>
               <Input
+                autoFocus={this.props.autoFocus}
                 width={width}
                 placeholder={placeholder}
                 onBlur={this.onBlurCascade}

public/app/plugins/datasource/loki/querybuilder/LokiQueryModeller.test.ts

@@ -119,7 +119,7 @@ describe('LokiQueryModeller', () => {
         operations: [],
       };
 
-      const def = modeller.getOperationDef('sum');
+      const def = modeller.getOperationDef('sum')!;
       const result = def.addOperationHandler(def, query, modeller);
       expect(result.operations[0].id).toBe('rate');
       expect(result.operations[1].id).toBe('sum');
@@ -131,7 +131,7 @@ describe('LokiQueryModeller', () => {
         operations: [{ id: 'json', params: [] }],
       };
 
-      const def = modeller.getOperationDef('sum');
+      const def = modeller.getOperationDef('sum')!;
       const result = def.addOperationHandler(def, query, modeller);
       expect(result.operations[0].id).toBe('json');
       expect(result.operations[1].id).toBe('rate');
@@ -144,7 +144,7 @@ describe('LokiQueryModeller', () => {
         operations: [{ id: 'rate', params: [] }],
       };
 
-      const def = modeller.getOperationDef('json');
+      const def = modeller.getOperationDef('json')!;
       const result = def.addOperationHandler(def, query, modeller);
       expect(result.operations[0].id).toBe('json');
       expect(result.operations[1].id).toBe('rate');
@@ -156,7 +156,7 @@ describe('LokiQueryModeller', () => {
         operations: [{ id: '__line_contains', params: ['error'] }],
       };
 
-      const def = modeller.getOperationDef('json');
+      const def = modeller.getOperationDef('json')!;
       const result = def.addOperationHandler(def, query, modeller);
       expect(result.operations[0].id).toBe('__line_contains');
       expect(result.operations[1].id).toBe('json');
@@ -168,7 +168,7 @@ describe('LokiQueryModeller', () => {
         operations: [{ id: 'json', params: [] }],
       };
 
-      const def = modeller.getOperationDef('__line_contains');
+      const def = modeller.getOperationDef('__line_contains')!;
       const result = def.addOperationHandler(def, query, modeller);
       expect(result.operations[0].id).toBe('__line_contains');
       expect(result.operations[1].id).toBe('json');
@@ -180,7 +180,7 @@ describe('LokiQueryModeller', () => {
         operations: [],
       };
 
-      const def = modeller.getOperationDef('rate');
+      const def = modeller.getOperationDef('rate')!;
       const result = def.addOperationHandler(def, query, modeller);
       expect(result.operations.length).toBe(1);
     });

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

@@ -220,7 +220,11 @@ function getIndexOfOrLast(
   condition: (def: QueryBuilderOperationDef) => boolean
 ) {
   const index = operations.findIndex((x) => {
-    return condition(queryModeller.getOperationDef(x.id));
+    const opDef = queryModeller.getOperationDef(x.id);
+    if (!opDef) {
+      return false;
+    }
+    return condition(opDef);
   });
 
   return index === -1 ? operations.length : index;
@@ -242,7 +246,11 @@ export function addLokiOperation(
     case LokiVisualQueryOperationCategory.Aggregations:
     case LokiVisualQueryOperationCategory.Functions: {
       const rangeVectorFunction = operations.find((x) => {
-        return isRangeVectorFunction(modeller.getOperationDef(x.id));
+        const opDef = modeller.getOperationDef(x.id);
+        if (!opDef) {
+          return false;
+        }
+        return isRangeVectorFunction(opDef);
       });
 
       // If we are adding a function but we have not range vector function yet add one

public/app/plugins/datasource/prometheus/querybuilder/PromQueryModeller.ts

@@ -48,7 +48,7 @@ export class PromQueryModeller extends LokiAndPromQueryModellerBase<PromVisualQu
     return (
       query.operations.find((op) => {
         const def = this.getOperationDef(op.id);
-        return def.category === PromVisualQueryOperationCategory.BinaryOps;
+        return def?.category === PromVisualQueryOperationCategory.BinaryOps;
       }) !== undefined
     );
   }

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

@@ -50,7 +50,7 @@ describe('PromQueryBuilder', () => {
     setup();
     // Add label
     expect(screen.getByLabelText('Add')).toBeInTheDocument();
-    expect(screen.getByLabelText('Add operation')).toBeInTheDocument();
+    expect(screen.getByTitle('Add operation')).toBeInTheDocument();
   });
 
   it('renders all the query sections', async () => {

public/app/plugins/datasource/prometheus/querybuilder/operations.ts

@@ -350,7 +350,8 @@ export function addOperationWithRangeVector(
   };
 
   if (query.operations.length > 0) {
-    const firstOp = modeller.getOperationDef(query.operations[0].id);
+    // If operation exists it has to be in the registry so no point to check if it was found
+    const firstOp = modeller.getOperationDef(query.operations[0].id)!;
 
     if (firstOp.addOperationHandler === addOperationWithRangeVector) {
       return {

public/app/plugins/datasource/prometheus/querybuilder/shared/LokiAndPromQueryModellerBase.ts

@@ -37,8 +37,8 @@ export abstract class LokiAndPromQueryModellerBase<T extends QueryWithOperations
     return this.categories;
   }
 
-  getOperationDef(id: string) {
+  getOperationDef(id: string): QueryBuilderOperationDef | undefined {
-    return this.operationsRegisty.get(id);
+    return this.operationsRegisty.getIfExists(id);
   }
 
   renderOperations(queryString: string, operations: QueryBuilderOperation[]) {

public/app/plugins/datasource/prometheus/querybuilder/shared/OperationEditor.tsx

@@ -39,6 +39,9 @@ export function OperationEditor({
 }: Props) {
   const styles = useStyles2(getStyles);
   const def = queryModeller.getOperationDef(operation.id);
+  if (!def) {
+    return <span>Operation {operation.id} not found</span>;
+  }
 
   const onParamValueChanged = (paramIdx: number, value: QueryBuilderOperationParamValue) => {
     const update: QueryBuilderOperation = { ...operation, params: [...operation.params] };

public/app/plugins/datasource/prometheus/querybuilder/shared/OperationList.tsx

@@ -1,8 +1,8 @@
 import { css } from '@emotion/css';
 import { DataSourceApi, GrafanaTheme2 } from '@grafana/data';
 import { Stack } from '@grafana/experimental';
-import { ButtonCascader, CascaderOption, useStyles2 } from '@grafana/ui';
+import { Button, Cascader, CascaderOption, useStyles2 } from '@grafana/ui';
-import React from 'react';
+import React, { useState } from 'react';
 import { DragDropContext, Droppable, DropResult } from 'react-beautiful-dnd';
 import { QueryBuilderOperation, QueryWithOperations, VisualQueryModeller } from '../shared/types';
 import { OperationEditor } from './OperationEditor';
@@ -26,6 +26,8 @@ export function OperationList<T extends QueryWithOperations>({
   const styles = useStyles2(getStyles);
   const { operations } = query;
 
+  const [cascaderOpen, setCascaderOpen] = useState(false);
+
   const onOperationChange = (index: number, update: QueryBuilderOperation) => {
     const updatedList = [...operations];
     updatedList.splice(index, 1, update);
@@ -41,7 +43,7 @@ export function OperationList<T extends QueryWithOperations>({
     return {
       value: category,
       label: category,
-      children: queryModeller.getOperationsForCategory(category).map((operation) => ({
+      items: queryModeller.getOperationsForCategory(category).map((operation) => ({
         value: operation.id,
         label: operation.name,
         isLeaf: true,
@@ -49,9 +51,13 @@ export function OperationList<T extends QueryWithOperations>({
     };
   });
 
-  const onAddOperation = (value: string[]) => {
+  const onAddOperation = (value: string) => {
-    const operationDef = queryModeller.getOperationDef(value[1]);
+    const operationDef = queryModeller.getOperationDef(value);
+    if (!operationDef) {
+      return;
+    }
     onChange(operationDef.addOperationHandler(operationDef, query, queryModeller));
+    setCascaderOpen(false);
   };
 
   const onDragEnd = (result: DropResult) => {
@@ -66,6 +72,10 @@ export function OperationList<T extends QueryWithOperations>({
     onChange({ ...query, operations: updatedList });
   };
 
+  const onCascaderBlur = () => {
+    setCascaderOpen(false);
+  };
+
   return (
     <Stack gap={1} direction="column">
       <Stack gap={1}>
@@ -94,15 +104,19 @@ export function OperationList<T extends QueryWithOperations>({
           </DragDropContext>
         )}
         <div className={styles.addButton}>
-          <ButtonCascader
+          {cascaderOpen ? (
-            key="cascader"
+            <Cascader
-            icon="plus"
               options={addOptions}
-            onChange={onAddOperation}
+              onSelect={onAddOperation}
-            variant="secondary"
+              onBlur={onCascaderBlur}
-            hideDownIcon={true}
+              autoFocus={true}
-            buttonProps={{ 'aria-label': 'Add operation', title: 'Add operation' }}
+              alwaysOpen={true}
+              hideActiveLevelLabel={true}
+              placeholder={'Search'}
             />
+          ) : (
+            <Button icon={'plus'} variant={'secondary'} onClick={() => setCascaderOpen(true)} title={'Add operation'} />
+          )}
         </div>
       </Stack>
     </Stack>
@@ -122,6 +136,7 @@ const getStyles = (theme: GrafanaTheme2) => {
       gap: theme.spacing(2),
     }),
     addButton: css({
+      width: 150,
       paddingBottom: theme.spacing(1),
     }),
   };

public/app/plugins/datasource/prometheus/querybuilder/shared/OperationListExplained.tsx

@@ -14,6 +14,9 @@ export function OperationListExplained<T extends QueryWithOperations>({ query, q
     <>
       {query.operations.map((op, index) => {
         const def = queryModeller.getOperationDef(op.id);
+        if (!def) {
+          return `Operation ${op.id} not found`;
+        }
         const title = def.renderer(op, def, '<expr>');
         const body = def.explainHandler ? def.explainHandler(op, def) : def.documentation ?? 'no docs';
 

public/app/plugins/datasource/prometheus/querybuilder/shared/OperationName.tsx

@@ -60,7 +60,8 @@ export const OperationName = React.memo<Props>(({ operation, def, index, onChang
           onCloseMenu={onToggleSwitcher}
           onChange={(value) => {
             if (value.value) {
-              const newDef = queryModeller.getOperationDef(value.value.id);
+              // Operation should exist if it is selectable
+              const newDef = queryModeller.getOperationDef(value.value.id)!;
               let changedOp = { ...operation, id: value.value.id };
               onChange(index, def.changeTypeHandler ? def.changeTypeHandler(changedOp, newDef) : changedOp);
             }

public/app/plugins/datasource/prometheus/querybuilder/shared/types.ts

@@ -96,5 +96,5 @@ export interface VisualQueryModeller {
   getOperationsForCategory(category: string): QueryBuilderOperationDef[];
   getAlternativeOperations(key: string): QueryBuilderOperationDef[];
   getCategories(): string[];
-  getOperationDef(id: string): QueryBuilderOperationDef;
+  getOperationDef(id: string): QueryBuilderOperationDef | undefined;
 }