Add ability to override the Datagrid header row

docs/List.md

@@ -2016,6 +2016,7 @@ The `Datagrid` component renders a list of records as a table. It is usually use
 Here are all the props accepted by the component:
 
 * [`body`](#body-element)
+* [`header`](#header-element)
 * [`rowStyle`](#row-style-function)
 * [`rowClick`](#rowclick)
 * [`expand`](#expand)
@@ -2101,6 +2102,39 @@ const PostList = props => (
 export default PostList;
 ```
 
+### Header Element
+
+By default, `<Datagrid>` renders its header using `<DatagridHeader>`, an internal react-admin component. You can pass a custom component as the `header` prop to override that default. This can be useful e.g. to add a second header row, or to create headers spanning multiple columns.
+
+For instance, here is a simple datagrid header that displays column names with no sort and no "select all" button:
+
+```jsx
+import { TableHead, TableRow, TableCell } from '@material-ui/core';
+
+const DatagridHeader = ({ children }) => (
+    <TableHead>
+        <TableRow>
+            <TableCell></TableCell> {/* empty cell to account for the select row checkbox in the body */}
+            {Children.map(children, child => (
+                <TableCell key={child.props.source}>
+                    {child.props.source}
+                </TableCell>
+            ))}
+        </TableRow>
+    </TableHead>
+);
+
+const PostList = props => (
+    <List {...props}>
+        <Datagrid header={<DatagridHeader />}>
+            {/* ... */}
+        </Datagrid>
+    </List>
+);
+```
+
+**Tip**: To handle sorting in your custom Datagrid header component, check out the [Building a custom sort control](#building-a-custom-sort-control) section.
+
 ### Row Style Function
 
 You can customize the `<Datagrid>` row style (applied to the `<tr>` element) based on the record, thanks to the `rowStyle` prop, which expects a function. React-admin calls this function for each row, passing the current record and index as arguments. The function should return a style object, which react-admin uses as a `<tr style>` prop. 

packages/ra-ui-materialui/src/list/datagrid/Datagrid.tsx

@@ -1,12 +1,13 @@
 import * as React from 'react';
 import {
-    isValidElement,
-    Children,
     cloneElement,
+    createElement,
+    isValidElement,
     useCallback,
     useRef,
     useEffect,
     FC,
+    ComponentType,
     ReactElement,
     useMemo,
 } from 'react';
@@ -20,19 +21,12 @@ import {
     RecordMap,
     SortPayload,
 } from 'ra-core';
-import {
-    Checkbox,
-    Table,
-    TableProps,
-    TableCell,
-    TableHead,
-    TableRow,
-} from '@material-ui/core';
+import { Table, TableProps } from '@material-ui/core';
 import classnames from 'classnames';
 import union from 'lodash/union';
 import difference from 'lodash/difference';
 
-import DatagridHeaderCell from './DatagridHeaderCell';
+import { DatagridHeader } from './DatagridHeader';
 import DatagridLoading from './DatagridLoading';
 import DatagridBody, { PureDatagridBody } from './DatagridBody';
 import useDatagridStyles from './useDatagridStyles';
@@ -113,7 +107,8 @@ const Datagrid: FC<DatagridProps> = React.forwardRef((props, ref) => {
     const classes = useDatagridStyles(props);
     const {
         optimized = false,
-        body = optimized ? <PureDatagridBody /> : <DatagridBody />,
+        body = optimized ? PureDatagridBody : DatagridBody,
+        header = DatagridHeader,
         children,
         classes: classesOverride,
         className,
@@ -148,42 +143,6 @@ const Datagrid: FC<DatagridProps> = React.forwardRef((props, ref) => {
         isRowExpandable,
     ]);
 
-    const updateSortCallback = useCallback(
-        event => {
-            event.stopPropagation();
-            const newField = event.currentTarget.dataset.field;
-            const newOrder =
-                currentSort.field === newField
-                    ? currentSort.order === 'ASC'
-                        ? 'DESC'
-                        : 'ASC'
-                    : event.currentTarget.dataset.order;
-
-            setSort(newField, newOrder);
-        },
-        [currentSort.field, currentSort.order, setSort]
-    );
-
-    const updateSort = setSort ? updateSortCallback : null;
-
-    const handleSelectAll = useCallback(
-        event => {
-            if (event.target.checked) {
-                const all = ids.concat(
-                    selectedIds.filter(id => !ids.includes(id))
-                );
-                onSelect(
-                    isRowSelectable
-                        ? all.filter(id => isRowSelectable(data[id]))
-                        : all
-                );
-            } else {
-                onSelect([]);
-            }
-        },
-        [data, ids, onSelect, isRowSelectable, selectedIds]
-    );
-
     const lastSelected = useRef(null);
 
     useEffect(() => {
@@ -253,10 +212,6 @@ const Datagrid: FC<DatagridProps> = React.forwardRef((props, ref) => {
         return null;
     }
 
-    const all = isRowSelectable
-        ? ids.filter(id => isRowSelectable(data[id]))
-        : ids;
-
     /**
      * After the initial load, if the data for the list isn't empty,
      * and even if the data is refreshing (e.g. after a filter change),
@@ -270,58 +225,26 @@ const Datagrid: FC<DatagridProps> = React.forwardRef((props, ref) => {
                 size={size}
                 {...sanitizeListRestProps(rest)}
             >
-                <TableHead className={classes.thead}>
-                    <TableRow
-                        className={classnames(classes.row, classes.headerRow)}
-                    >
-                        {expand && (
-                            <TableCell
-                                padding="none"
-                                className={classnames(
-                                    classes.headerCell,
-                                    classes.expandHeader
-                                )}
-                            />
-                        )}
-                        {hasBulkActions && selectedIds && (
-                            <TableCell
-                                padding="checkbox"
-                                className={classes.headerCell}
-                            >
-                                <Checkbox
-                                    className="select-all"
-                                    color="primary"
-                                    checked={
-                                        selectedIds.length > 0 &&
-                                        all.length > 0 &&
-                                        all.every(id =>
-                                            selectedIds.includes(id)
-                                        )
-                                    }
-                                    onChange={handleSelectAll}
-                                />
-                            </TableCell>
-                        )}
-                        {Children.map(children, (field, index) =>
-                            isValidElement(field) ? (
-                                <DatagridHeaderCell
-                                    className={classes.headerCell}
-                                    currentSort={currentSort}
-                                    field={field}
-                                    isSorting={
-                                        currentSort.field ===
-                                        ((field.props as any).sortBy ||
-                                            (field.props as any).source)
-                                    }
-                                    key={(field.props as any).source || index}
-                                    resource={resource}
-                                    updateSort={updateSort}
-                                />
-                            ) : null
-                        )}
-                    </TableRow>
-                </TableHead>
-                {cloneElement(
+                {createOrCloneElement(
+                    header,
+                    {
+                        children,
+                        classes,
+                        className,
+                        currentSort,
+                        data,
+                        hasExpand: !!expand,
+                        hasBulkActions,
+                        ids,
+                        isRowSelectable,
+                        onSelect,
+                        resource,
+                        selectedIds,
+                        setSort,
+                    },
+                    children
+                )}
+                {createOrCloneElement(
                     body,
                     {
                         basePath,
@@ -347,9 +270,15 @@ const Datagrid: FC<DatagridProps> = React.forwardRef((props, ref) => {
     );
 });
 
+const createOrCloneElement = (element, props, children) =>
+    isValidElement(element)
+        ? cloneElement(element, props, children)
+        : createElement(element, props, children);
+
 Datagrid.propTypes = {
     basePath: PropTypes.string,
-    body: PropTypes.element,
+    // @ts-ignore
+    body: PropTypes.oneOfType([PropTypes.element, PropTypes.elementType]),
     children: PropTypes.node.isRequired,
     classes: PropTypes.object,
     className: PropTypes.string,
@@ -362,6 +291,8 @@ Datagrid.propTypes = {
     // @ts-ignore
     expand: PropTypes.oneOfType([PropTypes.element, PropTypes.elementType]),
     hasBulkActions: PropTypes.bool,
+    // @ts-ignore
+    header: PropTypes.oneOfType([PropTypes.element, PropTypes.elementType]),
     hover: PropTypes.bool,
     ids: PropTypes.arrayOf(PropTypes.any),
     loading: PropTypes.bool,
@@ -380,7 +311,7 @@ Datagrid.propTypes = {
 
 export interface DatagridProps<RecordType extends Record = Record>
     extends Omit<TableProps, 'size' | 'classes' | 'onSelect'> {
-    body?: ReactElement;
+    body?: ReactElement | ComponentType;
     classes?: ClassesOverride<typeof useDatagridStyles>;
     className?: string;
     expand?:
@@ -392,6 +323,7 @@ export interface DatagridProps<RecordType extends Record = Record>
               resource: string;
           }>;
     hasBulkActions?: boolean;
+    header?: ReactElement | ComponentType;
     hover?: boolean;
     empty?: ReactElement;
     isRowSelectable?: (record: Record) => boolean;