[Lens] Do not reset columns on incomplete switch before closing flyout (

#102876)

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

@@ -117,21 +117,14 @@ export function DimensionEditor(props: DimensionEditorProps) {
   const setStateWrapper = (
     setter: IndexPatternLayer | ((prevLayer: IndexPatternLayer) => IndexPatternLayer)
   ) => {
-    const prevOperationType =
-      operationDefinitionMap[state.layers[layerId].columns[columnId]?.operationType]?.input;
-
     const hypotheticalLayer = typeof setter === 'function' ? setter(state.layers[layerId]) : setter;
-    const hasIncompleteColumns = Boolean(hypotheticalLayer.incompleteColumns?.[columnId]);
     setState(
       (prevState) => {
         const layer = typeof setter === 'function' ? setter(prevState.layers[layerId]) : setter;
         return mergeLayer({ state: prevState, layerId, newLayer: layer });
       },
       {
-        isDimensionComplete:
-          prevOperationType === 'fullReference'
-            ? !hasIncompleteColumns
-            : Boolean(hypotheticalLayer.columns[columnId]),
+        isDimensionComplete: Boolean(hypotheticalLayer.columns[columnId]),
       }
     );
   };

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

@@ -908,20 +908,21 @@ describe('IndexPatternDimensionEditorPanel', () => {
       });
     });
 
