[Lens] Retain column config (#90048)

flash1293 · flash1293 · commit 6136d59db2ea · 2021-02-08T13:49:22.000+01:00
diff --git a/x-pack/plugins/lens/public/datatable_visualization/visualization.test.tsx b/x-pack/plugins/lens/public/datatable_visualization/visualization.test.tsx
@@ -108,6 +108,40 @@ describe('Datatable Visualization', () => {
       expect(suggestions.length).toBeGreaterThan(0);
     });
 
+    it('should retain width and hidden config from existing state', () => {
+      const suggestions = datatableVisualization.getSuggestions({
+        state: {
+          layerId: 'first',
+          columns: [
+            { columnId: 'col1', width: 123 },
+            { columnId: 'col2', hidden: true },
+          ],
+          sorting: {
+            columnId: 'col1',
+            direction: 'asc',
+          },
+        },
+        table: {
+          isMultiRow: true,
+          layerId: 'first',
+          changeType: 'initial',
+          columns: [numCol('col1'), strCol('col2'), strCol('col3')],
+        },
+        keptLayerIds: [],
+      });
+
+      expect(suggestions.length).toBeGreaterThan(0);
+      expect(suggestions[0].state.columns).toEqual([
+        { columnId: 'col1', width: 123 },
+        { columnId: 'col2', hidden: true },
+        { columnId: 'col3' },
+      ]);
+      expect(suggestions[0].state.sorting).toEqual({
+        columnId: 'col1',
+        direction: 'asc',
+      });
+    });
+
     it('should not make suggestions when the table is unchanged', () => {
       const suggestions = datatableVisualization.getSuggestions({
         state: {
diff --git a/x-pack/plugins/lens/public/datatable_visualization/visualization.tsx b/x-pack/plugins/lens/public/datatable_visualization/visualization.tsx
@@ -98,6 +98,12 @@ export const datatableVisualization: Visualization<DatatableVisualizationState>
     ) {
       return [];
     }
+    const oldColumnSettings: Record<string, ColumnState> = {};
+    if (state) {
+      state.columns.forEach((column) => {
+        oldColumnSettings[column.columnId] = column;
+      });
+    }
     const title =
       table.changeType === 'unchanged'
         ? i18n.translate('xpack.lens.datatable.suggestionLabel', {
@@ -126,8 +132,12 @@ export const datatableVisualization: Visualization<DatatableVisualizationState>
         // table with >= 10 columns will have a score of 0.4, fewer columns reduce score
         score: (Math.min(table.columns.length, 10) / 10) * 0.4,
         state: {
+          ...(state || {}),
           layerId: table.layerId,
-          columns: table.columns.map((col) => ({ columnId: col.columnId })),
+          columns: table.columns.map((col) => ({
+            ...(oldColumnSettings[col.columnId] || {}),
+            columnId: col.columnId,
+          })),
         },
         previewIcon: LensIconChartDatatable,
         // tables are hidden from suggestion bar, but used for drag & drop and chart switching
