From 883d1f82e249b977f72f083119c6ac014bed183e Mon Sep 17 00:00:00 2001
From: Bilal Shafi <bilalshafidev@gmail.com>
Date: Wed, 7 Aug 2024 14:36:42 +0500
Subject: [PATCH] [DataGrid] Fix `checkboxSelectionVisibleOnly` behavior with
 server-side pagination (#14083)

---
 .../tests/rowSelection.DataGridPro.test.tsx   | 43 +++++++++++++++++++
 .../rowSelection/useGridRowSelection.ts       | 12 +++---
 2 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/packages/x-data-grid-pro/src/tests/rowSelection.DataGridPro.test.tsx b/packages/x-data-grid-pro/src/tests/rowSelection.DataGridPro.test.tsx
index 3c5eda242358..8d0af590260d 100644
--- a/packages/x-data-grid-pro/src/tests/rowSelection.DataGridPro.test.tsx
+++ b/packages/x-data-grid-pro/src/tests/rowSelection.DataGridPro.test.tsx
@@ -210,6 +210,49 @@ describe('<DataGridPro /> - Row selection', () => {
       });
       expect(selectAllCheckbox).to.have.attr('data-indeterminate', 'false');
     });
+
+    it('should allow to select all the current page rows when props.paginationMode="server"', () => {
+      function TestDataGridSelectionServerSide({
+        rowLength = 4,
+      }: Omit<DataGridProProps, 'rows' | 'columns' | 'apiRef'> &
+        Partial<Pick<DataGridProProps, 'rows' | 'columns'>> & { rowLength?: number }) {
+        apiRef = useGridApiRef();
+        const paginationModel = { pageSize: 2, page: 1 };
+
+        const data = React.useMemo(() => getBasicGridData(rowLength, 2), [rowLength]);
+
+        const rows = data.rows.slice(
+          paginationModel.pageSize * paginationModel.page,
+          paginationModel.pageSize * (paginationModel.page + 1),
+        );
+
+        return (
+          <div style={{ width: 300, height: 300 }}>
+            <DataGridPro
+              {...data}
+              rows={rows}
+              checkboxSelection
+              checkboxSelectionVisibleOnly
+              initialState={{ pagination: { paginationModel } }}
+              pagination
+              paginationMode="server"
+              pageSizeOptions={[2]}
+              apiRef={apiRef}
+              rowCount={rowLength}
+              disableVirtualization
+            />
+          </div>
+        );
+      }
+      render(<TestDataGridSelectionServerSide />);
+
+      const selectAllCheckbox = screen.getByRole('checkbox', {
+        name: /select all rows/i,
+      });
+
+      fireEvent.click(selectAllCheckbox);
+      expect(apiRef.current.getSelectedRows()).to.have.length(2);
+    });
   });
 
   describe('apiRef: getSelectedRows', () => {
diff --git a/packages/x-data-grid/src/hooks/features/rowSelection/useGridRowSelection.ts b/packages/x-data-grid/src/hooks/features/rowSelection/useGridRowSelection.ts
index 2888ec5929a1..fe7495699b33 100644
--- a/packages/x-data-grid/src/hooks/features/rowSelection/useGridRowSelection.ts
+++ b/packages/x-data-grid/src/hooks/features/rowSelection/useGridRowSelection.ts
@@ -443,17 +443,15 @@ export const useGridRowSelection = (
     GridEventListener<'headerSelectionCheckboxChange'>
   >(
     (params) => {
-      const shouldLimitSelectionToCurrentPage =
-        props.checkboxSelectionVisibleOnly && props.pagination;
-
-      const rowsToBeSelected = shouldLimitSelectionToCurrentPage
-        ? gridPaginatedVisibleSortedGridRowIdsSelector(apiRef)
-        : gridExpandedSortedRowIdsSelector(apiRef);
+      const rowsToBeSelected =
+        props.pagination && props.checkboxSelectionVisibleOnly && props.paginationMode === 'client'
+          ? gridPaginatedVisibleSortedGridRowIdsSelector(apiRef)
+          : gridExpandedSortedRowIdsSelector(apiRef);
 
       const filterModel = gridFilterModelSelector(apiRef);
       apiRef.current.selectRows(rowsToBeSelected, params.value, filterModel?.items.length > 0);
     },
-    [apiRef, props.checkboxSelectionVisibleOnly, props.pagination],
+    [apiRef, props.checkboxSelectionVisibleOnly, props.pagination, props.paginationMode],
   );
 
   const handleCellKeyDown = React.useCallback<GridEventListener<'cellKeyDown'>>(