fix(structure): remove useEffect and set table columns as initial state

packages/sanity/src/structure/panes/documentList/ColumnsControl.tsx

@@ -1,53 +1,29 @@
 /* eslint-disable @sanity/i18n/no-attribute-string-literals */
 /* eslint-disable i18next/no-literal-string */
 import {Box, Button, Card, Checkbox, Flex, Menu, MenuButton, Stack, Text} from '@sanity/ui'
-import {type Column, type Table} from '@tanstack/react-table'
-import {useEffect, useRef, useState} from 'react'
+import {type Column, type Table, type VisibilityState} from '@tanstack/react-table'
+import {useCallback} from 'react'
 import {type SanityDocument} from 'sanity'
 
-const VISIBLE_COLUMN_LIMIT = 5
+import {VISIBLE_COLUMN_LIMIT} from './useDocumentSheetColumns'
 
 type Props = {
  table: Table<SanityDocument>
+ initialState: VisibilityState
 }
 
-export function ColumnsControl({table}: Props) {
- const tableRef = useRef(table)
- const [reset, setReset] = useState(0)
-
+export function ColumnsControl({table, initialState}: Props) {
  const isVisibleLimitReached =
  table.getVisibleLeafColumns().filter((col) => col.getCanHide()).length >= VISIBLE_COLUMN_LIMIT
 
- // set the initial visible columns state
- useEffect(() => {
- const [newColumns]: [Record<string, boolean>, number] = tableRef.current
- .getAllLeafColumns()
- .reduce(
- ([accCols, countAllowedVisible], column) => {
- // this column is always visible
- if (!column.getCanHide()) {
- return [{...accCols, [column.id]: true}, countAllowedVisible]
- }
-
- // have already reached column visibility limit, hide column by default
- if (countAllowedVisible === VISIBLE_COLUMN_LIMIT) {
- return [{...accCols, [column.id]: false}, countAllowedVisible]
- }
-
- return [{...accCols, [column.id]: true}, countAllowedVisible + 1]
- },
- [{}, 0],
- )
-
- tableRef.current.setColumnVisibility(newColumns)
- }, [reset])
+ const setInitialColumns = useCallback(() => {
+ table.setColumnVisibility(initialState)
+ }, [table, initialState])
 
  const handleColumnOnChange = (column: Column<SanityDocument, unknown>) => () => {
  column.toggleVisibility()
  }
 
- const handleResetColumns = () => setReset((prev) => prev + 1)
-
  const getColumnVisibilityDisabled = (column: Column<SanityDocument, unknown>) => {
  const isColumnVisible = column.getIsVisible()
  const isSingleColumnVisible =
@@ -62,7 +38,7 @@ export function ColumnsControl({table}: Props) {
  id="columns-control"
  menu={
  <Menu padding={3} paddingBottom={1} style={{maxHeight: 300, overflow: 'scroll'}}>
- <Button size={0} text="Reset" onClick={handleResetColumns} />
+ <Button size={0} text="Reset" onClick={setInitialColumns} />
  <Stack>
  {table
  .getAllLeafColumns()

packages/sanity/src/structure/panes/documentList/DocumentSheetListPane.tsx

@@ -68,7 +68,7 @@ function DocumentSheetListPaneInner({
  schemaType,
 }: DocumentSheetListPaneProps & {schemaType: SchemaType}) {
  const {dispatch, state} = useSearchState()
- const columns = useDocumentSheetColumns(schemaType)
+ const {columns, initialColumnsVisibility} = useDocumentSheetColumns(schemaType)
  const {data} = useDocumentSheetList({
  typeName: schemaType.name,
  })
@@ -84,6 +84,7 @@ function DocumentSheetListPaneInner({
  autoResetPageIndex: false,
  initialState: {
  pagination: {pageSize: 25},
+ columnVisibility: initialColumnsVisibility,
  },
  meta: {
  updateData: () => {
@@ -138,7 +139,7 @@ function DocumentSheetListPaneInner({
  </Text>
  </Flex>
  <TableContainer>
- <ColumnsControl table={table} />
+ <ColumnsControl table={table} initialState={initialColumnsVisibility} />
  <Table>
  <thead>
  {table.getHeaderGroups().map((headerGroup) => (

packages/sanity/src/structure/panes/documentList/__tests__/ColumnsControl.test.tsx

@@ -1,4 +1,6 @@
-import {beforeEach, describe, expect, it} from '@jest/globals'
+import {afterEach} from 'node:test'
+
+import {beforeEach, describe, expect, it, jest} from '@jest/globals'
 import {studioTheme, ThemeProvider} from '@sanity/ui'
 import {useReactTable} from '@tanstack/react-table'
 import {fireEvent, render, screen} from '@testing-library/react'
@@ -7,15 +9,27 @@ import {type SanityDocument} from 'sanity'
 import {ColumnsControl} from '../ColumnsControl'
 
 const TableHarness = ({columns}) => {
+ const initialVisibilityState = {
+ 'First Column': true,
+ 'Second Column': true,
+ 'Third Column': true,
+ 'Nested First Column': true,
+ 'Nested Second Column': true,
+ 'Fifth Column': true,
+ 'Sixth Column': false,
+ }
  const table = useReactTable<SanityDocument>({
  columns,
  data: [],
  getCoreRowModel: () => {
  throw new Error('getCoreRowModel not implemented.')
  },
+ initialState: {
+ columnVisibility: initialVisibilityState,
+ },
  })
 
- return <ColumnsControl table={table} />
+ return <ColumnsControl table={table} initialState={initialVisibilityState} />
 }
 
 describe('ColumnsControl', () => {
@@ -43,6 +57,10 @@ describe('ColumnsControl', () => {
  )
  })
 
+ afterEach(() => {
+ jest.clearAllMocks()
+ })
+
  it('should set default column visibilities', () => {
  fireEvent.click(screen.getByText('Columns'))
  expect(screen.getByRole('checkbox', {name: 'First Column'})).toBeChecked()

packages/sanity/src/structure/panes/documentList/useDocumentSheetColumns.tsx

@@ -1,5 +1,10 @@
 import {Checkbox, Flex, Text} from '@sanity/ui'
-import {type AccessorKeyColumnDef, createColumnHelper} from '@tanstack/react-table'
+import {
+ type AccessorKeyColumnDef,
+ createColumnHelper,
+ type GroupColumnDef,
+ type VisibilityState,
+} from '@tanstack/react-table'
 import {useMemo} from 'react'
 import {useMemoObservable} from 'react-rx'
 import {
@@ -14,6 +19,8 @@ import {
 import {type PaneItemPreviewState} from '../../components/paneItem/types'
 import {SheetListCell} from './SheetListCell'
 
+export const VISIBLE_COLUMN_LIMIT = 5
+
 const PreviewCell = (props: {
  documentPreviewStore: DocumentPreviewStore
  schemaType: SchemaType
@@ -47,42 +54,68 @@ const PreviewCell = (props: {
 const columnHelper = createColumnHelper<SanityDocument>()
 const SUPPORTED_FIELDS = ['string', 'number', 'boolean']
 
-const getColsFromSchemaType = (schemaType: SchemaType, parentalField?: string) => {
+type Columns = (
+ | AccessorKeyColumnDef<SanityDocument, unknown>
+ | GroupColumnDef<SanityDocument, unknown>
+)[]
+
+const getColsFromSchemaType = (schemaType: SchemaType, parentalField?: string): Columns => {
  //@ts-expect-error - wip.
- return schemaType.fields.reduce(
- (cols: AccessorKeyColumnDef<SanityDocument, unknown>[], field: any) => {
- const {type, name} = field
- if (SUPPORTED_FIELDS.includes(type.name)) {
- const nextCol = columnHelper.accessor(
- parentalField ? `${parentalField}.${field.name}` : field.name,
- {
- header: field.type.title,
- enableHiding: true,
- cell: (info) => <SheetListCell {...info} type={type} />,
- },
- )
-
- return [...cols, nextCol]
- }
-
- // if first layer nested object
- if (type.name === 'object' && !parentalField) {
- return [
- ...cols,
- columnHelper.group({header: name, columns: getColsFromSchemaType(type, field.name)}),
- ]
- }
-
- return cols
- },
- [],
- )
+ return schemaType.fields.reduce((cols: Columns[], field: any) => {
+ const {type, name} = field
+ if (SUPPORTED_FIELDS.includes(type.name)) {
+ const nextCol = columnHelper.accessor(
+ parentalField ? `${parentalField}.${field.name}` : field.name,
+ {
+ header: field.type.title,
+ enableHiding: true,
+ cell: (info) => <SheetListCell {...info} type={type} />,
+ },
+ )
+
+ return [...cols, nextCol]
+ }
+
+ // if first layer nested object
+ if (type.name === 'object' && !parentalField) {
+ return [
+ ...cols,
+ columnHelper.group({header: name, columns: getColsFromSchemaType(type, field.name)}),
+ ]
+ }
+
+ return cols
+ }, [])
+}
+
+// Type guard function to check if a column is of type GroupColumnDef
+function isAccessorKeyColumnDef(
+ col: Columns[number],
+): col is AccessorKeyColumnDef<SanityDocument, unknown> {
+ return (col as AccessorKeyColumnDef<SanityDocument, unknown>).accessorKey !== undefined
+}
+function isGroupColumnDef(
+ column: AccessorKeyColumnDef<SanityDocument, unknown> | GroupColumnDef<SanityDocument, unknown>,
+): column is GroupColumnDef<SanityDocument, unknown> {
+ return (column as GroupColumnDef<SanityDocument, unknown>).columns !== undefined
+}
+
+const flatColumns = (cols: Columns): AccessorKeyColumnDef<SanityDocument, unknown>[] => {
+ return cols.flatMap((col) => {
+ if (isAccessorKeyColumnDef(col)) {
+ return col
+ }
+ if (isGroupColumnDef(col)) {
+ return col.columns ? flatColumns(col.columns) : []
+ }
+ return []
+ })
 }
 
 export function useDocumentSheetColumns(schemaType?: SchemaType) {
  const documentPreviewStore = useDocumentPreviewStore()
 
- const columns: AccessorKeyColumnDef<SanityDocument, unknown>[] = useMemo(() => {
+ const columns: Columns = useMemo(() => {
  if (!schemaType) {
  return []
  }
@@ -119,5 +152,26 @@ export function useDocumentSheetColumns(schemaType?: SchemaType) {
  ]
  }, [documentPreviewStore, schemaType])
 
- return columns
+ const [newColumns]: [VisibilityState, number] = useMemo(
+ () =>
+ flatColumns(columns).reduce(
+ ([accCols, countAllowedVisible], column) => {
+ // this column is always visible
+ if (!column.enableHiding) {
+ return [{...accCols, [column.accessorKey]: true}, countAllowedVisible]
+ }
+
+ // have already reached column visibility limit, hide column by default
+ if (countAllowedVisible === VISIBLE_COLUMN_LIMIT) {
+ return [{...accCols, [column.accessorKey]: false}, countAllowedVisible]
+ }
+
+ return [{...accCols, [column.accessorKey]: true}, countAllowedVisible + 1]
+ },
+ [{}, 0],
+ ),
+ [columns],
+ )
+
+ return {columns, initialColumnsVisibility: newColumns}
 }