fix(table-core): Do not select subrows if enableSubRowSelection is set to false

packages/table-core/__tests__/RowSelection.test.ts

@@ -1,5 +1,6 @@
 import {
   ColumnDef,
+  TableState,
   createColumnHelper,
   createTable,
   getCoreRowModel,
@@ -319,4 +320,111 @@ describe('RowSelection', () => {
       expect(result).toEqual('some')
     })
   })
+
+  describe('toggleAllRowsSelected', () => {
+    test('it should not select subrows if enableSubRowSelection is set to false', () => {
+      const data = makeData(2, 1)
+      const columns = generateColumns(data)
+
+      const state: Partial<TableState> = {
+        rowSelection: {},
+      }
+
+      const table = createTable<Person>({
+        enableRowSelection: true,
+        onStateChange() {},
+        renderFallbackValue: '',
+        data,
+        columns,
+        getSubRows: row => row.subRows,
+        state,
+        onRowSelectionChange(updater) {
+          state.rowSelection =
+            typeof updater === 'function'
+              ? updater(state.rowSelection ?? {})
+              : updater
+        },
+        enableSubRowSelection: false,
+        getCoreRowModel: getCoreRowModel(),
+      })
+
+      table.toggleAllRowsSelected(true)
+      const allRows = table
+        .getCoreRowModel()
+        .flatRows.map(v => v.getIsSelected())
+      expect(allRows).toEqual([true, false, true, false])
+    })
+
+    test('it should not select subrows if enableSubRowSelection is a function and it evaluates to false', () => {
+      const data = makeData(2, 2, 1)
+      const columns = generateColumns(data)
+
+      const state: Partial<TableState> = {
+        rowSelection: {},
+      }
+
+      const table = createTable<Person>({
+        onStateChange() {},
+        renderFallbackValue: '',
+        data,
+        columns,
+        getSubRows: row => row.subRows,
+        state,
+        onRowSelectionChange(updater) {
+          state.rowSelection =
+            typeof updater === 'function'
+              ? updater(state.rowSelection ?? {})
+              : updater
+        },
+        enableRowSelection: true,
+        enableSubRowSelection: row =>
+          row.id === '0' || row.id.startsWith('0.0'),
+        getCoreRowModel: getCoreRowModel(),
+      })
+
+      table.toggleAllRowsSelected(true)
+      const selectionState = table.getState().rowSelection
+
+      expect(selectionState).toEqual({
+        '0': true,
+        '0.0': true,
+        '0.0.0': true,
+        '0.1': true,
+        '1': true,
+      })
+    })
+
+    test('it should select subrows if enableSubRowSelection is set to true', () => {
+      const data = makeData(2, 1)
+      const columns = generateColumns(data)
+
+      const state: Partial<TableState> = {
+        rowSelection: {},
+      }
+
+      const table = createTable<Person>({
+        enableRowSelection: true,
+        onStateChange() {},
+        renderFallbackValue: '',
+        data,
+        columns,
+        getSubRows: row => row.subRows,
+        state,
+        onRowSelectionChange(updater) {
+          state.rowSelection =
+            typeof updater === 'function'
+              ? updater(state.rowSelection ?? {})
+              : updater
+        },
+        enableSubRowSelection: true,
+        getCoreRowModel: getCoreRowModel(),
+      })
+
+      table.toggleAllRowsSelected(true)
+      const allRows = table
+        .getCoreRowModel()
+        .flatRows.map(v => v.getIsSelected())
+      expect(allRows).toEqual([true, true, true, true])
+    })
+  })
 })

packages/table-core/src/features/RowSelection.ts

@@ -225,21 +225,40 @@ export const RowSelection: TableFeature = {
 
         const rowSelection = { ...old }
 
-        const preGroupedFlatRows = table.getPreGroupedRowModel().flatRows
-
-        // We don't use `mutateRowIsSelected` here for performance reasons.
-        // All of the rows are flat already, so it wouldn't be worth it
-        if (value) {
-          preGroupedFlatRows.forEach(row => {
-            if (!row.getCanSelect()) {
-              return
-            }
-            rowSelection[row.id] = true
+        if (typeof table.options.enableSubRowSelection === 'function') {
+          // There's no easy workaround like using flat rows here in the case enableSubRowSelection is a function
+          // We need to recursively check if the current row/subrow can be selected so just
+          // using mutateRowIsSelected is easier
+          table.getFilteredRowModel().rows.forEach(row => {
+            mutateRowIsSelected(
+              rowSelection,
+              row.id,
+              value ?? false,
+              true,
+              table
+            )
           })
         } else {
-          preGroupedFlatRows.forEach(row => {
-            delete rowSelection[row.id]
-          })
+          // Here we check if only the first level of rows should be selected (enableSubRowSelection is true.)
+          // If that's the case we just get the first level of rows, otherwise we get the flat rows.
+          const preGroupedFlatRows = table.options.enableSubRowSelection
+            ? table.getPreGroupedRowModel().flatRows
+            : table.getPreGroupedRowModel().rows
+
+          // We don't use `mutateRowIsSelected` here for performance reasons.
+          // The rows are either flat or just the first level is being selected, so it wouldn't be worth it
+          if (value) {
+            preGroupedFlatRows.forEach(row => {
+              if (!row.getCanSelect()) {
+                return
+              }
+              rowSelection[row.id] = true
+            })
+          } else {
+            preGroupedFlatRows.forEach(row => {
+              delete rowSelection[row.id]
+            })
+          }
         }
 
         return rowSelection