[Lens] Clean full reference operation error state when switching to other operation (#87064)

* 🐛 Make sure to check incomplete columns before saved ones

* ⚗️ Try a different approach

* 🐛 Isolate the fullRef to field transition and trigger the removal + tests

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

x-pack/plugins/lens/public/editor_frame_service/editor_frame/config_panel/layer_panel.test.tsx

@@ -332,7 +332,7 @@ describe('LayerPanel', () => {
             columns: {},
             columnOrder: [],
           },
-          true
+          { shouldReplaceDimension: true }
         );
       });
       expect(updateAll).toHaveBeenCalled();

x-pack/plugins/lens/public/editor_frame_service/editor_frame/config_panel/layer_panel.tsx

@@ -506,17 +506,32 @@ export function LayerPanel(
                       columnId: activeId,
                       filterOperations: activeGroup.filterOperations,
                       dimensionGroups: groups,
-                      setState: (newState: unknown, shouldUpdateVisualization?: boolean) => {
-                        if (shouldUpdateVisualization) {
+                      setState: (
+                        newState: unknown,
+                        {
+                          shouldReplaceDimension,
+                          shouldRemoveDimension,
+                        }: {
+                          shouldReplaceDimension?: boolean;
+                          shouldRemoveDimension?: boolean;
+                        } = {}
+                      ) => {
+                        if (shouldReplaceDimension || shouldRemoveDimension) {
                           props.updateAll(
                             datasourceId,
                             newState,
-                            activeVisualization.setDimension({
-                              layerId,
-                              groupId: activeGroup.groupId,
-                              columnId: activeId,
-                              prevState: props.visualizationState,
-                            })
+                            shouldRemoveDimension
+                              ? activeVisualization.removeDimension({
+                                  layerId,
+                                  columnId: activeId,
+                                  prevState: props.visualizationState,
+                                })
+                              : activeVisualization.setDimension({
+                                  layerId,
+                                  groupId: activeGroup.groupId,
+                                  columnId: activeId,
+                                  prevState: props.visualizationState,
+                                })
                           );
                         } else {
                           props.updateDatasource(datasourceId, newState);

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

@@ -122,7 +122,14 @@ export function DimensionEditor(props: DimensionEditorProps) {
   const { fieldByOperation, operationWithoutField } = operationSupportMatrix;
 
   const setStateWrapper = (layer: IndexPatternLayer) => {
-    setState(mergeLayer({ state, layerId, newLayer: layer }), Boolean(layer.columns[columnId]));
+    const hasIncompleteColumns = Boolean(layer.incompleteColumns?.[columnId]);
+    const prevOperationType =
+      operationDefinitionMap[state.layers[layerId].columns[columnId]?.operationType]?.input;
+    setState(mergeLayer({ state, layerId, newLayer: layer }), {
+      shouldReplaceDimension: Boolean(layer.columns[columnId]),
+      // clear the dimension if there's an incomplete column pending && previous operation was a fullReference operation
+      shouldRemoveDimension: Boolean(hasIncompleteColumns && prevOperationType === 'fullReference'),
+    });
   };
 
   const selectedOperationDefinition =

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

@@ -492,7 +492,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      true
+      { shouldRemoveDimension: false, shouldReplaceDimension: true }
     );
   });
 
@@ -525,7 +525,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      true
+      { shouldRemoveDimension: false, shouldReplaceDimension: true }
     );
   });
 
@@ -559,7 +559,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      true
+      { shouldRemoveDimension: false, shouldReplaceDimension: true }
     );
   });
 
@@ -629,7 +629,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      true
+      { shouldRemoveDimension: false, shouldReplaceDimension: true }
     );
   });
 
@@ -667,7 +667,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      true
+      { shouldRemoveDimension: false, shouldReplaceDimension: true }
     );
   });
 
@@ -696,7 +696,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
             },
           },
         },
-        true
+        { shouldRemoveDimension: false, shouldReplaceDimension: true }
       );
     });
 
@@ -816,7 +816,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
             },
           },
         },
-        false
+        { shouldRemoveDimension: false, shouldReplaceDimension: false }
       );
 
       const comboBox = wrapper
@@ -847,7 +847,47 @@ describe('IndexPatternDimensionEditorPanel', () => {
             },
           },
         },
-        true
+        { shouldRemoveDimension: false, shouldReplaceDimension: true }
+      );
+    });
+
+    it('should clean up when transitioning from incomplete reference-based operations to field operation', () => {
+      wrapper = mount(
+        <IndexPatternDimensionEditorComponent
+          {...defaultProps}
+          state={getStateWithColumns({
+            ...defaultProps.state.layers.first.columns,
+            col2: {
+              label: 'Counter rate',
+              dataType: 'number',
+              isBucketed: false,
+              operationType: 'counter_rate',
+              references: ['ref'],
+            },
+          })}
+          columnId={'col2'}
+        />
+      );
+
+      // Transition to a field operation (incompatible)
+      wrapper
+        .find('button[data-test-subj="lns-indexPatternDimension-avg incompatible"]')
+        .simulate('click');
+
+      // Now check that the dimension gets cleaned up on state update
+      expect(setState).toHaveBeenCalledWith(
+        {
+          ...state,
+          layers: {
+            first: {
+              ...state.layers.first,
+              incompleteColumns: {
+                col2: { operationType: 'avg' },
+              },
+            },
+          },
+        },
+        { shouldRemoveDimension: true, shouldReplaceDimension: false }
       );
     });
 