-    it('should clean up when transitioning from incomplete reference-based operations to field operation', () => {
+    it('should keep current state and write incomplete column when transitioning from incomplete reference-based operations to field operation', () => {
+      const baseState = getStateWithColumns({
+        ...defaultProps.state.layers.first.columns,
+        col2: {
+          label: 'Counter rate',
+          dataType: 'number',
+          isBucketed: false,
+          operationType: 'counter_rate',
+          references: ['ref'],
+        },
+      });
       wrapper = mount(
         <IndexPatternDimensionEditorComponent
           {...defaultProps}
-          state={getStateWithColumns({
-            ...defaultProps.state.layers.first.columns,
-            col2: {
-              label: 'Counter rate',
-              dataType: 'number',
-              isBucketed: false,
-              operationType: 'counter_rate',
-              references: ['ref'],
-            },
-          })}
+          state={baseState}
           columnId={'col2'}
         />
       );
@@ -932,15 +933,12 @@ describe('IndexPatternDimensionEditorPanel', () => {
         .simulate('click');
 
       // Now check that the dimension gets cleaned up on state update
-      expect(setState.mock.calls[0]).toEqual([
-        expect.any(Function),
-        { isDimensionComplete: false },
-      ]);
+      expect(setState.mock.calls[0]).toEqual([expect.any(Function), { isDimensionComplete: true }]);
       expect(setState.mock.calls[0][0](state)).toEqual({
-        ...state,
+        ...baseState,
         layers: {
           first: {
-            ...state.layers.first,
+            ...baseState.layers.first,
             incompleteColumns: {
               col2: { operationType: 'average' },
             },

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

@@ -1917,6 +1917,54 @@ describe('state_helpers', () => {
           })
         );
       });
+
+      it('should keep state and set incomplete column on incompatible switch', () => {
+        const layer: IndexPatternLayer = {
+          indexPatternId: '1',
+          columnOrder: ['metric', 'ref'],
+          columns: {
+            metric: {
+              dataType: 'number' as const,
+              isBucketed: false,
+              sourceField: 'source',
+              operationType: 'unique_count' as const,
+              filter: { language: 'kuery', query: 'bytes > 4000' },
+              timeShift: '3h',
+              label: 'Cardinality',
+              customLabel: true,
+            },
+            ref: {
+              label: 'Reference',
+              dataType: 'number',
+              isBucketed: false,
+              operationType: 'differences',
+              references: ['metric'],
+              filter: { language: 'kuery', query: 'bytes > 4000' },
+              timeShift: '3h',
+            },
+          },
+        };
+        const result = replaceColumn({
+          layer,
+          indexPattern,
+          columnId: 'ref',
+          op: 'sum',
+          visualizationGroups: [],
+        });
+        expect(result.columnOrder).toEqual(layer.columnOrder);
+        expect(result.columns).toEqual(layer.columns);
+        expect(result.incompleteColumns).toEqual({
+          ref: {
+            operationType: 'sum',
+            filter: {
+              language: 'kuery',
+              query: 'bytes > 4000',
+            },
+            timeScale: undefined,
+            timeShift: '3h',
+          },
+        });
+      });
     });
 
     it('should allow making a replacement on an operation that is being referenced, even if it ends up invalid', () => {

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

@@ -19,6 +19,7 @@ import {
   OperationType,
   IndexPatternColumn,
   RequiredReference,
+  OperationDefinition,
   GenericOperationDefinition,
 } from './definitions';
 import type {
@@ -532,20 +533,15 @@ export function replaceColumn({
       );
     }
 
-    // This logic comes after the transitions because they need to look at previous columns
-    if (previousDefinition.input === 'fullReference') {
-      (previousColumn as ReferenceBasedIndexPatternColumn).references.forEach((id: string) => {
-        tempLayer = deleteColumn({
-          layer: tempLayer,
-          columnId: id,
-          indexPattern,
-        });
-      });
-    }
-
     if (operationDefinition.input === 'none') {
       let newColumn = operationDefinition.buildColumn({ ...baseOptions, layer: tempLayer });
       newColumn = copyCustomLabel(newColumn, previousColumn);
+      tempLayer = removeOrphanedColumns(
+        previousDefinition,
+        previousColumn,
+        tempLayer,
+        indexPattern
+      );
 
       const newLayer = { ...tempLayer, columns: { ...tempLayer.columns, [columnId]: newColumn } };
       return updateDefaultLabels(
@@ -564,7 +560,6 @@ export function replaceColumn({
       } & ColumnAdvancedParams = { operationType: op };
       // if no field is available perform a full clean of the column from the layer
       if (previousDefinition.input === 'fullReference') {
-        tempLayer = deleteColumn({ layer: tempLayer, columnId, indexPattern });
         const previousReferenceId = (previousColumn as ReferenceBasedIndexPatternColumn)
           .references[0];
         const referenceColumn = layer.columns[previousReferenceId];
@@ -598,6 +593,8 @@ export function replaceColumn({
       };
     }
 
+    tempLayer = removeOrphanedColumns(previousDefinition, previousColumn, tempLayer, indexPattern);
+
     let newColumn = operationDefinition.buildColumn({ ...baseOptions, layer: tempLayer, field });
     if (!shouldResetLabel) {
       newColumn = copyCustomLabel(newColumn, previousColumn);
@@ -637,6 +634,27 @@ export function replaceColumn({
   }
 }
 
+function removeOrphanedColumns(
+  previousDefinition:
+    | OperationDefinition<IndexPatternColumn, 'field'>
+    | OperationDefinition<IndexPatternColumn, 'none'>
+    | OperationDefinition<IndexPatternColumn, 'fullReference'>,
+  previousColumn: IndexPatternColumn,
+  tempLayer: IndexPatternLayer,
+  indexPattern: IndexPattern
+) {
+  if (previousDefinition.input === 'fullReference') {
+    (previousColumn as ReferenceBasedIndexPatternColumn).references.forEach((id: string) => {
+      tempLayer = deleteColumn({
+        layer: tempLayer,
+        columnId: id,
+        indexPattern,
+      });
+    });
+  }
+  return tempLayer;
+}
+
 export function canTransition({
   layer,
   columnId,

x-pack/test/functional/apps/lens/smokescreen.ts

@@ -604,7 +604,7 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
       );
     });
 
-    it('should not leave an incomplete column in the visualization config with reference-based operations', async () => {
+    it('should revert to previous configuration and not leave an incomplete column in the visualization config with reference-based operations', async () => {
       await PageObjects.visualize.navigateToNewVisualization();
       await PageObjects.visualize.clickVisType('lens');
       await PageObjects.lens.goToTimeRange();
@@ -636,7 +636,7 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
       await PageObjects.lens.closeDimensionEditor();
 
       expect(await PageObjects.lens.getDimensionTriggerText('lnsXY_yDimensionPanel')).to.eql(
-        undefined
+        'Moving average of Count of records'
       );
     });