From e5b0088698bd5d16e64c6cc3232acab44ec6437e Mon Sep 17 00:00:00 2001
From: Anan Zhuang <ananzh@amazon.com>
Date: Wed, 28 Sep 2022 04:15:49 +0000
Subject: [PATCH] [Table Visualization] make table vis column resizable

* add resizable state to column

Partially resolve:
https://github.com/opensearch-project/OpenSearch-Dashboards/issues/2305

Signed-off-by: Anan Zhuang <ananzh@amazon.com>
---
 .../public/components/table_vis_component.tsx | 57 ++++++++++++++++++-
 .../components/table_vis_grid_columns.tsx     | 10 +++-
 src/plugins/vis_type_table/public/types.ts    |  5 ++
 3 files changed, 67 insertions(+), 5 deletions(-)

diff --git a/src/plugins/vis_type_table/public/components/table_vis_component.tsx b/src/plugins/vis_type_table/public/components/table_vis_component.tsx
index a9b8c6a8ba36..6819caae308c 100644
--- a/src/plugins/vis_type_table/public/components/table_vis_component.tsx
+++ b/src/plugins/vis_type_table/public/components/table_vis_component.tsx
@@ -3,12 +3,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { useState, useCallback, useMemo } from 'react';
+import React, { useState, useCallback, useMemo, useRef, useEffect } from 'react';
+import { isEqual } from 'lodash';
 import { EuiDataGridProps, EuiDataGrid } from '@elastic/eui';
 
 import { IInterpreterRenderHandlers } from 'src/plugins/expressions';
 import { Table } from '../table_vis_response_handler';
-import { TableVisConfig } from '../types';
+import { TableVisConfig, ColumnWidth } from '../types';
 import { getDataGridColumns } from './table_vis_grid_columns';
 
 interface TableVisComponentProps {
@@ -45,7 +46,56 @@ export const TableVisComponent = ({ table, visConfig, handlers }: TableVisCompon
     }) as EuiDataGridProps['renderCellValue'];
   }, [rows, pagination.pageIndex, pagination.pageSize]);
 
-  const dataGridColumns = getDataGridColumns(table, visConfig, handlers);
+  // resize column
+  const [columnsWidth, setColumnsWidth] = useState<ColumnWidth[]>(
+    handlers.uiState.get('vis.columnsWidth') || []
+  );
+  const curColumnsWidth = useRef<{
+    columnsWidth: ColumnWidth[];
+  }>({
+    columnsWidth: handlers.uiState?.get('vis.columnsWidth'),
+  });
+
+  const onColumnResize: EuiDataGridProps['onColumnResize'] = useCallback(
+    ({ columnId, width }) => {
+      setColumnsWidth((prevState) => {
+        const nextColIndex = columns.findIndex((c) => c.id === columnId);
+        const prevColIndex = prevState.findIndex((c) => c.colIndex === nextColIndex);
+        const nextState = [...prevState];
+        const updatedColWidth = { colIndex: nextColIndex, width };
+
+        // if updated column index is not found, then add it to nextState
+        // else reset it in nextState
+        if (prevColIndex < 0) nextState.push(updatedColWidth);
+        else nextState[prevColIndex] = updatedColWidth;
+
+        // update uiState
+        handlers.uiState?.set('vis.columnsWidth', nextState);
+        return nextState;
+      });
+    },
+    [columns, setColumnsWidth, handlers.uiState]
+  );
+
+  useEffect(() => {
+    const updateTable = () => {
+      const updatedVisState = handlers.uiState?.getChanges()?.vis;
+      if (!isEqual(updatedVisState?.columnsWidth, curColumnsWidth.current.columnsWidth)) {
+        curColumnsWidth.current.columnsWidth = updatedVisState?.columnsWidth;
+        setColumnsWidth(updatedVisState?.columnsWidth || []);
+      }
+    };
+
+    if (handlers.uiState) {
+      handlers.uiState.on('change', updateTable);
+    }
+
+    return () => {
+      handlers.uiState?.off('change', updateTable);
+    };
+  }, [handlers.uiState]);
+
+  const dataGridColumns = getDataGridColumns(table, visConfig, handlers, columnsWidth);
   return (
     <EuiDataGrid
       aria-label="tableVis"
@@ -61,6 +111,7 @@ export const TableVisComponent = ({ table, visConfig, handlers }: TableVisCompon
         onChangeItemsPerPage,
         onChangePage,
       }}
+      onColumnResize={onColumnResize}
     />
   );
 };
diff --git a/src/plugins/vis_type_table/public/components/table_vis_grid_columns.tsx b/src/plugins/vis_type_table/public/components/table_vis_grid_columns.tsx
index 5541e39d018c..b1a9443fe059 100644
--- a/src/plugins/vis_type_table/public/components/table_vis_grid_columns.tsx
+++ b/src/plugins/vis_type_table/public/components/table_vis_grid_columns.tsx
@@ -8,13 +8,14 @@ import { i18n } from '@osd/i18n';
 import { EuiDataGridColumn, EuiDataGridColumnCellActionProps } from '@elastic/eui';
 import { IInterpreterRenderHandlers } from 'src/plugins/expressions';
 import { Table } from '../table_vis_response_handler';
-import { TableVisConfig } from '../types';
+import { TableVisConfig, ColumnWidth } from '../types';
 import { convertToFormattedData } from '../utils';
 
 export const getDataGridColumns = (
   table: Table,
   visConfig: TableVisConfig,
-  handlers: IInterpreterRenderHandlers
+  handlers: IInterpreterRenderHandlers,
+  columnsWidth: ColumnWidth[]
 ) => {
   const { formattedRows, formattedColumns } = convertToFormattedData(table, visConfig);
 
@@ -117,6 +118,8 @@ export const getDataGridColumns = (
         ]
       : undefined;
 
+    const initialWidth = columnsWidth.find((c) => c.colIndex === colIndex);
+
     const dataGridColumn: EuiDataGridColumn = {
       id: col.id,
       display: col.title,
@@ -138,6 +141,9 @@ export const getDataGridColumns = (
       },
       cellActions,
     };
+    if (initialWidth) {
+      dataGridColumn.initialWidth = initialWidth.width;
+    }
     return dataGridColumn;
   });
 };
diff --git a/src/plugins/vis_type_table/public/types.ts b/src/plugins/vis_type_table/public/types.ts
index 170683719461..b3a377065140 100644
--- a/src/plugins/vis_type_table/public/types.ts
+++ b/src/plugins/vis_type_table/public/types.ts
@@ -64,3 +64,8 @@ export interface FormattedColumn {
   sumTotal?: number;
   total?: number;
 }
+
+export interface ColumnWidth {
+  colIndex: number;
+  width: number;
+}