@@ -945,7 +985,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
             },
           },
         },
-        true
+        { shouldRemoveDimension: false, shouldReplaceDimension: true }
       );
     });
   });
@@ -1037,7 +1077,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
             },
           },
         },
-        true
+        { shouldRemoveDimension: false, shouldReplaceDimension: true }
       );
     });
 
@@ -1068,7 +1108,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
             },
           },
         },
-        true
+        { shouldRemoveDimension: false, shouldReplaceDimension: true }
       );
     });
 
@@ -1097,7 +1137,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
             },
           },
         },
-        true
+        { shouldRemoveDimension: false, shouldReplaceDimension: true }
       );
     });
 
@@ -1126,7 +1166,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
             },
           },
         },
-        true
+        { shouldRemoveDimension: false, shouldReplaceDimension: true }
       );
     });
 
@@ -1155,7 +1195,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
             },
           },
         },
-        true
+        { shouldRemoveDimension: false, shouldReplaceDimension: true }
       );
     });
 
@@ -1185,7 +1225,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
             },
           },
         },
-        true
+        { shouldRemoveDimension: false, shouldReplaceDimension: true }
       );
     });
   });
@@ -1238,7 +1278,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      false
+      { shouldRemoveDimension: false, shouldReplaceDimension: false }
     );
 
     const comboBox = wrapper
@@ -1268,7 +1308,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      true
+      { shouldRemoveDimension: false, shouldReplaceDimension: true }
     );
   });
 
@@ -1311,7 +1351,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      true
+      { shouldRemoveDimension: false, shouldReplaceDimension: true }
     );
   });
 
@@ -1337,7 +1377,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      true
+      { shouldRemoveDimension: false, shouldReplaceDimension: true }
     );
   });
 
@@ -1474,7 +1514,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      true
+      { shouldRemoveDimension: false, shouldReplaceDimension: true }
     );
   });
 
@@ -1523,7 +1563,7 @@ describe('IndexPatternDimensionEditorPanel', () => {
           },
         },
       },
-      false
+      { shouldRemoveDimension: false, shouldReplaceDimension: false }
     );
   });
 

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

@@ -2174,5 +2174,34 @@ describe('state_helpers', () => {
       expect(mock).toHaveBeenCalled();
       expect(errors).toHaveLength(1);
     });
+
+    it('should consider incompleteColumns before layer columns', () => {
+      const savedRef = jest.fn().mockReturnValue(['error 1']);
+      const incompleteRef = jest.fn();
+      operationDefinitionMap.testReference.getErrorMessage = savedRef;
+      // @ts-expect-error invalid type, just need a single function on it
+      operationDefinitionMap.testIncompleteReference = {
+        getErrorMessage: incompleteRef,
+      };
+
+      const errors = getErrorMessages({
+        indexPatternId: '1',
+        columnOrder: [],
+        columns: {
+          col1:
+            // @ts-expect-error not statically analyzed
+            { operationType: 'testReference', references: [] },
+        },
+        incompleteColumns: {
+          // @ts-expect-error not statically analyzed
+          col1: { operationType: 'testIncompleteReference' },
+        },
+      });
+      expect(savedRef).not.toHaveBeenCalled();
+      expect(incompleteRef).toHaveBeenCalled();
+      expect(errors).toBeUndefined();
+
+      delete operationDefinitionMap.testIncompleteReference;
+    });
   });
 });

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

@@ -316,6 +316,10 @@ export function replaceColumn({
     }
 
     if (!field) {
+      // 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 });
+      }
       return {
         ...tempLayer,
         incompleteColumns: {
@@ -862,9 +866,12 @@ export function updateLayerIndexPattern(
  */
 export function getErrorMessages(layer: IndexPatternLayer): string[] | undefined {
   const errors: string[] = [];
-
   Object.entries(layer.columns).forEach(([columnId, column]) => {
-    const def = operationDefinitionMap[column.operationType];
+    // If we're transitioning to another operation, check for "new" incompleteColumns rather
+    // than "old" saved operation on the layer
+    const columnFinalRef =
+      layer.incompleteColumns?.[columnId]?.operationType || column.operationType;
+    const def = operationDefinitionMap[columnFinalRef];
     if (def.getErrorMessage) {
       errors.push(...(def.getErrorMessage(layer, columnId) ?? []));
     }

x-pack/plugins/lens/public/types.ts

@@ -243,7 +243,10 @@ export type DatasourceDimensionProps<T> = SharedDimensionProps & {
 // The only way a visualization has to restrict the query building
 export type DatasourceDimensionEditorProps<T = unknown> = DatasourceDimensionProps<T> & {
   // Not a StateSetter because we have this unique use case of determining valid columns
-  setState: (newState: Parameters<StateSetter<T>>[0], publishToVisualization?: boolean) => void;
+  setState: (
+    newState: Parameters<StateSetter<T>>[0],
+    publishToVisualization?: { shouldReplaceDimension?: boolean; shouldRemoveDimension?: boolean }
+  ) => void;
   core: Pick<CoreSetup, 'http' | 'notifications' | 'uiSettings'>;
   dateRange: DateRange;
   dimensionGroups: VisualizationDimensionGroupConfig[];