feat: keep column visibility on query params (#19013)

* feat: keep column visibility on query params

* fix resizing logic

---------

Co-authored-by: Peer Richelsen <peeroke@gmail.com>

packages/features/data-table/components/DataTableWrapper.tsx

@@ -54,7 +54,7 @@ export function DataTableWrapper<TData, TValue>({
     fetchNextPage,
     isFetching,
   });
-  const { sorting, setSorting } = useDataTable();
+  const { sorting, setSorting, columnVisibility, setColumnVisibility } = useDataTable();
   const columnFilters = useColumnFilters();
 
   useEffect(() => {
@@ -62,12 +62,14 @@ export function DataTableWrapper<TData, TValue>({
       ...prev,
       sorting,
       columnFilters,
+      columnVisibility,
     }));
     table.setOptions((prev) => ({
       ...prev,
       onSortingChange: setSorting,
+      onColumnVisibilityChange: setColumnVisibility,
     }));
-  }, [table, sorting, columnFilters]);
+  }, [table, sorting, columnFilters, columnVisibility]);
 
   return (
     <DataTable

packages/features/data-table/lib/context.tsx

@@ -1,11 +1,11 @@
 "use client";
 
-import type { SortingState, OnChangeFn } from "@tanstack/react-table";
+import type { SortingState, OnChangeFn, VisibilityState } from "@tanstack/react-table";
 import { useQueryState, parseAsArrayOf, parseAsJson } from "nuqs";
 import { createContext, useCallback } from "react";
 import { z } from "zod";
 
-import { type FilterValue, ZFilterValue, ZSorting } from "./types";
+import { type FilterValue, ZFilterValue, ZSorting, ZColumnVisibility } from "./types";
 
 const ZActiveFilter = z.object({
   f: z.string(),
@@ -23,6 +23,9 @@ export type DataTableContextType = {
 
   sorting: SortingState;
   setSorting: OnChangeFn<SortingState>;
+
+  columnVisibility: VisibilityState;
+  setColumnVisibility: OnChangeFn<VisibilityState>;
 };
 
 export const DataTableContext = createContext<DataTableContextType | null>(null);
@@ -36,6 +39,10 @@ export function DataTableProvider({ children }: { children: React.ReactNode }) {
     "sorting",
     parseAsArrayOf(parseAsJson(ZSorting.parse)).withDefault([])
   );
+  const [columnVisibility, setColumnVisibility] = useQueryState<VisibilityState>(
+    "cols",
+    parseAsJson(ZColumnVisibility.parse).withDefault({})
+  );
 
   const clearAll = useCallback(
     (exclude?: string[]) => {
@@ -81,6 +88,8 @@ export function DataTableProvider({ children }: { children: React.ReactNode }) {
         removeFilter,
         sorting,
         setSorting,
+        columnVisibility,
+        setColumnVisibility,
       }}>
       {children}
     </DataTableContext.Provider>

packages/features/data-table/lib/resizing.ts

@@ -49,8 +49,13 @@ function getAdjustedColumnSizing<TData>({
     return currentColumnSizing;
   }
 
-  const getColumnSize = (id: string) =>
-    resizedColumns.has(id) ? currentColumnSizing[id] : initialColumnSizing[id];
+  const getColumnSize = (header: Header<TData, unknown>) => {
+    const id = header.id;
+    if (resizedColumns.has(id)) {
+      return currentColumnSizing[id] ?? initialColumnSizing[id] ?? header.getSize();
+    }
+    return initialColumnSizing[id] ?? header.getSize();
+  };
 
   const isAdjustable = (header: Header<TData, unknown>) =>
     header.column.columnDef.enableResizing !== false && !resizedColumns.has(header.id);
@@ -59,7 +64,7 @@ function getAdjustedColumnSizing<TData>({
   const numberOfAdjustableColumns = headers.filter(isAdjustable).length;
 
   // Calculate widths and required adjustments
-  const totalColumnsWidth = headers.reduce((total, header) => total + getColumnSize(header.id), 0);
+  const totalColumnsWidth = headers.reduce((total, header) => total + getColumnSize(header), 0);
   const widthToFill = Math.max(0, containerWidth - totalColumnsWidth);
   const widthPerAdjustableColumn =
     numberOfAdjustableColumns === 0 ? 0 : Math.floor(widthToFill / numberOfAdjustableColumns);
@@ -69,7 +74,7 @@ function getAdjustedColumnSizing<TData>({
 
   // Build the new column sizing object
   const newColumnSizing = headers.reduce((acc, header, index) => {
-    const baseWidth = getColumnSize(header.id);
+    const baseWidth = getColumnSize(header);
     const adjustmentWidth = isAdjustable(header) ? widthPerAdjustableColumn : 0;
     const isLastColumn = index === headers.length - 1;
     const extraWidth = isLastColumn ? leftoverWidth : 0;

packages/features/data-table/lib/types.ts

@@ -146,3 +146,5 @@ export const ZSorting = z.object({
 });
 
 export type SortingState = Array<z.infer<typeof ZSorting>>;
+
+export const ZColumnVisibility = z.record(z.string(), z.boolean());