feat: allow sorting and persist column state in infinite table

package.json

@@ -36,6 +36,7 @@
   "dependencies": {
     "ag-grid-community": "^31.2.1",
     "ag-grid-react": "^31.2.1",
+    "dequal": "^2.0.3",
     "lodash.debounce": "^4.0.8",
     "react": "18.2.0",
     "react-dom": "18.2.0",

src/components/InfiniteTable/InfiniteTable.tsx

@@ -13,36 +13,44 @@ import "@/styles/ag-theme-quartz.css";
 import {
   BodyScrollEvent,
   ColDef,
-  ColumnMovedEvent,
   ColumnResizedEvent,
   ColumnState,
   FirstDataRenderedEvent,
   GridReadyEvent,
   IGetRowsParams,
   RowDoubleClickedEvent,
+  SortDirection,
 } from "ag-grid-community";
 import { TableProps } from "@/types";
 import { useDeepArrayMemo } from "@/hooks/useDeepArrayMemo";
-import debounce from "lodash/debounce";
 import { HeaderCheckbox } from "./HeaderCheckbox";
 import { useRowSelection } from "./useRowSelection";
-import { useAutoFitColumns } from "./useAutoFitColumns";
-import { getPersistedColumnState } from "./columnStateHelper";
+import { areStatesEqual, useColumnState } from "./useColumnState";
+import { CHECKBOX_COLUMN, STATUS_COLUMN } from "./columnStateHelper";
 
 const DEBOUNCE_TIME = 50;
+const DEFAULT_TOTAL_ROWS_VALUE = Number.MAX_SAFE_INTEGER;
 
 export type InfiniteTableProps = Omit<
   TableProps,
   "dataSource" & "loading" & "loadingComponent" & "height"
 > & {
-  onRequestData: (startRow: number, endRow: number) => Promise<any[]>;
+  onRequestData: ({
+    startRow,
+    endRow,
+    sortFields,
+  }: {
+    startRow: number;
+    endRow: number;
+    sortFields?: Record<string, SortDirection>;
+  }) => Promise<any[] | undefined>;
   height?: number;
   onColumnChanged?: (columnsState: ColumnState[]) => void;
   onGetColumnsState?: () => ColumnState[] | undefined;
   onGetFirstVisibleRowIndex?: () => number | undefined;
   onChangeFirstVisibleRowIndex?: (index: number) => void;
   onGetSelectedRowKeys?: () => any[] | undefined;
-  totalRows: number;
+  totalRows?: number;
   allRowSelectedMode?: boolean;
   onAllRowSelectedModeChange?: (allRowSelectedMode: boolean) => void;
   footer?: React.ReactNode;
@@ -71,7 +79,7 @@ const InfiniteTableComp = forwardRef<InfiniteTableRef, InfiniteTableProps>(
       onChangeFirstVisibleRowIndex,
       onGetFirstVisibleRowIndex,
       onGetSelectedRowKeys,
-      totalRows,
+      totalRows = DEFAULT_TOTAL_ROWS_VALUE,
       onAllRowSelectedModeChange,
       allRowSelectedMode: allRowSelectedModeProps,
       footer,
@@ -84,39 +92,10 @@ const InfiniteTableComp = forwardRef<InfiniteTableRef, InfiniteTableProps>(
     const gridRef = useRef<AgGridReact>(null);
     const firstTimeOnBodyScroll = useRef(true);
     const allRowSelectedModeRef = useRef<boolean>(false);
-    const columnsPersistedStateRef = useRef<any>();
     const containerRef = useRef<HTMLDivElement>(null);
-    const columnChangeListenerReady = useRef(false);
     const totalHeight = footer ? heightProps + footerHeight : heightProps;
     const tableHeight = footer ? heightProps - footerHeight : heightProps;
 
-    const { autoSizeColumnsIfNecessary } = useAutoFitColumns({
-      gridRef,
-      containerRef,
-      columnsPersistedStateRef,
-      hasStatusColumn,
-    });
-
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-    const debouncedOnColumnChanged = useCallback(
-      debounce((state) => {
-        if (!columnChangeListenerReady.current) {
-          columnChangeListenerReady.current = true;
-          return;
-        }
-        onColumnsChangedProps?.(state);
-      }, DEBOUNCE_TIME),
-      [onColumnsChangedProps],
-    );
-
-    const onColumnChanged = useCallback(
-      (event: ColumnResizedEvent | ColumnMovedEvent) => {
-        const state = event.api.getColumnState();
-        debouncedOnColumnChanged(state);
-      },
-      [debouncedOnColumnChanged],
-    );
-
     useImperativeHandle(ref, () => ({
       unselectAll: () => {
         setSelectedRowKeysPendingToRender([]);
@@ -148,17 +127,79 @@ const InfiniteTableComp = forwardRef<InfiniteTableRef, InfiniteTableProps>(
 
     const columns = useDeepArrayMemo(columnsProps, "key");
 
-    const defaultColDef = useMemo<ColDef>(() => ({}), []);
+    const {
+      loadPersistedColumnState,
+      columnsPersistedStateRef,
+      applyAndUpdateNewState,
+    } = useColumnState({
+      gridRef,
+      containerRef,
+      hasStatusColumn,
+      columns,
+      onGetColumnsState,
+    });
+
+    const onColumnChanged = useCallback(() => {
+      const state = gridRef?.current?.api.getColumnState();
+      if (!state) {
+        return;
+      }
+      if (areStatesEqual(state, columnsPersistedStateRef.current)) {
+        return;
+      }
+      applyAndUpdateNewState(state);
+      onColumnsChangedProps?.(state);
+    }, [
+      applyAndUpdateNewState,
+      columnsPersistedStateRef,
+      onColumnsChangedProps,
+    ]);
+
+    const onColumnMoved = useCallback(() => {
+      onColumnChanged();
+    }, [onColumnChanged]);
+
+    const onColumnResized = useCallback(
+      (event: ColumnResizedEvent) => {
+        if (!event.finished) {
+          return;
+        }
+        onColumnChanged();
+      },
+      [onColumnChanged],
+    );
+
+    const getSortedFields = useCallback(():
+      | Record<string, SortDirection>
+      | undefined => {
+      const state = gridRef?.current?.api.getColumnState()!;
+
+      const columnsWithSort = state.filter((col) => col.sort);
+      if (columnsWithSort.length === 0) {
+        return undefined;
+      }
+      const sortFields = columnsWithSort.reduce(
+        (acc, col) => ({
+          ...acc,
+          [col.colId]: col.sort,
+        }),
+        {},
+      );
+
+      return sortFields;
+    }, []);
 
     const colDefs = useMemo((): ColDef[] => {
       const checkboxColumn = {
         checkboxSelection: true,
         suppressMovable: true,
         sortable: false,
-        lockPosition: true,
         pinned: "left",
+        lockPosition: "left",
+        lockPinned: true,
         maxWidth: 50,
         resizable: false,
+        field: CHECKBOX_COLUMN,
         headerComponent: () => (
           <HeaderCheckbox
             totalRows={totalRows}
@@ -172,36 +213,51 @@ const InfiniteTableComp = forwardRef<InfiniteTableRef, InfiniteTableProps>(
         ),
       } as ColDef;
 
-      const restOfColumns = columns.map((column) => ({
+      const storedState = columnsPersistedStateRef.current;
+      const storedStateKeys = storedState?.map((col: any) => col.colId);
+
+      const restOfColumns: ColDef[] = columns.map((column) => ({
         field: column.key,
-        sortable: false,
+        sortable: column.isSortable,
         headerName: column.title,
         cellRenderer: column.render
           ? (cell: any) => column.render(cell.value)
           : undefined,
       }));
 
+      // restOfColumns should be sorted by the order of the storedState
+      storedState &&
+        storedStateKeys &&
+        restOfColumns.sort((a, b) => {
+          const aIndex = storedStateKeys.indexOf(a.field);
+          const bIndex = storedStateKeys.indexOf(b.field);
+          return aIndex - bIndex;
+        });
+
       const statusColumn = {
-        field: "$status",
+        field: STATUS_COLUMN,
         suppressMovable: true,
         sortable: false,
-        lockPosition: true,
+        lockPosition: "left",
+        lockPinned: true,
         maxWidth: 30,
-        type: "leftAligned",
         pinned: "left",
         resizable: false,
         headerComponent: () => null,
         cellRenderer: (cell: any) => statusComponent?.(cell.value),
       } as ColDef;
 
-      return [
+      const finalColumns = [
         checkboxColumn,
         ...(hasStatusColumn ? [statusColumn] : []),
         ...restOfColumns,
       ];
+
+      return finalColumns;
     }, [
       allRowSelectedMode,
       columns,
+      columnsPersistedStateRef,
       hasStatusColumn,
       internalSelectedRowKeys.length,
       onHeaderCheckboxChange,
@@ -213,7 +269,15 @@ const InfiniteTableComp = forwardRef<InfiniteTableRef, InfiniteTableProps>(
       async (params: IGetRowsParams) => {
         gridRef.current?.api.showLoadingOverlay();
         const { startRow, endRow } = params;
-        const data = await onRequestData(startRow, endRow);
+        const data = await onRequestData({
+          startRow,
+          endRow,
+          sortFields: getSortedFields(),
+        });
+        if (!data) {
+          params.failCallback();
+          return;
+        }
         let lastRow = -1;
         if (data.length < endRow - startRow) {
           lastRow = startRow + data.length;
@@ -257,10 +321,9 @@ const InfiniteTableComp = forwardRef<InfiniteTableRef, InfiniteTableProps>(
           }
         }
         gridRef.current?.api.hideOverlay();
-        autoSizeColumnsIfNecessary();
       },
       [
-        autoSizeColumnsIfNecessary,
+        getSortedFields,
         hasStatusColumn,
         onGetSelectedRowKeys,
         onRequestData,
@@ -272,22 +335,12 @@ const InfiniteTableComp = forwardRef<InfiniteTableRef, InfiniteTableProps>(
 
     const onGridReady = useCallback(
       (params: GridReadyEvent) => {
-        columnsPersistedStateRef.current = getPersistedColumnState({
-          actualColumnKeys: columns.map((column) => column.key),
-          persistedColumnState: onGetColumnsState?.(),
-        });
-        if (columnsPersistedStateRef.current) {
-          params.api.applyColumnState({
-            state: columnsPersistedStateRef.current,
-            applyOrder: true,
-          });
-        }
-
+        loadPersistedColumnState();
         params.api.setGridOption("datasource", {
           getRows,
         });
       },
-      [columns, getRows, onGetColumnsState],
+      [getRows, loadPersistedColumnState],
     );
 
     const onRowDoubleClicked = useCallback(
@@ -335,7 +388,6 @@ const InfiniteTableComp = forwardRef<InfiniteTableRef, InfiniteTableProps>(
           <AgGridReact
             ref={gridRef}
             columnDefs={colDefs}
-            defaultColDef={defaultColDef}
             onRowDoubleClicked={onRowDoubleClicked}
             rowStyle={{
               cursor: onRowDoubleClick ? "pointer" : "auto",
@@ -345,10 +397,10 @@ const InfiniteTableComp = forwardRef<InfiniteTableRef, InfiniteTableProps>(
             suppressRowClickSelection={true}
             rowBuffer={0}
             rowSelection={"multiple"}
-            onColumnMoved={onColumnChanged}
-            onColumnResized={onColumnChanged}
+            onDragStopped={onColumnMoved}
+            onColumnResized={onColumnResized}
             rowModelType={"infinite"}
-            cacheBlockSize={20}
+            cacheBlockSize={200}
             onSelectionChanged={onSelectionChangedDebounced}
             cacheOverflowSize={2}
             maxConcurrentDatasourceRequests={1}
@@ -358,6 +410,7 @@ const InfiniteTableComp = forwardRef<InfiniteTableRef, InfiniteTableProps>(
             onFirstDataRendered={onFirstDataRendered}
             onBodyScroll={onBodyScroll}
             blockLoadDebounceMillis={DEBOUNCE_TIME}
+            suppressDragLeaveHidesColumns={true}
           />
         </div>
         {footer && <div style={{ height: footerHeight }}>{footer}</div>}

src/components/InfiniteTable/columnStateHelper.ts

@@ -1,3 +1,11 @@
+export const STATUS_COLUMN = "$status";
+export const CHECKBOX_COLUMN = "$checkbox";
+export const FIXED_COLUMNS_TO_IGNORE = [CHECKBOX_COLUMN];
+export const ALL_COLUMNS_TO_IGNORE = [
+  ...FIXED_COLUMNS_TO_IGNORE,
+  STATUS_COLUMN,
+];
+
 export const getPersistedColumnState = ({
   actualColumnKeys,
   persistedColumnState,
@@ -15,7 +23,9 @@ export const getPersistedColumnState = ({
   // we have to remove the "0" (checkbox column) from the persistedColumnKeys
   const sortedActualColumnKeys = [...actualColumnKeys].sort();
   const sortedPersistedColumnKeys = [
-    ...persistedColumnKeys.filter((key) => key !== "0"),
+    ...persistedColumnKeys.filter(
+      (key) => !ALL_COLUMNS_TO_IGNORE.includes(key),
+    ),
   ].sort();
   const areColumnKeysEqual =
     JSON.stringify(sortedActualColumnKeys) ===