layers removed by datasource get removed from visualization

elastic · Oct 21, 2022 · 2e249a4 · 2e249a4
1 parent d6c5ec9
commit 2e249a4
Show file tree

Hide file tree

Showing 7 changed files with 204 additions and 36 deletions.
diff --git a/x-pack/plugins/lens/public/datasources/form_based/form_based.test.ts b/x-pack/plugins/lens/public/datasources/form_based/form_based.test.ts
@@ -1665,12 +1665,74 @@ describe('IndexPattern Data Source', () => {
         currentIndexPatternId: '1',
       };
       expect(FormBasedDatasource.removeLayer(state, 'first')).toEqual({
-        ...state,
+        removedLayerIds: ['first'],
+        newState: {
+          ...state,
+          layers: {
+            second: {
+              indexPatternId: '2',
+              columnOrder: [],
+              columns: {},
+            },
+          },
+        },
+      });
+    });
+
+    it('should remove linked layers', () => {
+      const state = {
         layers: {
+          first: {
+            indexPatternId: '1',
+            columnOrder: [],
+            columns: {},
+          },
           second: {
             indexPatternId: '2',
             columnOrder: [],
             columns: {},
+            linkToLayers: ['first'],
+          },
+        },
+        currentIndexPatternId: '1',
+      };
+      expect(FormBasedDatasource.removeLayer(state, 'first')).toEqual({
+        removedLayerIds: ['first', 'second'],
+        newState: {
+          ...state,
+          layers: {},
+        },
+      });
+    });
+  });
+
+  describe('#clearLayer', () => {
+    it('should clear a layer', () => {
+      const state = {
+        layers: {
+          first: {
+            indexPatternId: '1',
+            columnOrder: ['some', 'order'],
+            columns: {
+              some: {} as GenericIndexPatternColumn,
+              columns: {} as GenericIndexPatternColumn,
+            },
+            linkToLayers: ['some-layer'],
+          },
+        },
+        currentIndexPatternId: '1',
+      };
+      expect(FormBasedDatasource.clearLayer(state, 'first')).toEqual({
+        removedLayerIds: [],
+        newState: {
+          ...state,
+          layers: {
+            first: {
+              indexPatternId: '1',
+              columnOrder: [],
+              columns: {},
+              linkToLayers: ['some-layer'],
+            },
           },
         },
       });
@@ -1693,9 +1755,14 @@ describe('IndexPattern Data Source', () => {
         },
         currentIndexPatternId: '1',
       };
-      expect(FormBasedDatasource.removeLayer(state, 'first')).toEqual({
-        ...state,
-        layers: {},
+      expect(FormBasedDatasource.clearLayer(state, 'first')).toEqual({
+        removedLayerIds: ['second'],
+        newState: {
+          ...state,
+          layers: {
+            first: state.layers.first,
+          },
+        },
       });
     });
   });

diff --git a/x-pack/plugins/lens/public/datasources/form_based/form_based.tsx b/x-pack/plugins/lens/public/datasources/form_based/form_based.tsx
@@ -218,27 +218,47 @@ export function getFormBasedDatasource({
     removeLayer(state: FormBasedPrivateState, layerId: string) {
       const newLayers = { ...state.layers };
       delete newLayers[layerId];
+      const removedLayerIds: string[] = [layerId];
 
       // delete layers linked to this layer
       Object.keys(newLayers).forEach((id) => {
         const linkedLayers = newLayers[id]?.linkToLayers;
         if (linkedLayers && linkedLayers.includes(layerId)) {
           delete newLayers[id];
+          removedLayerIds.push(id);
         }
       });
 
       return {
-        ...state,
-        layers: newLayers,
+        removedLayerIds,
+        newState: {
+          ...state,
+          layers: newLayers,
+        },
       };
     },
 
     clearLayer(state: FormBasedPrivateState, layerId: string) {
+      const newLayers = { ...state.layers };
+
+      const removedLayerIds: string[] = [];
+      // delete layers linked to this layer
+      Object.keys(newLayers).forEach((id) => {
+        const linkedLayers = newLayers[id]?.linkToLayers;
+        if (linkedLayers && linkedLayers.includes(layerId)) {
+          delete newLayers[id];
+          removedLayerIds.push(id);
+        }
+      });
+
       return {
-        ...state,
-        layers: {
-          ...state.layers,
-          [layerId]: blankLayer(state.currentIndexPatternId, state.layers[layerId].linkToLayers),
+        removedLayerIds,
+        newState: {
+          ...state,
+          layers: {
+            ...newLayers,
+            [layerId]: blankLayer(state.currentIndexPatternId, state.layers[layerId].linkToLayers),
+          },
         },
       };
     },

diff --git a/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.tsx b/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.tsx
@@ -281,18 +281,24 @@ export function getTextBasedDatasource({
       };
 
       return {
-        ...state,
-        layers: newLayers,
-        fieldList: state.fieldList,
+        removedLayerIds: [layerId],
+        newState: {
+          ...state,
+          layers: newLayers,
+          fieldList: state.fieldList,
+        },
       };
     },
 
     clearLayer(state: TextBasedPrivateState, layerId: string) {
       return {
-        ...state,
-        layers: {
-          ...state.layers,
-          [layerId]: { ...state.layers[layerId], columns: [] },
+        removedLayerIds: [],
+        newState: {
+          ...state,
+          layers: {
+            ...state.layers,
+            [layerId]: { ...state.layers[layerId], columns: [] },
+          },
         },
       };
     },

diff --git a/x-pack/plugins/lens/public/state_management/lens_slice.test.ts b/x-pack/plugins/lens/public/state_management/lens_slice.test.ts
@@ -39,6 +39,7 @@ describe('lensSlice', () => {
   let store: EnhancedStore<{ lens: LensAppState }>;
   beforeEach(() => {
     store = makeLensStore({}).store;
+    jest.clearAllMocks();
   });
   const customQuery = { query: 'custom' } as Query;
 
@@ -275,17 +276,21 @@ describe('lensSlice', () => {
         return {
           id: datasourceId,
           getPublicAPI: () => ({
-            datasourceId: 'testDatasource',
+            datasourceId,
             getOperationForColumnId: jest.fn(),
             getTableSpec: jest.fn(),
           }),
           getLayers: () => ['layer1'],
-          clearLayer: (layerIds: unknown, layerId: string) =>
-            (layerIds as string[]).map((id: string) =>
+          clearLayer: (layerIds: unknown, layerId: string) => ({
+            removedLayerIds: [],
+            newState: (layerIds as string[]).map((id: string) =>
               id === layerId ? `${datasourceId}_clear_${layerId}` : id
             ),
-          removeLayer: (layerIds: unknown, layerId: string) =>
-            (layerIds as string[]).filter((id: string) => id !== layerId),
+          }),
+          removeLayer: (layerIds: unknown, layerId: string) => ({
+            newState: (layerIds as string[]).filter((id: string) => id !== layerId),
+            removedLayerIds: [layerId],
+          }),
           insertLayer: (layerIds: unknown, layerId: string, layersToLinkTo: string[]) => [
             ...(layerIds as string[]),
             layerId,
@@ -317,8 +322,9 @@ describe('lensSlice', () => {
             (layerIds as string[]).map((id: string) =>
               id === layerId ? `vis_clear_${layerId}` : id
             ),
-          removeLayer: (layerIds: unknown, layerId: string) =>
-            (layerIds as string[]).filter((id: string) => id !== layerId),
+          removeLayer: jest.fn((layerIds: unknown, layerId: string) =>
+            (layerIds as string[]).filter((id: string) => id !== layerId)
+          ),
           getLayerIds: (layerIds: unknown) => layerIds as string[],
           getLayersToLinkTo: (state, newLayerId) => ['linked-layer-id'],
           appendLayer: (layerIds: unknown, layerId: string) => [...(layerIds as string[]), layerId],
@@ -482,6 +488,54 @@ describe('lensSlice', () => {
         expect(state.datasourceStates.testDatasource2.state).toEqual(['layer2']);
         expect(state.stagedPreview).not.toBeDefined();
       });
+
+      it('removeLayer: should remove all layers from visualization that were removed by datasource', () => {
+        const removedLayerId = 'other-removed-layer';
+
+        const testDatasource3 = testDatasource('testDatasource3');
+        testDatasource3.removeLayer = (layerIds: unknown, layerId: string) => ({
+          newState: (layerIds as string[]).filter((id: string) => id !== layerId),
+          removedLayerIds: [layerId, removedLayerId],
+        });
+
+        const localStore = makeLensStore({
+          preloadedState: {
+            activeDatasourceId: 'testDatasource',
+            datasourceStates: {
+              ...datasourceStates,
+              testDatasource3: {
+                isLoading: false,
+                state: [],
+              },
+            },
+            visualization: {
+              activeId: activeVisId,
+              state: ['layer1', 'layer2'],
+            },
+            stagedPreview: {
+              visualization: {
+                activeId: activeVisId,
+                state: ['layer1', 'layer2'],
+              },
+              datasourceStates,
+            },
+          },
+          storeDeps: mockStoreDeps({
+            visualizationMap: visualizationMap as unknown as VisualizationMap,
+            datasourceMap: { ...datasourceMap, testDatasource3 } as unknown as DatasourceMap,
+          }),
+        }).store;
+
+        localStore.dispatch(
+          removeOrClearLayer({
+            visualizationId: 'testVis',
+            layerId: 'layer1',
+            layerIds: ['layer1', 'layer2'],
+          })
+        );
+
+        expect(visualizationMap[activeVisId].removeLayer).toHaveBeenCalledTimes(2);
+      });
     });
 
     describe('removing a dimension', () => {
@@ -546,8 +600,6 @@ describe('lensSlice', () => {
             datasourceMap: datasourceMap as unknown as DatasourceMap,
           }),
         }).store;
-
-        jest.clearAllMocks();
       });
 
       it('removes a dimension', () => {

diff --git a/x-pack/plugins/lens/public/state_management/lens_slice.ts b/x-pack/plugins/lens/public/state_management/lens_slice.ts
@@ -7,7 +7,7 @@
 
 import { createAction, createReducer, current, PayloadAction } from '@reduxjs/toolkit';
 import { VisualizeFieldContext } from '@kbn/ui-actions-plugin/public';
-import { mapValues } from 'lodash';
+import { mapValues, uniq } from 'lodash';
 import { Query } from '@kbn/es-query';
 import { History } from 'history';
 import { LensEmbeddableInput } from '..';
@@ -400,16 +400,23 @@ export const makeLensReducer = (storeDeps: LensStoreDeps) => {
           layerIds.length
         ) === 'clear';
 
+      let removedLayerIds: string[] = [];
+
       state.datasourceStates = mapValues(
         state.datasourceStates,
         (datasourceState, datasourceId) => {
           const datasource = datasourceMap[datasourceId!];
+
+          const { newState, removedLayerIds: removedLayerIdsForThisDatasource } = isOnlyLayer
+            ? datasource.clearLayer(datasourceState.state, layerId)
+            : datasource.removeLayer(datasourceState.state, layerId);
+
+          removedLayerIds = [...removedLayerIds, ...removedLayerIdsForThisDatasource];
+
           return {
             ...datasourceState,
             ...(datasourceId === state.activeDatasourceId && {
-              state: isOnlyLayer
-                ? datasource.clearLayer(datasourceState.state, layerId)
-                : datasource.removeLayer(datasourceState.state, layerId),
+              state: newState,
             }),
           };
         }
@@ -419,10 +426,22 @@ export const makeLensReducer = (storeDeps: LensStoreDeps) => {
       const currentDataViewsId = activeDataSource.getUsedDataView(
         state.datasourceStates[state.activeDatasourceId!].state
       );
-      state.visualization.state =
-        isOnlyLayer || !activeVisualization.removeLayer
-          ? activeVisualization.clearLayer(state.visualization.state, layerId, currentDataViewsId)
-          : activeVisualization.removeLayer(state.visualization.state, layerId);
+
+      if (isOnlyLayer || !activeVisualization.removeLayer) {
+        state.visualization.state = activeVisualization.clearLayer(
+          state.visualization.state,
+          layerId,
+          currentDataViewsId
+        );
+      }
+
+      uniq(removedLayerIds).forEach(
+        (removedId) =>
+          (state.visualization.state = activeVisualization.removeLayer?.(
+            state.visualization.state,
+            removedId
+          ))
+      );
     },
     [changeIndexPattern.type]: (
       state,

diff --git a/x-pack/plugins/lens/public/types.ts b/x-pack/plugins/lens/public/types.ts
@@ -265,8 +265,8 @@ export interface Datasource<T = unknown, P = unknown> {
 
   insertLayer: (state: T, newLayerId: string, linkToLayers?: string[]) => T;
   createEmptyLayer: (indexPatternId: string) => T;
-  removeLayer: (state: T, layerId: string) => T;
-  clearLayer: (state: T, layerId: string) => T;
+  removeLayer: (state: T, layerId: string) => { newState: T; removedLayerIds: string[] };
+  clearLayer: (state: T, layerId: string) => { newState: T; removedLayerIds: string[] };
   cloneLayer: (
     state: T,
     layerId: string,

diff --git a/x-pack/plugins/lens/public/visualizations/metric/visualization.tsx b/x-pack/plugins/lens/public/visualizations/metric/visualization.tsx
@@ -454,6 +454,10 @@ export const getMetricVisualization = ({
     return newState;
   },
 
+  getRemoveOperation(_state, _layerId) {
+    return 'clear';
+  },
+
   getLayersToLinkTo(state, newLayerId: string): string[] {
     return newLayerId === state.trendlineLayerId ? [state.layerId] : [];
   },