From 3f48aa721d7977232f322bb846b3fa0454780ae4 Mon Sep 17 00:00:00 2001
From: Jan Potoms <2109932+Janpot@users.noreply.github.com>
Date: Mon, 3 Jun 2024 10:26:19 +0200
Subject: [PATCH] [DataGrid] Data provider create and update fixes (#3621)

---
 .../toolpad-core/src/DataGrid/DataGrid.tsx    | 65 ++++++++++---------
 1 file changed, 35 insertions(+), 30 deletions(-)

diff --git a/packages/toolpad-core/src/DataGrid/DataGrid.tsx b/packages/toolpad-core/src/DataGrid/DataGrid.tsx
index d908351f663..d4a44d70a96 100644
--- a/packages/toolpad-core/src/DataGrid/DataGrid.tsx
+++ b/packages/toolpad-core/src/DataGrid/DataGrid.tsx
@@ -21,6 +21,7 @@ import {
   GridPaginationModel,
   gridClasses,
   GridRowIdGetter,
+  GridApi,
 } from '@mui/x-data-grid';
 import PropTypes from 'prop-types';
 import * as React from 'react';
@@ -231,7 +232,7 @@ interface GridState {
 }
 
 type GridAction =
-  | { kind: 'START_ROW_EDIT'; rowId: GridRowId }
+  | { kind: 'START_ROW_EDIT'; rowId: GridRowId; fieldToFocus: string | undefined }
   | { kind: 'CANCEL_ROW_EDIT' }
   | { kind: 'START_ROW_UPDATE' }
   | { kind: 'END_ROW_UPDATE' };
@@ -248,6 +249,7 @@ function gridEditingReducer(state: GridState, action: GridAction): GridState {
         rowModesModel: {
           [action.rowId]: {
             mode: GridRowModes.Edit,
+            fieldToFocus: action.fieldToFocus,
           },
         },
       };
@@ -325,6 +327,7 @@ export function getColumnsFromDataProviderFields<R extends Datum>(
 }
 
 function updateColumnsWithDataProviderEditing<R extends Datum>(
+  apiRef: React.MutableRefObject<GridApi>,
   dataProvider: ResolvedDataProvider<R>,
   baseColumns: readonly GridColDef<R>[],
   state: GridState,
@@ -363,7 +366,7 @@ function updateColumnsWithDataProviderEditing<R extends Datum>(
         const rowId = params.row[idField] as GridRowId;
 
         const isEditing = state.editedRowId !== null || state.isProcessingRowUpdate;
-        const isEditedRow = rowId === state.editedRowId;
+        const isEditedRow = isDraftRow(params.row) || rowId === state.editedRowId;
 
         if (isEditedRow) {
           actions.push(
@@ -395,7 +398,11 @@ function updateColumnsWithDataProviderEditing<R extends Datum>(
                 label="Edit"
                 disabled={isEditing}
                 onClick={() => {
-                  dispatch({ kind: 'START_ROW_EDIT', rowId });
+                  dispatch({
+                    kind: 'START_ROW_EDIT',
+                    rowId,
+                    fieldToFocus: getEditedRowFieldToFocus(apiRef, idField),
+                  });
                 }}
               />,
             );
@@ -436,23 +443,6 @@ function ToolbarGridToolbar() {
   );
 }
 
-function usePatchedRowModesModel(rowModesModel: GridRowModesModel): GridRowModesModel {
-  const prevRowModesModel = React.useRef(rowModesModel);
-  React.useEffect(() => {
-    prevRowModesModel.current = rowModesModel;
-  }, [rowModesModel]);
-
-  return React.useMemo(() => {
-    if (rowModesModel === prevRowModesModel.current) {
-      return rowModesModel;
-    }
-    const base = Object.fromEntries(
-      Object.keys(prevRowModesModel.current).map((rowId) => [rowId, { mode: GridRowModes.View }]),
-    );
-    return { ...base, ...rowModesModel };
-  }, [rowModesModel]);
-}
-
 function diffRows<R extends Record<PropertyKey, unknown>>(original: R, changed: R): Partial<R> {
   const keys = new Set([...Object.keys(original), ...Object.keys(changed)]);
   const diff: Partial<R> = {};
@@ -474,6 +464,14 @@ function diffRows<R extends Record<PropertyKey, unknown>>(original: R, changed:
   return diff;
 }
 
+function getEditedRowFieldToFocus(
+  apiRef: React.MutableRefObject<GridApi>,
+  idField: ValidId,
+): string | undefined {
+  const firstNonIdColumn = apiRef.current.getVisibleColumns().find((col) => col.field !== idField);
+  return firstNonIdColumn?.field;
+}
+
 /**
  *
  * Demos:
@@ -526,8 +524,12 @@ const DataGrid = React.forwardRef(function DataGrid<R extends Datum>(
   });
 
   const handleCreateRowRequest = React.useCallback(() => {
-    dispatchEditingAction({ kind: 'START_ROW_EDIT', rowId: DRAFT_ROW_ID });
-  }, []);
+    dispatchEditingAction({
+      kind: 'START_ROW_EDIT',
+      rowId: DRAFT_ROW_ID,
+      fieldToFocus: getEditedRowFieldToFocus(apiRef, idField),
+    });
+  }, [apiRef, idField]);
 
   const useGetManyParams = React.useMemo<GetManyParams<R>>(
     () => ({
@@ -662,17 +664,19 @@ const DataGrid = React.forwardRef(function DataGrid<R extends Datum>(
     [getRowIdProp, idField],
   );
 
-  // Remove when https://github.com/mui/mui-x/issues/11423 is fixed
-  const rowModesModelPatched = usePatchedRowModesModel(editingState.rowModesModel ?? {});
-
   const handleRowEditStart = React.useCallback<GridEventListener<'rowEditStart'>>(
     (params) => {
       const rowId = params.row[idField] as GridRowId;
-      if (params.reason === 'cellDoubleClick') {
-        dispatchEditingAction({ kind: 'START_ROW_EDIT', rowId });
+      const canEdit = !!dataProvider?.updateOne;
+      if (params.reason === 'cellDoubleClick' && canEdit) {
+        dispatchEditingAction({
+          kind: 'START_ROW_EDIT',
+          rowId,
+          fieldToFocus: getEditedRowFieldToFocus(apiRef, idField),
+        });
       }
     },
-    [idField],
+    [apiRef, dataProvider?.updateOne, idField],
   );
 
   // Calculate separately to avoid re-calculating columns on every render
@@ -702,6 +706,7 @@ const DataGrid = React.forwardRef(function DataGrid<R extends Datum>(
     );
 
     gridColumns = updateColumnsWithDataProviderEditing(
+      apiRef,
       dataProvider,
       gridColumns,
       editingState,
@@ -709,7 +714,7 @@ const DataGrid = React.forwardRef(function DataGrid<R extends Datum>(
     );
 
     return gridColumns;
-  }, [columnsProp, dataProvider, editingState, inferredFields]);
+  }, [apiRef, columnsProp, dataProvider, editingState, inferredFields]);
 
   const isSsr = useSsr();
 
@@ -743,7 +748,7 @@ const DataGrid = React.forwardRef(function DataGrid<R extends Datum>(
                 loading={loading}
                 processRowUpdate={processRowUpdate}
                 slots={slots}
-                rowModesModel={rowModesModelPatched}
+                rowModesModel={editingState.rowModesModel}
                 onRowEditStart={handleRowEditStart}
                 getRowId={getRowId}
                 {...(restProps.paginationMode === 'server'