From ed477054c5560b50ede4245106954509617d65db Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Mon, 31 May 2021 11:42:53 +0200
Subject: [PATCH] keep current quick function

---
 .../dimension_panel/dimension_editor.tsx      | 28 ++++++++++++++++---
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/x-pack/plugins/lens/public/indexpattern_datasource/dimension_panel/dimension_editor.tsx b/x-pack/plugins/lens/public/indexpattern_datasource/dimension_panel/dimension_editor.tsx
index 0540605eecf12..5be05338fdc2e 100644
--- a/x-pack/plugins/lens/public/indexpattern_datasource/dimension_panel/dimension_editor.tsx
+++ b/x-pack/plugins/lens/public/indexpattern_datasource/dimension_panel/dimension_editor.tsx
@@ -22,6 +22,7 @@ import {
   EuiTab,
   EuiCallOut,
 } from '@elastic/eui';
+import useUnmount from 'react-use/lib/useUnmount';
 import { IndexPatternDimensionEditorProps } from './dimension_panel';
 import { OperationSupportMatrix } from './operation_support';
 import { IndexPatternColumn } from '../indexpattern';
@@ -40,7 +41,7 @@ import { mergeLayer } from '../state_helpers';
 import { FieldSelect } from './field_select';
 import { hasField, fieldIsInvalid } from '../utils';
 import { BucketNestingEditor } from './bucket_nesting_editor';
-import { IndexPattern, IndexPatternLayer } from '../types';
+import { IndexPattern, IndexPatternLayer, IndexPatternPrivateState } from '../types';
 import { trackUiEvent } from '../../lens_ui_telemetry';
 import { FormatSelector } from './format_selector';
 import { ReferenceEditor } from './reference_editor';
@@ -84,7 +85,7 @@ const LabelInput = ({ value, onChange }: { value: string; onChange: (value: stri
 
 export function DimensionEditor(props: DimensionEditorProps) {
   const {
-    selectedColumn,
+    selectedColumn: upstreamSelectedColumn,
     operationSupportMatrix,
     state,
     columnId,
@@ -106,6 +107,17 @@ export function DimensionEditor(props: DimensionEditorProps) {
   };
   const { fieldByOperation, operationWithoutField } = operationSupportMatrix;
 
+  const [previousQuickFunctionState, setPreviousQuickFunctionState] = useState<
+    IndexPatternPrivateState | undefined
+  >(undefined);
+
+  const [temporaryQuickFunction, setQuickFunction] = useState(false);
+
+  const selectedColumn =
+    temporaryQuickFunction && previousQuickFunctionState
+      ? previousQuickFunctionState.layers[layerId].columns[columnId]
+      : upstreamSelectedColumn;
+
   const selectedOperationDefinition =
     selectedColumn && operationDefinitionMap[selectedColumn.operationType];
 
@@ -113,12 +125,21 @@ export function DimensionEditor(props: DimensionEditorProps) {
     Boolean(selectedOperationDefinition?.type === 'formula')
   );
 
+  useUnmount(() => {
+    if (temporaryQuickFunction && previousQuickFunctionState) {
+      setState(previousQuickFunctionState, {
+        shouldClose: true,
+      });
+    }
+  });
+
   const setStateWrapper = (
     setter: IndexPatternLayer | ((prevLayer: IndexPatternLayer) => IndexPatternLayer),
     shouldClose?: boolean
   ) => {
     if (selectedOperationDefinition?.type === 'formula' && !temporaryQuickFunction) {
       setChangedFormula(true);
+      setPreviousQuickFunctionState(undefined);
     } else {
       setChangedFormula(false);
     }
@@ -152,8 +173,6 @@ export function DimensionEditor(props: DimensionEditorProps) {
 
   const ParamEditor = selectedOperationDefinition?.paramEditor;
 
-  const [temporaryQuickFunction, setQuickFunction] = useState(false);
-
   const possibleOperations = useMemo(() => {
     return Object.values(operationDefinitionMap)
       .filter(({ hidden }) => !hidden)
@@ -615,6 +634,7 @@ export function DimensionEditor(props: DimensionEditorProps) {
               onClick={() => {
                 if (selectedColumn?.operationType !== 'formula') {
                   setQuickFunction(false);
+                  setPreviousQuickFunctionState(state);
                   const newLayer = insertOrReplaceColumn({
                     layer: props.state.layers[props.layerId],
                     indexPattern: currentIndexPattern,