From 916f152a6f10ab09fb717ec14718b9c2eef5551d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torkel=20=C3=96degaard?= <torkel@grafana.com>
Date: Wed, 16 Feb 2022 08:05:16 +0100
Subject: [PATCH] Prometheus: Fixes issue adding rest params  (#45364)

* Prometheus: Fixing issue adding rest param and adding test

* Rebuild when expr changes from outside
---
 .../components/LabelParamEditor.tsx           | 11 ++-
 .../components/PromQueryBuilder.test.tsx      |  7 +-
 .../PromQueryBuilderContainer.test.tsx        | 68 ++++++++++++-------
 .../components/PromQueryBuilderContainer.tsx  | 25 +++++--
 .../components/PromQueryBuilderContext.tsx    | 12 ----
 .../querybuilder/shared/OperationEditor.tsx   | 21 ++++--
 .../shared/OperationParamEditor.tsx           | 11 ++-
 .../querybuilder/shared/operationUtils.ts     |  4 ++
 .../prometheus/querybuilder/shared/types.ts   |  2 +
 9 files changed, 110 insertions(+), 51 deletions(-)
 delete mode 100644 public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContext.tsx

diff --git a/public/app/plugins/datasource/prometheus/querybuilder/components/LabelParamEditor.tsx b/public/app/plugins/datasource/prometheus/querybuilder/components/LabelParamEditor.tsx
index 51a3bde9d59..e2be9939a7e 100644
--- a/public/app/plugins/datasource/prometheus/querybuilder/components/LabelParamEditor.tsx
+++ b/public/app/plugins/datasource/prometheus/querybuilder/components/LabelParamEditor.tsx
@@ -3,10 +3,18 @@ import { Select } from '@grafana/ui';
 import React, { useState } from 'react';
 import { PrometheusDatasource } from '../../datasource';
 import { promQueryModeller } from '../PromQueryModeller';
+import { getOperationParamId } from '../shared/operationUtils';
 import { QueryBuilderOperationParamEditorProps } from '../shared/types';
 import { PromVisualQuery } from '../types';
 
-export function LabelParamEditor({ onChange, index, value, query, datasource }: QueryBuilderOperationParamEditorProps) {
+export function LabelParamEditor({
+  onChange,
+  index,
+  operationIndex,
+  value,
+  query,
+  datasource,
+}: QueryBuilderOperationParamEditorProps) {
   const [state, setState] = useState<{
     options?: Array<SelectableValue<any>>;
     isLoading?: boolean;
@@ -14,6 +22,7 @@ export function LabelParamEditor({ onChange, index, value, query, datasource }:
 
   return (
     <Select
+      inputId={getOperationParamId(operationIndex, index)}
       menuShouldPortal
       autoFocus={value === '' ? true : undefined}
       openMenuOnFocus
diff --git a/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilder.test.tsx b/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilder.test.tsx
index 83745261467..ed69157ce1f 100644
--- a/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilder.test.tsx
+++ b/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilder.test.tsx
@@ -36,7 +36,7 @@ const bugQuery: PromVisualQuery = {
         labels: [{ label: 'foo', op: '=', value: 'bar' }],
         operations: [
           {
-            id: '__sum_by',
+            id: '__avg_by',
             params: ['app'],
           },
         ],
@@ -58,11 +58,12 @@ describe('PromQueryBuilder', () => {
     expect(screen.getByText('random_metric')).toBeInTheDocument();
     expect(screen.getByText('localhost:9090')).toBeInTheDocument();
     expect(screen.getByText('Rate')).toBeInTheDocument();
-    const sumBys = screen.getAllByTestId('operation-wrapper-for-__sum_by');
+    const sumBys = screen.getAllByTestId('operations.1.wrapper');
     expect(getByText(sumBys[0], 'instance')).toBeInTheDocument();
     expect(getByText(sumBys[0], 'job')).toBeInTheDocument();
 
-    expect(getByText(sumBys[1], 'app')).toBeInTheDocument();
+    const avgBys = screen.getAllByTestId('operations.0.wrapper');
+    expect(getByText(avgBys[1], 'app')).toBeInTheDocument();
     expect(screen.getByText('Operator')).toBeInTheDocument();
     expect(screen.getByText('Vector matches')).toBeInTheDocument();
   });
diff --git a/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContainer.test.tsx b/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContainer.test.tsx
index e8b7a7a7b4d..164ce6db537 100644
--- a/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContainer.test.tsx
+++ b/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContainer.test.tsx
@@ -1,37 +1,18 @@
 import React from 'react';
-import { render, screen } from '@testing-library/react';
+import { render, screen, waitFor } from '@testing-library/react';
 import { PromQueryBuilderContainer } from './PromQueryBuilderContainer';
 import { PrometheusDatasource } from '../../datasource';
 import { EmptyLanguageProviderMock } from '../../language_provider.mock';
 import PromQlLanguageProvider from '../../language_provider';
 import { addOperation } from '../shared/OperationList.testUtils';
+import { PromQuery } from '../../types';
+import userEvent from '@testing-library/user-event';
+import { getOperationParamId } from '../shared/operationUtils';
 
 describe('PromQueryBuilderContainer', () => {
   it('translates query between string and model', async () => {
-    const props = {
-      query: {
-        expr: 'metric_test{job="testjob"}',
-        refId: 'A',
-      },
-      datasource: new PrometheusDatasource(
-        {
-          id: 1,
-          uid: '',
-          type: 'prometheus',
-          name: 'prom-test',
-          access: 'proxy',
-          url: '',
-          jsonData: {},
-          meta: {} as any,
-        },
-        undefined,
-        undefined,
-        new EmptyLanguageProviderMock() as unknown as PromQlLanguageProvider
-      ),
-      onChange: jest.fn(),
-      onRunQuery: () => {},
-    };
-    render(<PromQueryBuilderContainer {...props} />);
+    const { props } = setup({ expr: 'rate(metric_test{job="testjob"}[$__rate_interval])' });
+
     expect(screen.getByText('metric_test')).toBeInTheDocument();
     addOperation('Range functions', 'Rate');
     expect(props.onChange).toBeCalledWith({
@@ -39,4 +20,41 @@ describe('PromQueryBuilderContainer', () => {
       refId: 'A',
     });
   });
+
+  it('Can add rest param', async () => {
+    const { container } = setup({ expr: 'sum(ALERTS)' });
+    userEvent.click(screen.getByTestId('operations.0.add-rest-param'));
+
+    waitFor(() => {
+      expect(container.querySelector(`${getOperationParamId(0, 0)}`)).toBeInTheDocument();
+    });
+  });
 });
+
+function setup(queryOverrides: Partial<PromQuery> = {}) {
+  const languageProvider = new EmptyLanguageProviderMock() as unknown as PromQlLanguageProvider;
+  const datasource = new PrometheusDatasource(
+    {
+      url: '',
+      jsonData: {},
+      meta: {} as any,
+    } as any,
+    undefined,
+    undefined,
+    languageProvider
+  );
+
+  const props = {
+    datasource,
+    query: {
+      refId: 'A',
+      expr: '',
+      ...queryOverrides,
+    },
+    onRunQuery: jest.fn(),
+    onChange: jest.fn(),
+  };
+
+  const { container } = render(<PromQueryBuilderContainer {...props} />);
+  return { languageProvider, datasource, container, props };
+}
diff --git a/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContainer.tsx b/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContainer.tsx
index 791e97f3393..5df01f76347 100644
--- a/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContainer.tsx
+++ b/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContainer.tsx
@@ -1,5 +1,5 @@
 import { PanelData } from '@grafana/data';
-import React from 'react';
+import React, { useEffect, useState } from 'react';
 
 import { PrometheusDatasource } from '../../datasource';
 import { PromQuery } from '../../types';
@@ -17,25 +17,40 @@ export interface Props {
   data?: PanelData;
 }
 
+export interface State {
+  visQuery?: PromVisualQuery;
+  expr: string;
+}
+
 /**
  * This component is here just to contain the translation logic between string query and the visual query builder model.
- * @param props
- * @constructor
  */
 export function PromQueryBuilderContainer(props: Props) {
   const { query, onChange, onRunQuery, datasource, data } = props;
+  const [state, setState] = useState<State>({ expr: query.expr });
 
-  const visQuery = buildVisualQueryFromString(query.expr || '').query;
+  // Only rebuild visual query if expr changes from outside
+  useEffect(() => {
+    if (!state.visQuery || state.expr !== query.expr) {
+      const result = buildVisualQueryFromString(query.expr || '');
+      setState({ visQuery: result.query, expr: query.expr });
+    }
+  }, [query.expr, state.visQuery, state.expr]);
 
   const onVisQueryChange = (newVisQuery: PromVisualQuery) => {
     const rendered = promQueryModeller.renderQuery(newVisQuery);
     onChange({ ...query, expr: rendered });
+    setState({ visQuery: newVisQuery, expr: rendered });
   };
 
+  if (!state.visQuery) {
+    return null;
+  }
+
   return (
     <>
       <PromQueryBuilder
-        query={visQuery}
+        query={state.visQuery}
         datasource={datasource}
         onChange={onVisQueryChange}
         onRunQuery={onRunQuery}
diff --git a/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContext.tsx b/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContext.tsx
deleted file mode 100644
index 43c7e3f42a7..00000000000
--- a/public/app/plugins/datasource/prometheus/querybuilder/components/PromQueryBuilderContext.tsx
+++ /dev/null
@@ -1,12 +0,0 @@
-import React from 'react';
-import { PrometheusDatasource } from '../../datasource';
-import { PromVisualQuery } from '../types';
-
-export interface PromQueryBuilderContextType {
-  query: PromVisualQuery;
-  datasource: PrometheusDatasource;
-}
-
-export const PromQueryBuilderContext = React.createContext<PromQueryBuilderContextType>(
-  {} as any as PromQueryBuilderContextType
-);
diff --git a/public/app/plugins/datasource/prometheus/querybuilder/shared/OperationEditor.tsx b/public/app/plugins/datasource/prometheus/querybuilder/shared/OperationEditor.tsx
index a657a69d88b..f848f645fac 100644
--- a/public/app/plugins/datasource/prometheus/querybuilder/shared/OperationEditor.tsx
+++ b/public/app/plugins/datasource/prometheus/querybuilder/shared/OperationEditor.tsx
@@ -14,6 +14,7 @@ import {
 import { OperationInfoButton } from './OperationInfoButton';
 import { OperationName } from './OperationName';
 import { getOperationParamEditor } from './OperationParamEditor';
+import { getOperationParamId } from './operationUtils';
 
 export interface Props {
   operation: QueryBuilderOperation;
@@ -66,7 +67,9 @@ export function OperationEditor({
 
     operationElements.push(
       <div className={styles.paramRow} key={`${paramIndex}-1`}>
-        <div className={styles.paramName}>{paramDef.name}</div>
+        <label className={styles.paramName} htmlFor={getOperationParamId(index, paramIndex)}>
+          {paramDef.name}
+        </label>
         <div className={styles.paramValue}>
           <Stack gap={0.5} direction="row" alignItems="center" wrap={false}>
             <Editor
@@ -74,6 +77,7 @@ export function OperationEditor({
               paramDef={paramDef}
               value={operation.params[paramIndex]}
               operation={operation}
+              operationIndex={index}
               onChange={onParamValueChanged}
               onRunQuery={onRunQuery}
               query={query}
@@ -81,6 +85,7 @@ export function OperationEditor({
             />
             {paramDef.restParam && (operation.params.length > def.params.length || paramDef.optional) && (
               <Button
+                data-testid={`operations.${index}.remove-rest-param`}
                 size="sm"
                 fill="text"
                 icon="times"
@@ -100,7 +105,7 @@ export function OperationEditor({
   if (def.params.length > 0) {
     const lastParamDef = def.params[def.params.length - 1];
     if (lastParamDef.restParam) {
-      restParam = renderAddRestParamButton(lastParamDef, onAddRestParam, operation.params.length, styles);
+      restParam = renderAddRestParamButton(lastParamDef, onAddRestParam, index, operation.params.length, styles);
     }
   }
 
@@ -111,7 +116,7 @@ export function OperationEditor({
           className={styles.card}
           ref={provided.innerRef}
           {...provided.draggableProps}
-          data-testid={`operation-wrapper-for-${operation.id}`}
+          data-testid={`operations.${index}.wrapper`}
         >
           <div className={styles.header} {...provided.dragHandleProps}>
             <OperationName
@@ -151,12 +156,20 @@ export function OperationEditor({
 function renderAddRestParamButton(
   paramDef: QueryBuilderOperationParamDef,
   onAddRestParam: () => void,
+  operationIndex: number,
   paramIndex: number,
   styles: OperationEditorStyles
 ) {
   return (
     <div className={styles.restParam} key={`${paramIndex}-2`}>
-      <Button size="sm" icon="plus" title={`Add ${paramDef.name}`} variant="secondary" onClick={onAddRestParam}>
+      <Button
+        size="sm"
+        icon="plus"
+        title={`Add ${paramDef.name}`}
+        variant="secondary"
+        onClick={onAddRestParam}
+        data-testid={`operations.${operationIndex}.add-rest-param`}
+      >
         {paramDef.name}
       </Button>
     </div>
diff --git a/public/app/plugins/datasource/prometheus/querybuilder/shared/OperationParamEditor.tsx b/public/app/plugins/datasource/prometheus/querybuilder/shared/OperationParamEditor.tsx
index d89ef0f34b5..168d880342c 100644
--- a/public/app/plugins/datasource/prometheus/querybuilder/shared/OperationParamEditor.tsx
+++ b/public/app/plugins/datasource/prometheus/querybuilder/shared/OperationParamEditor.tsx
@@ -2,6 +2,7 @@ import { toOption } from '@grafana/data';
 import { Input, Select } from '@grafana/ui';
 import React, { ComponentType } from 'react';
 import { QueryBuilderOperationParamDef, QueryBuilderOperationParamEditorProps } from '../shared/types';
+import { getOperationParamId } from './operationUtils';
 
 export function getOperationParamEditor(
   paramDef: QueryBuilderOperationParamDef
@@ -20,6 +21,7 @@ export function getOperationParamEditor(
 function SimpleInputParamEditor(props: QueryBuilderOperationParamEditorProps) {
   return (
     <Input
+      id={getOperationParamId(props.operationIndex, props.index)}
       defaultValue={props.value ?? ''}
       onKeyDown={(evt) => {
         if (evt.key === 'Enter') {
@@ -36,7 +38,13 @@ function SimpleInputParamEditor(props: QueryBuilderOperationParamEditorProps) {
   );
 }
 
-function SelectInputParamEditor({ paramDef, value, index, onChange }: QueryBuilderOperationParamEditorProps) {
+function SelectInputParamEditor({
+  paramDef,
+  value,
+  index,
+  operationIndex,
+  onChange,
+}: QueryBuilderOperationParamEditorProps) {
   const selectOptions = paramDef.options!.map((option) => ({
     label: option as string,
     value: option as string,
@@ -44,6 +52,7 @@ function SelectInputParamEditor({ paramDef, value, index, onChange }: QueryBuild
 
   return (
     <Select
+      id={getOperationParamId(operationIndex, index)}
       menuShouldPortal
       value={toOption(value as string)}
       options={selectOptions}
diff --git a/public/app/plugins/datasource/prometheus/querybuilder/shared/operationUtils.ts b/public/app/plugins/datasource/prometheus/querybuilder/shared/operationUtils.ts
index 8d991f3ee55..dc3187e34df 100644
--- a/public/app/plugins/datasource/prometheus/querybuilder/shared/operationUtils.ts
+++ b/public/app/plugins/datasource/prometheus/querybuilder/shared/operationUtils.ts
@@ -49,3 +49,7 @@ export function defaultAddOperationHandler<T extends QueryWithOperations>(def: Q
 export function getPromAndLokiOperationDisplayName(funcName: string) {
   return capitalize(funcName.replace(/_/g, ' '));
 }
+
+export function getOperationParamId(operationIndex: number, paramIndex: number) {
+  return `operations.${operationIndex}.param.${paramIndex}`;
+}
diff --git a/public/app/plugins/datasource/prometheus/querybuilder/shared/types.ts b/public/app/plugins/datasource/prometheus/querybuilder/shared/types.ts
index db088500f77..3a219472e4b 100644
--- a/public/app/plugins/datasource/prometheus/querybuilder/shared/types.ts
+++ b/public/app/plugins/datasource/prometheus/querybuilder/shared/types.ts
@@ -75,8 +75,10 @@ export interface QueryBuilderOperationEditorProps {
 export interface QueryBuilderOperationParamEditorProps {
   value?: QueryBuilderOperationParamValue;
   paramDef: QueryBuilderOperationParamDef;
+  /** Parameter index */
   index: number;
   operation: QueryBuilderOperation;
+  operationIndex: number;
   query: any;
   datasource: DataSourceApi;
   onChange: (index: number, value: QueryBuilderOperationParamValue) => void;