[Lens] Use a setter function for the dimension panel

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

@@ -119,6 +119,25 @@ describe('ConfigPanel', () => {
     expect(component.find(LayerPanel).exists()).toBe(false);
   });
 
+  it('allow datasources and visualizations to use setters', () => {
+    const props = getDefaultProps();
+    const component = mountWithIntl(<LayerPanels {...props} />);
+    const { updateDatasource, updateAll } = component.find(LayerPanel).props();
+
+    const updater = () => 'updated';
+    updateDatasource('ds1', updater);
+    expect(props.dispatch).toHaveBeenCalledTimes(1);
+    expect(props.dispatch.mock.calls[0][0].updater(props.datasourceStates.ds1.state)).toEqual(
+      'updated'
+    );
+
+    updateAll('ds1', updater, props.visualizationState);
+    expect(props.dispatch).toHaveBeenCalledTimes(2);
+    expect(props.dispatch.mock.calls[0][0].updater(props.datasourceStates.ds1.state)).toEqual(
+      'updated'
+    );
+  });
+
   describe('focus behavior when adding or removing layers', () => {
     it('should focus the only layer when resetting the layer', () => {
       const component = mountWithIntl(<LayerPanels {...getDefaultProps()} />, {

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

@@ -63,7 +63,8 @@ export function LayerPanels(
     () => (datasourceId: string, newState: unknown) => {
       dispatch({
         type: 'UPDATE_DATASOURCE_STATE',
-        updater: () => newState,
+        updater: (prevState: unknown) =>
+          typeof newState === 'function' ? newState(prevState) : newState,
         datasourceId,
         clearStagedPreview: false,
       });
@@ -76,12 +77,16 @@ export function LayerPanels(
         type: 'UPDATE_STATE',
         subType: 'UPDATE_ALL_STATES',
         updater: (prevState) => {
+          const updatedDatasourceState =
+            typeof newDatasourceState === 'function'
+              ? newDatasourceState(prevState.datasourceStates[datasourceId].state)
+              : newDatasourceState;
           return {
             ...prevState,
             datasourceStates: {
               ...prevState.datasourceStates,
               [datasourceId]: {
-                state: newDatasourceState,
+                state: updatedDatasourceState,
                 isLoading: false,
               },
             },

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

@@ -100,15 +100,26 @@ export function DimensionEditor(props: DimensionEditorProps) {
   };
   const { fieldByOperation, operationWithoutField } = operationSupportMatrix;
 
-  const setStateWrapper = (layer: IndexPatternLayer) => {
-    const hasIncompleteColumns = Boolean(layer.incompleteColumns?.[columnId]);
+  const setStateWrapper = (
+    setter: IndexPatternLayer | ((prevLayer: IndexPatternLayer) => IndexPatternLayer)
+  ) => {
+    const hypotheticalLayer = typeof setter === 'function' ? setter(state.layers[layerId]) : setter;
+    const hasIncompleteColumns = Boolean(hypotheticalLayer.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'),
-    });
+      operationDefinitionMap[hypotheticalLayer.columns[columnId]?.operationType]?.input;
+    setState(
+      (prevState) => {
+        const layer = typeof setter === 'function' ? setter(prevState.layers[layerId]) : setter;
+        return mergeLayer({ state: prevState, layerId, newLayer: layer });
+      },
+      {
+        shouldReplaceDimension: Boolean(hypotheticalLayer.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 =
@@ -337,8 +348,19 @@ export function DimensionEditor(props: DimensionEditorProps) {
                   key={index}
                   layer={state.layers[layerId]}
                   columnId={referenceId}
-                  updateLayer={(newLayer: IndexPatternLayer) => {
-                    setState(mergeLayer({ state, layerId, newLayer }));
+                  updateLayer={(
+                    setter:
+                      | IndexPatternLayer
+                      | ((prevLayer: IndexPatternLayer) => IndexPatternLayer)
+                  ) => {
+                    setState(
+                      mergeLayer({
+                        state,
+                        layerId,
+                        newLayer:
+                          typeof setter === 'function' ? setter(state.layers[layerId]) : setter,
+                      })
+                    );
                   }}
                   validation={validation}
                   currentIndexPattern={currentIndexPattern}