[7.x] [Lens] Keep the custom label when transitioning to/from Formula (#114270)

x-pack/plugins/lens/public/indexpattern_datasource/dimension_panel/dimension_editor.tsx

@@ -737,9 +737,9 @@ export function DimensionEditor(props: DimensionEditorProps) {
       />
       {TabContent}
 
-      {!isFullscreen && !currentFieldIsInvalid && temporaryState === 'none' && (
+      {!isFullscreen && !currentFieldIsInvalid && (
         <div className="lnsIndexPatternDimensionEditor__section lnsIndexPatternDimensionEditor__section--padded">
-          {!incompleteInfo && selectedColumn && (
+          {!incompleteInfo && selectedColumn && temporaryState === 'none' && (
             <LabelInput
               value={selectedColumn.label}
               onChange={(value) => {
@@ -762,7 +762,7 @@ export function DimensionEditor(props: DimensionEditorProps) {
             />
           )}
 
-          {!isFullscreen && !incompleteInfo && !hideGrouping && (
+          {!isFullscreen && !incompleteInfo && !hideGrouping && temporaryState === 'none' && (
             <BucketNestingEditor
               layer={state.layers[props.layerId]}
               columnId={props.columnId}

x-pack/plugins/lens/public/indexpattern_datasource/operations/layer_helpers.test.ts

@@ -1139,6 +1139,122 @@ describe('state_helpers', () => {
           }).columns.col1
         ).toEqual(expect.objectContaining({ label: 'Average of bytes' }));
       });
+
+      it('should carry over a custom label when transitioning to a managed reference', () => {
+        expect(
+          replaceColumn({
+            layer: {
+              indexPatternId: '1',
+              columnOrder: ['col1', 'col2'],
+              columns: {
+                col1: {
+                  label: 'MY CUSTOM LABEL',
+                  customLabel: true,
+                  dataType: 'string',
+                  isBucketed: true,
+                  operationType: 'terms',
+                  sourceField: 'source',
+                  params: {
+                    orderBy: { type: 'alphabetical' },
+                    orderDirection: 'asc',
+                    size: 5,
+                  },
+                },
+              },
+            },
+            indexPattern,
+            columnId: 'col1',
+            op: 'formula',
+            field: indexPattern.fields[2], // bytes field
+            visualizationGroups: [],
+            shouldResetLabel: undefined,
+          }).columns.col1
+        ).toEqual(expect.objectContaining({ label: 'MY CUSTOM LABEL' }));
+      });
+
+      it('should overwrite the current label when transitioning to a managed reference operation when not custom', () => {
+        expect(
+          replaceColumn({
+            layer: {
+              indexPatternId: '1',
+              columnOrder: ['col1', 'col2'],
+              columns: {
+                col1: {
+                  label: 'Average of bytes',
+                  dataType: 'number',
+                  isBucketed: false,
+                  operationType: 'average',
+                  sourceField: 'bytes',
+                },
+              },
+            },
+            indexPattern,
+            columnId: 'col1',
+            op: 'formula',
+            field: indexPattern.fields[2], // bytes field
+            visualizationGroups: [],
+            shouldResetLabel: undefined,
+          }).columns.col1
+        ).toEqual(expect.objectContaining({ label: 'average(bytes)' }));
+      });
+
+      it('should carry over a custom label when transitioning from a managed reference', () => {
+        expect(
+          replaceColumn({
+            layer: {
+              indexPatternId: '1',
+              columnOrder: ['col1', 'col2'],
+              columns: {
+                col1: {
+                  label: 'MY CUSTOM LABEL',
+                  customLabel: true,
+                  dataType: 'number',
+                  operationType: 'formula',
+                  isBucketed: false,
+                  scale: 'ratio',
+                  params: { isFormulaBroken: false, formula: 'average(bytes)' },
+                  references: [],
+                },
+              },
+            },
+            indexPattern,
+            columnId: 'col1',
+            op: 'average',
+            field: indexPattern.fields[2], // bytes field
+            visualizationGroups: [],
+            shouldResetLabel: undefined,
+          }).columns.col1
+        ).toEqual(expect.objectContaining({ label: 'MY CUSTOM LABEL' }));
+      });
+
+      it('should not carry over the managed reference default label to the new operation', () => {
+        expect(
+          replaceColumn({
+            layer: {
+              indexPatternId: '1',
+              columnOrder: ['col1', 'col2'],
+              columns: {
+                col1: {
+                  label: 'average(bytes)',
+                  customLabel: true,
+                  dataType: 'number',
+                  operationType: 'formula',
+                  isBucketed: false,
+                  scale: 'ratio',
+                  params: { isFormulaBroken: false, formula: 'average(bytes)' },
+                  references: [],
+                },
+              },
+            },
+            indexPattern,
+            columnId: 'col1',
+            op: 'average',
+            field: indexPattern.fields[2], // bytes field
+            visualizationGroups: [],
+            shouldResetLabel: undefined,
+          }).columns.col1
+        ).toEqual(expect.objectContaining({ label: 'Average of bytes' }));
+      });
     });
 
     it('should execute adjustments for other columns', () => {

x-pack/plugins/lens/public/indexpattern_datasource/operations/layer_helpers.ts

@@ -417,6 +417,17 @@ export function replaceColumn({
         field,
         visualizationGroups,
       });
+
+      // if the formula label is not the default one, propagate it to the new operation
+      if (
+        !shouldResetLabel &&
+        previousColumn.customLabel &&
+        previousColumn.label !==
+          previousDefinition.getDefaultLabel(previousColumn, indexPattern, tempLayer.columns)
+      ) {
+        hypotheticalLayer.columns[columnId].customLabel = true;
+        hypotheticalLayer.columns[columnId].label = previousColumn.label;
+      }
       if (hypotheticalLayer.incompleteColumns && hypotheticalLayer.incompleteColumns[columnId]) {
         return {
           ...layer,
@@ -500,13 +511,10 @@ export function replaceColumn({
     // TODO: Refactor all this to be more generic and know less about Formula
     // if managed it has to look at the full picture to have a seamless transition
     if (operationDefinition.input === 'managedReference') {
-      const newColumn = copyCustomLabel(
-        operationDefinition.buildColumn(
-          { ...baseOptions, layer: tempLayer },
-          previousColumn.params,
-          operationDefinitionMap
-        ),
-        previousColumn
+      const newColumn = operationDefinition.buildColumn(
+        { ...baseOptions, layer: tempLayer },
+        previousColumn.params,
+        operationDefinitionMap
       ) as FormulaIndexPatternColumn;
 
       // now remove the previous references
@@ -535,6 +543,17 @@ export function replaceColumn({
         newLayer = basicLayer;
       }
 
+      // when coming to Formula keep the custom label
+      const regeneratedColumn = newLayer.columns[columnId];
+      if (
+        !shouldResetLabel &&
+        regeneratedColumn.operationType !== previousColumn.operationType &&
+        previousColumn.customLabel
+      ) {
+        regeneratedColumn.customLabel = true;
+        regeneratedColumn.label = previousColumn.label;
+      }
+
       return updateDefaultLabels(
         {
           ...tempLayer,