fix: associate DataTable TableCells with column headers (#22169)

adamalston · web-flow · commit feca78681865 · 2026-05-27T09:04:55.000Z
diff --git a/packages/react/src/components/DataTable/DataTable.tsx b/packages/react/src/components/DataTable/DataTable.tsx
@@ -135,13 +135,15 @@ export interface DataTableRenderProps<RowType, ColTypes extends any[]> {
 
   getHeaderProps: (options: {
     header: DataTableHeader;
+    id?: string;
     isSortable?: boolean;
     onClick?: (
       event: MouseEvent<HTMLButtonElement>,
       sortState: { sortHeaderKey: string; sortDirection: DataTableSortState }
     ) => void;
     [key: string]: unknown;
   }) => {
+    id: string;
     isSortable: boolean | undefined;
     isSortHeader: boolean;
     key: string;
@@ -235,9 +237,14 @@ export interface DataTableRenderProps<RowType, ColTypes extends any[]> {
     useStaticWidth?: boolean;
   };
 
-  getCellProps: (options: { cell: DataTableCell<ColTypes> }) => {
+  getCellProps: (options: {
+    cell: DataTableCell<ColTypes>;
+    headers?: string;
+    [key: string]: unknown;
+  }) => {
     [key: string]: unknown;
     hasAILabelHeader?: boolean;
+    headers: string;
     key: string;
   };
 
@@ -402,9 +409,14 @@ export const DataTable = <RowType, ColTypes extends any[]>(
   }) => {
     const { sortDirection, sortHeaderKey } = state;
     const { key, slug, decorator } = header;
+    const id =
+      typeof rest.id === 'string' && rest.id.length
+        ? rest.id
+        : getHeaderId(key);
 
     return {
       ...rest,
+      id,
       key,
       sortDirection,
       isSortable: headerIsSortable ?? header.isSortable ?? isSortable,
@@ -609,12 +621,23 @@ export const DataTable = <RowType, ColTypes extends any[]>(
   };
 
   const getCellProps: RenderProps['getCellProps'] = ({
-    cell: { hasAILabelHeader, id },
+    cell: {
+      hasAILabelHeader,
+      id,
+      info: { header },
+    },
+    headers: customHeaders,
     ...rest
   }) => {
+    const headers =
+      typeof customHeaders === 'string' && customHeaders.length
+        ? customHeaders
+        : getHeaderId(header);
+
     return {
       ...rest,
       hasAILabelHeader,
+      headers,
       key: id,
     };
   };
@@ -644,6 +667,9 @@ export const DataTable = <RowType, ColTypes extends any[]>(
    */
   const getTablePrefix = () => `data-table-${instanceId}`;
 
+  const getHeaderId = (headerKey: string) =>
+    `${getTablePrefix()}__header-${headerKey}`;
+
   /**
    * Generates a new `rowsById` object with updated selection state.
    */
diff --git a/packages/react/src/components/DataTable/__tests__/DataTable-test.js b/packages/react/src/components/DataTable/__tests__/DataTable-test.js
@@ -155,6 +155,75 @@ describe('DataTable', () => {
       const { container } = render(<DataTable {...mockProps} />);
       expect(container).toMatchSnapshot();
     });
+
+    it('should associate data cells with their column headers', () => {
+      render(<DataTable {...mockProps} />);
+
+      const [firstHeader, secondHeader] = screen.getAllByRole('columnheader');
+      const cells = screen.getAllByRole('cell');
+
+      expect(firstHeader).toHaveAttribute(
+        'id',
+        expect.stringMatching(/^data-table-\d+__header-fieldA$/)
+      );
+      expect(secondHeader).toHaveAttribute(
+        'id',
+        expect.stringMatching(/^data-table-\d+__header-fieldB$/)
+      );
+      expect(cells[0]).toHaveAttribute('headers', firstHeader.id);
+      expect(cells[1]).toHaveAttribute('headers', secondHeader.id);
+      expect(cells[2]).toHaveAttribute('headers', firstHeader.id);
+      expect(cells[3]).toHaveAttribute('headers', secondHeader.id);
+    });
+
+    it('should preserve custom header and cell associations', () => {
+      const { render: _render, ...props } = mockProps;
+
+      render(
+        <DataTable {...props}>
+          {({ getCellProps, getHeaderProps, getTableProps, headers, rows }) => (
+            <Table {...getTableProps()}>
+              <TableHead>
+                <TableRow>
+                  {headers.map((header) => (
+                    <TableHeader
+                      {...getHeaderProps({
+                        header,
+                        id: `custom-header-${header.key}`,
+                      })}>
+                      {header.header}
+                    </TableHeader>
+                  ))}
+                </TableRow>
+              </TableHead>
+              <TableBody>
+                {rows.map((row) => (
+                  <TableRow key={row.id}>
+                    {row.cells.map((cell) => (
+                      <TableCell
+                        {...getCellProps({
+                          cell,
+                          headers: `custom-headers-${cell.info.header}`,
+                        })}>
+                        {cell.value}
+                      </TableCell>
+                    ))}
+                  </TableRow>
+                ))}
+              </TableBody>
+            </Table>
+          )}
+        </DataTable>
+      );
+
+      const [firstHeader, secondHeader] = screen.getAllByRole('columnheader');
+      const cells = screen.getAllByRole('cell');
+
+      expect(firstHeader).toHaveAttribute('id', 'custom-header-fieldA');
+      expect(secondHeader).toHaveAttribute('id', 'custom-header-fieldB');
+      expect(cells[0]).toHaveAttribute('headers', 'custom-headers-fieldA');
+      expect(cells[1]).toHaveAttribute('headers', 'custom-headers-fieldB');
+    });
   });
 
   describe('behaves as expected', () => {
diff --git a/packages/react/src/components/DataTable/__tests__/__snapshots__/DataTable-test.js.snap b/packages/react/src/components/DataTable/__tests__/__snapshots__/DataTable-test.js.snap
