[DataGrid] Always refetch lazy-loading rows

packages/x-data-grid-pro/src/hooks/features/serverSideLazyLoader/useGridDataSourceLazyLoader.ts

@@ -151,6 +151,7 @@ export const useGridDataSourceLazyLoader = (
 
   const addSkeletonRows = React.useCallback(() => {
     const tree = privateApiRef.current.state.rows.tree;
+    const dataRowIdToModelLookup = privateApiRef.current.state.rows.dataRowIdToModelLookup;
     const rootGroup = tree[GRID_ROOT_GROUP_ID] as GridGroupNode;
     const rootGroupChildren = [...rootGroup.children];
 
@@ -159,23 +160,31 @@ export const useGridDataSourceLazyLoader = (
 
     /**
      * Do nothing if
-     * - rowCount is unknown
      * - children count is 0
-     * - children count is equal to rowCount
      */
-    if (
-      pageRowCount === -1 ||
-      pageRowCount === undefined ||
-      rootChildrenCount === 0 ||
-      rootChildrenCount === pageRowCount
-    ) {
+    if (rootChildrenCount === 0) {
       return;
     }
 
-    // fill the grid with skeleton rows
-    for (let i = 0; i < pageRowCount - rootChildrenCount; i += 1) {
-      const skeletonId = getSkeletonRowId(i + rootChildrenCount); // to avoid duplicate keys on rebuild
-      rootGroupChildren.push(skeletonId);
+    const pageToSkip = adjustRowParams({
+      start: renderedRowsIntervalCache.current.firstRowToRender,
+      end: renderedRowsIntervalCache.current.lastRowToRender,
+    } as GridGetRowsParams);
+
+    let hasChanged = false;
+
+    for (let i = 0; i < rootChildrenCount; i += 1) {
+      // replace the rows not in the viewport with skeleton rows
+      if (
+        ((pageToSkip.start as number) <= i && i <= pageToSkip.end) ||
+        tree[rootGroupChildren[i]]?.type === 'skeletonRow'
+      ) {
+        continue;
+      }
+
+      const skeletonId = getSkeletonRowId(i); // to avoid duplicate keys on rebuild
+      const removedRow = rootGroupChildren[i];
+      rootGroupChildren[i] = skeletonId;
 
       const skeletonRowNode: GridSkeletonRowNode = {
         type: 'skeletonRow',
@@ -185,6 +194,32 @@ export const useGridDataSourceLazyLoader = (
       };
 
       tree[skeletonId] = skeletonRowNode;
+      delete tree[removedRow];
+      delete dataRowIdToModelLookup[removedRow];
+      hasChanged = true;
+    }
+
+    // Should only happen with VIEWPORT loading trigger
+    if (pageRowCount !== -1 && pageRowCount !== undefined) {
+      // fill the grid with skeleton rows
+      for (let i = 0; i < pageRowCount - rootChildrenCount; i += 1) {
+        const skeletonId = getSkeletonRowId(i + rootChildrenCount); // to avoid duplicate keys on rebuild
+        rootGroupChildren.push(skeletonId);
+
+        const skeletonRowNode: GridSkeletonRowNode = {
+          type: 'skeletonRow',
+          id: skeletonId,
+          parent: GRID_ROOT_GROUP_ID,
+          depth: 0,
+        };
+
+        tree[skeletonId] = skeletonRowNode;
+        hasChanged = true;
+      }
+    }
+
+    if (!hasChanged) {
+      return;
     }
 
     tree[GRID_ROOT_GROUP_ID] = { ...rootGroup, children: rootGroupChildren };
@@ -195,11 +230,12 @@ export const useGridDataSourceLazyLoader = (
         rows: {
           ...state.rows,
           tree,
+          dataRowIdToModelLookup,
         },
       }),
       'addSkeletonRows',
     );
-  }, [privateApiRef]);
+  }, [privateApiRef, adjustRowParams]);
 
   const updateLoadingTrigger = React.useCallback(
     (rowCount: number) => {
@@ -313,7 +349,7 @@ export const useGridDataSourceLazyLoader = (
     GridEventListener<'renderedRowsIntervalChange'>
   >(
     (params) => {
-      if (rowsStale.current || loadingTrigger.current !== LoadingTrigger.VIEWPORT) {
+      if (rowsStale.current) {
         return;
       }