From d1e861520f03ba525982973dd1cceadfa35de467 Mon Sep 17 00:00:00 2001
From: Andrei Zhaleznichenka <zhalea@amazon.de>
Date: Fri, 15 Nov 2024 09:35:50 +0100
Subject: [PATCH] refactor: Componentizes table cell types and adds
 permutations

---
 pages/table/cell-permutations.page.tsx        | 499 ++++++++++++++++++
 .../body-cell/disabled-inline-editor.tsx      |   3 +-
 src/table/body-cell/index.tsx                 |   8 +-
 src/table/body-cell/td-element.tsx            |   4 +-
 src/table/header-cell/index.tsx               |  12 +
 src/table/header-cell/th-element.tsx          |  20 +-
 src/table/internal.tsx                        |  60 +--
 src/table/no-data-cell.tsx                    |   2 +-
 .../progressive-loading/items-loader.tsx      |   2 +-
 src/table/progressive-loading/loader-cell.tsx |  34 ++
 src/table/selection/selection-cell.tsx        |  71 +++
 src/table/thead.tsx                           |  66 +--
 12 files changed, 685 insertions(+), 96 deletions(-)
 create mode 100644 pages/table/cell-permutations.page.tsx
 create mode 100644 src/table/progressive-loading/loader-cell.tsx
 create mode 100644 src/table/selection/selection-cell.tsx

diff --git a/pages/table/cell-permutations.page.tsx b/pages/table/cell-permutations.page.tsx
new file mode 100644
index 0000000000..93ecf0b0f3
--- /dev/null
+++ b/pages/table/cell-permutations.page.tsx
@@ -0,0 +1,499 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import React, { useContext, useRef, useState } from 'react';
+import { range } from 'lodash';
+
+import {
+  Box,
+  Button,
+  Checkbox,
+  Container,
+  ExpandableSection,
+  FormField,
+  Header,
+  Input,
+  Slider,
+  SpaceBetween,
+  StatusIndicator,
+  TableProps,
+} from '~components';
+import { useMergeRefs } from '~components/internal/hooks/use-merge-refs';
+import { TableBodyCell, TableBodyCellProps } from '~components/table/body-cell';
+import { TableHeaderCell, TableHeaderCellProps } from '~components/table/header-cell';
+import { NoDataCell } from '~components/table/no-data-cell';
+import { TableLoaderCell, TableLoaderCellProps } from '~components/table/progressive-loading/loader-cell';
+import { TableBodySelectionCell, TableHeaderSelectionCell } from '~components/table/selection/selection-cell';
+import { StickyColumnsModel, useStickyColumns } from '~components/table/sticky-columns';
+
+import AppContext, { AppContextType } from '../app/app-context';
+import ScreenshotArea from '../utils/screenshot-area';
+
+type PageContext = React.Context<
+  AppContextType<{
+    wrapLines?: boolean;
+    verticalAlignTop?: boolean;
+    sortingDisabled?: boolean;
+    resizableColumns?: boolean;
+    resizableColumnWidth?: string;
+    isExpandable?: boolean;
+    isExpanded?: boolean;
+    isEditable?: boolean;
+    stripedRows?: boolean;
+    hasSelection?: boolean;
+    multiSelection?: boolean;
+    hasFooter?: boolean;
+    stickyColumnsFirst?: string;
+    stickyColumnsLast?: string;
+    tableEmpty?: boolean;
+    tableLoading?: boolean;
+  }>
+>;
+
+const columns = range(0, 10).map(index => index + 1);
+
+export default function InlineEditorPermutations() {
+  const { settings, setUrlParams } = usePageSettings();
+  return (
+    <Box margin="m">
+      <h1>Table cell permutations</h1>
+
+      <ScreenshotArea gutters={false}>
+        <ExpandableSection variant="stacked" headerText="Settings" headingTagOverride="h2" defaultExpanded={true}>
+          <SpaceBetween size="m">
+            <SpaceBetween size="m" direction="horizontal" alignItems="center">
+              <Checkbox
+                checked={settings.wrapLines}
+                onChange={event => setUrlParams({ wrapLines: event.detail.checked })}
+              >
+                Wrap lines
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.isExpandable}
+                onChange={event => setUrlParams({ isExpandable: event.detail.checked })}
+              >
+                Is expandable
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.isExpanded}
+                onChange={event => setUrlParams({ isExpanded: event.detail.checked })}
+              >
+                Is expanded
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.isEditable}
+                onChange={event => setUrlParams({ isEditable: event.detail.checked })}
+              >
+                Editable
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.sortingDisabled}
+                onChange={event => setUrlParams({ sortingDisabled: event.detail.checked })}
+              >
+                Sorting disabled
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.verticalAlign === 'top'}
+                onChange={event => setUrlParams({ verticalAlignTop: event.detail.checked })}
+              >
+                Vertical align top
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.stripedRows}
+                onChange={event => setUrlParams({ stripedRows: event.detail.checked })}
+              >
+                Striped rows
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.hasSelection}
+                onChange={event => setUrlParams({ hasSelection: event.detail.checked })}
+              >
+                Has selection
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.multiSelection}
+                onChange={event => setUrlParams({ multiSelection: event.detail.checked })}
+              >
+                Multi selection
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.hasFooter}
+                onChange={event => setUrlParams({ hasFooter: event.detail.checked })}
+              >
+                Has footer
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.resizableColumns}
+                onChange={event => setUrlParams({ resizableColumns: event.detail.checked })}
+              >
+                Resizable columns
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.tableEmpty}
+                onChange={event => setUrlParams({ tableEmpty: event.detail.checked })}
+              >
+                Table empty
+              </Checkbox>
+
+              <Checkbox
+                checked={settings.tableLoading}
+                onChange={event => setUrlParams({ tableLoading: event.detail.checked })}
+              >
+                Table loading
+              </Checkbox>
+            </SpaceBetween>
+
+            <SpaceBetween size="m" direction="horizontal" alignItems="center">
+              <FormField label="Resizable column width">
+                <Slider
+                  value={settings.resizableColumnWidth}
+                  onChange={event => setUrlParams({ resizableColumnWidth: event.detail.value.toString() })}
+                  min={50}
+                  max={300}
+                />
+              </FormField>
+
+              <FormField label="Sticky columns first">
+                <Slider
+                  value={settings.stickyColumnsFirst}
+                  onChange={event => setUrlParams({ stickyColumnsFirst: event.detail.value.toString() })}
+                  min={0}
+                  max={3}
+                />
+              </FormField>
+
+              <FormField label="Sticky columns last">
+                <Slider
+                  value={settings.stickyColumnsLast}
+                  onChange={event => setUrlParams({ stickyColumnsLast: event.detail.value.toString() })}
+                  min={0}
+                  max={3}
+                />
+              </FormField>
+            </SpaceBetween>
+          </SpaceBetween>
+        </ExpandableSection>
+
+        <Container variant="stacked" header={<Header variant="h2">Standard cells</Header>}>
+          <StandardCellsTable />
+        </Container>
+      </ScreenshotArea>
+    </Box>
+  );
+}
+
+function StandardCellsTable() {
+  const { settings } = usePageSettings();
+  const containerRefObject = useRef<HTMLDivElement>(null);
+  const containerRef = useMergeRefs(containerRefObject, settings.stickyState.refs.wrapper);
+  const tableRefObject = useRef<HTMLTableElement>(null);
+  const tableRef = useMergeRefs(tableRefObject, settings.stickyState.refs.table);
+  return (
+    <div ref={containerRef} style={{ overflowX: 'auto', ...settings.stickyState.style.wrapper }}>
+      <table ref={tableRef} style={{ tableLayout: 'fixed', borderSpacing: 0, width: '100%' }}>
+        <thead>
+          <TextHeaders {...settings} />
+        </thead>
+        <tbody>
+          {settings.tableEmpty || settings.tableLoading ? (
+            <tr>
+              <NoDataCell
+                totalColumnsCount={columns.length}
+                hasFooter={settings.hasFooter}
+                loading={settings.tableLoading}
+                loadingText="Loading"
+                empty="Empty"
+                containerRef={containerRefObject}
+                tableRef={tableRefObject}
+              />
+            </tr>
+          ) : (
+            <>
+              <TextColumns {...settings} level={1} isEvenRow={false} isFirstRow={true} />
+              <TextColumns {...settings} level={2} isEvenRow={true} />
+              <TextColumns {...settings} level={3} isEvenRow={false} isNextSelected={true} />
+              <TextColumns {...settings} level={4} isEvenRow={true} isSelected={true} />
+              <TextColumns {...settings} level={5} isEvenRow={false} isPrevSelected={true} isNextSelected={true} />
+              <TextColumns {...settings} level={1} isEvenRow={true} isSelected={true} isNextSelected={true} />
+              <TextColumns {...settings} level={2} isEvenRow={false} isSelected={true} isPrevSelected={true} />
+              <TextColumns {...settings} level={3} isEvenRow={true} isPrevSelected={true} />
+              <LoaderColumns {...settings} level={3} isEvenRow={false} loadingStatus="pending" />
+              <LoaderColumns {...settings} level={2} isEvenRow={false} loadingStatus="loading" />
+              <LoaderColumns {...settings} level={1} isEvenRow={false} isLastRow={true} loadingStatus="error" />
+            </>
+          )}
+        </tbody>
+      </table>
+    </div>
+  );
+}
+
+function TextHeaders(
+  props: Partial<TableHeaderCellProps<string>> & {
+    stickyState: StickyColumnsModel;
+    resizableColumnWidth?: number;
+    isExpandable?: boolean;
+    hasSelection?: boolean;
+    multiSelection?: boolean;
+  }
+) {
+  return (
+    <tr>
+      {props.hasSelection && (
+        <TableHeaderSelectionCell
+          {...defaultCellProps()}
+          {...props}
+          onFocusMove={() => {}}
+          columnId="selection"
+          getSelectAllProps={
+            props.multiSelection
+              ? () => ({
+                  name: 'select-all',
+                  disabled: false,
+                  selectionType: 'multi',
+                  indeterminate: true,
+                  checked: false,
+                  onChange: () => {},
+                })
+              : undefined
+          }
+        />
+      )}
+      {columns.map(index => (
+        <PlaygroundHeaderCell
+          key={index}
+          {...props}
+          colIndex={index}
+          isExpandable={index === 1 && props.isExpandable}
+          sortable={[2, 3, 4, 5].includes(index)}
+          activeSorting={[4, 5].includes(index)}
+          sortingDescending={[3, 5].includes(index)}
+        >
+          {`Header cell content ${index}${index === 8 ? ' with longer text' : ''}`}
+        </PlaygroundHeaderCell>
+      ))}
+    </tr>
+  );
+}
+
+function TextColumns(
+  props: Partial<TableBodyCellProps<string>> & {
+    level: number;
+    stickyState: StickyColumnsModel;
+    isExpandable?: boolean;
+    isExpanded?: boolean;
+    multiSelection?: boolean;
+  }
+) {
+  return (
+    <tr>
+      {props.hasSelection && (
+        <TableBodySelectionCell
+          {...defaultCellProps()}
+          {...props}
+          columnId="selection"
+          isExpandable={false}
+          selectionControlProps={{
+            name: 'select-one',
+            disabled: false,
+            selectionType: props.multiSelection ? 'multi' : 'single',
+            checked: !!props.isSelected,
+            onChange: () => {},
+          }}
+        />
+      )}
+      {columns.map(index => (
+        <PlaygroundBodyCell
+          key={index}
+          {...props}
+          colIndex={index}
+          level={index === 1 && props.isExpandable ? props.level : 0}
+          isExpandable={index === 1 && props.isExpandable}
+          isExpanded={index === 1 && props.isExpanded}
+        >
+          {`Body cell content ${index}${index === 1 ? ` (L=${props.level})` : ''}${index === 8 ? ' with longer text' : ''}`}
+        </PlaygroundBodyCell>
+      ))}
+    </tr>
+  );
+}
+
+function LoaderColumns(
+  props: Partial<TableLoaderCellProps<string>> & {
+    level: number;
+    stickyState: StickyColumnsModel;
+    loadingStatus: TableProps.LoadingStatus;
+  }
+) {
+  return (
+    <tr>
+      {props.hasSelection && (
+        <TableBodySelectionCell {...defaultCellProps()} {...props} columnId="selection" isExpandable={false} />
+      )}
+      {columns.map(index => (
+        <TableLoaderCell
+          {...defaultCellProps()}
+          {...props}
+          key={index}
+          colIndex={index}
+          columnId={index.toString()}
+          level={index === 1 && props.isExpandable ? props.level - 1 : 0}
+          isRowHeader={index === 1}
+          isExpandable={false}
+          isExpanded={false}
+          item={''}
+          renderLoaderPending={() => <Button>Load more</Button>}
+          renderLoaderLoading={() => <StatusIndicator type="loading">Loading more</StatusIndicator>}
+          renderLoaderError={() => <StatusIndicator type="error">Error when loading more</StatusIndicator>}
+        />
+      ))}
+    </tr>
+  );
+}
+
+function PlaygroundHeaderCell(
+  props: Partial<TableHeaderCellProps<string>> & {
+    children: React.ReactNode;
+    colIndex: number;
+    resizableColumnWidth?: number;
+    stickyState: StickyColumnsModel;
+    sortable?: boolean;
+    activeSorting?: boolean;
+  }
+) {
+  const column: TableProps.ColumnDefinition<string> = {
+    id: props.colIndex.toString(),
+    header: props.children,
+    cell: () => '',
+    sortingField: props.sortable ? 'field' : undefined,
+  };
+  return (
+    <TableHeaderCell
+      style={{ width: props.resizableColumns ? props.resizableColumnWidth : undefined }}
+      tabIndex={-1}
+      column={column}
+      activeSortingColumn={props.activeSorting ? column : undefined}
+      columnId={props.colIndex.toString()}
+      hidden={false}
+      {...defaultCellProps()}
+      onResizeFinish={() => {}}
+      updateColumn={() => {}}
+      onClick={() => {}}
+      cellRef={() => {}}
+      tableRole="grid"
+      {...props}
+    />
+  );
+}
+
+function PlaygroundBodyCell(
+  props: Partial<TableBodyCellProps<string>> & {
+    children: React.ReactNode;
+    stickyState: StickyColumnsModel;
+    colIndex: number;
+    isEditable?: boolean;
+  }
+) {
+  const [isEditing, setIsEditing] = useState(false);
+  const [successfulEdit, setSuccessfulEdit] = useState(false);
+  const [editedValue, setEditedValue] = useState('');
+  return (
+    <TableBodyCell
+      ariaLabels={{
+        activateEditLabel: column => `Edit ${column.header}`,
+        cancelEditLabel: column => `Cancel editing ${column.header}`,
+        submitEditLabel: column => `Submit edit ${column.header}`,
+        successfulEditLabel: () => 'Edit successful',
+      }}
+      column={{
+        id: props.colIndex.toString(),
+        header: '',
+        cell: () => editedValue || props.children,
+        editConfig: {
+          ariaLabel: 'Edit dialog aria label',
+          editIconAriaLabel: 'Edit icon label',
+          errorIconAriaLabel: 'Edit error icon label',
+          validation(_item, value = '') {
+            if (value.trim() && value.toLowerCase().includes('content')) {
+              return 'Must not include "content"';
+            }
+          },
+          editingCell(_item, { currentValue, setValue }: TableProps.CellContext<string>) {
+            return (
+              <Input
+                autoFocus={true}
+                value={currentValue ?? (typeof props.children === 'string' ? props.children : '')}
+                onChange={event => setValue(event.detail.value)}
+              />
+            );
+          },
+        },
+      }}
+      isEditing={isEditing}
+      successfulEdit={successfulEdit}
+      onEditStart={() => setIsEditing(true)}
+      onEditEnd={() => {
+        setIsEditing(false);
+        setSuccessfulEdit(true);
+      }}
+      submitEdit={(_item, _column, newValue) => new Promise(resolve => resolve(setEditedValue(newValue as any)))}
+      columnId={props.colIndex.toString()}
+      item={''}
+      {...defaultCellProps()}
+      {...props}
+    />
+  );
+}
+
+function defaultCellProps() {
+  return {
+    isFirstRow: false,
+    isLastRow: false,
+    isSelected: false,
+    isNextSelected: false,
+    isPrevSelected: false,
+    wrapLines: false,
+    isEditable: false,
+    resizableColumns: false,
+    tableRole: 'grid' as const,
+    variant: 'container' as const,
+  };
+}
+
+function usePageSettings() {
+  const { urlParams, setUrlParams } = useContext(AppContext as PageContext);
+  const settings = {
+    isExpandable: urlParams.isExpandable ?? false,
+    isExpanded: urlParams.isExpanded ?? false,
+    isEditable: urlParams.isEditable ?? false,
+    sortingDisabled: urlParams.sortingDisabled ?? false,
+    resizableColumns: urlParams.resizableColumns ?? true,
+    resizableColumnWidth: parseInt(urlParams.resizableColumnWidth || '') || 150,
+    wrapLines: urlParams.wrapLines ?? false,
+    stripedRows: urlParams.stripedRows ?? false,
+    hasSelection: urlParams.hasSelection ?? true,
+    multiSelection: urlParams.multiSelection ?? true,
+    hasFooter: urlParams.hasFooter ?? false,
+    verticalAlign: (urlParams.verticalAlignTop ? 'top' : 'middle') as 'top' | 'middle',
+    stickyColumnsFirst: parseInt(urlParams.stickyColumnsFirst ?? '') || 0,
+    stickyColumnsLast: parseInt(urlParams.stickyColumnsLast ?? '') || 0,
+    tableEmpty: urlParams.tableEmpty ?? false,
+    tableLoading: urlParams.tableLoading ?? false,
+  };
+  const stickyState = useStickyColumns({
+    visibleColumns: columns.map(index => index.toString()),
+    stickyColumnsFirst: settings.stickyColumnsFirst,
+    stickyColumnsLast: settings.stickyColumnsLast,
+  });
+  return { settings: { ...settings, stickyState }, setUrlParams };
+}
diff --git a/src/table/body-cell/disabled-inline-editor.tsx b/src/table/body-cell/disabled-inline-editor.tsx
index 506d08b565..0cb96e7f88 100644
--- a/src/table/body-cell/disabled-inline-editor.tsx
+++ b/src/table/body-cell/disabled-inline-editor.tsx
@@ -8,6 +8,7 @@ import Portal from '../../internal/components/portal';
 import { useSingleTabStopNavigation } from '../../internal/context/single-tab-stop-navigation-context';
 import useHiddenDescription from '../../internal/hooks/use-hidden-description';
 import { usePortalModeClasses } from '../../internal/hooks/use-portal-mode-classes';
+import { useVisualRefresh } from '../../internal/hooks/use-visual-mode';
 import InternalLiveRegion from '../../live-region/internal';
 import Arrow from '../../popover/arrow';
 import PopoverBody from '../../popover/body';
@@ -31,10 +32,10 @@ export function DisabledInlineEditor<ItemType>({
   onEditStart,
   onEditEnd,
   editDisabledReason,
-  isVisualRefresh,
   resizableColumns = false,
   ...rest
 }: DisabledInlineEditorProps<ItemType>) {
+  const isVisualRefresh = useVisualRefresh();
   const clickAwayRef = useClickAway(() => {
     if (isEditing) {
       onEditEnd(true);
diff --git a/src/table/body-cell/index.tsx b/src/table/body-cell/index.tsx
index 63fe338d50..1aae5b3506 100644
--- a/src/table/body-cell/index.tsx
+++ b/src/table/body-cell/index.tsx
@@ -7,6 +7,7 @@ import { useInternalI18n } from '../../i18n/context';
 import Icon from '../../icon/internal';
 import { useSingleTabStopNavigation } from '../../internal/context/single-tab-stop-navigation-context.js';
 import { usePrevious } from '../../internal/hooks/use-previous';
+import { useVisualRefresh } from '../../internal/hooks/use-visual-mode';
 import InternalLiveRegion from '../../live-region/internal';
 import { TableProps } from '../interfaces';
 import { DisabledInlineEditor } from './disabled-inline-editor';
@@ -40,12 +41,12 @@ function TableCellEditable<ItemType>({
   onEditEnd,
   submitEdit,
   ariaLabels,
-  isVisualRefresh,
   resizableColumns = false,
   successfulEdit = false,
   ...rest
 }: TableBodyCellProps<ItemType>) {
   const i18n = useInternalI18n('table');
+  const isVisualRefresh = useVisualRefresh();
   const editActivateRef = useRef<HTMLButtonElement>(null);
   const tdNativeAttributes = {
     'data-inline-editing-active': isEditing.toString(),
@@ -155,15 +156,14 @@ export function TableBodyCell<ItemType>({
   isEditable,
   ...rest
 }: TableBodyCellProps<ItemType> & { isEditable: boolean }) {
-  const isExpandableColumnCell = rest.level !== undefined;
   const editDisabledReason = rest.column.editConfig?.disabledReason?.(rest.item);
 
   // Inline editing is deactivated for expandable column because editable cells are interactive
   // and cannot include interactive content such as expand toggles.
-  if (editDisabledReason && !isExpandableColumnCell) {
+  if (editDisabledReason) {
     return <DisabledInlineEditor editDisabledReason={editDisabledReason} {...rest} />;
   }
-  if ((isEditable || rest.isEditing) && !isExpandableColumnCell) {
+  if (isEditable || rest.isEditing) {
     return <TableCellEditable {...rest} />;
   }
 
diff --git a/src/table/body-cell/td-element.tsx b/src/table/body-cell/td-element.tsx
index 9c8c83512a..ff043e7733 100644
--- a/src/table/body-cell/td-element.tsx
+++ b/src/table/body-cell/td-element.tsx
@@ -7,6 +7,7 @@ import { copyAnalyticsMetadataAttribute } from '@cloudscape-design/component-too
 
 import { useSingleTabStopNavigation } from '../../internal/context/single-tab-stop-navigation-context';
 import { useMergeRefs } from '../../internal/hooks/use-merge-refs';
+import { useVisualRefresh } from '../../internal/hooks/use-visual-mode';
 import { ExpandToggleButton } from '../expandable-rows/expand-toggle-button';
 import { TableProps } from '../interfaces.js';
 import { StickyColumnsModel, useStickyCellStyles } from '../sticky-columns';
@@ -40,7 +41,6 @@ export interface TableTdElementProps {
   columnId: PropertyKey;
   colIndex: number;
   stickyState: StickyColumnsModel;
-  isVisualRefresh?: boolean;
   tableRole: TableRole;
   level?: number;
   isExpandable?: boolean;
@@ -70,7 +70,6 @@ export const TableTdElement = React.forwardRef<HTMLTableCellElement, TableTdElem
       onMouseLeave,
       isEvenRow,
       stripedRows,
-      isVisualRefresh,
       hasSelection,
       hasFooter,
       columnId,
@@ -89,6 +88,7 @@ export const TableTdElement = React.forwardRef<HTMLTableCellElement, TableTdElem
     ref
   ) => {
     const Element = isRowHeader ? 'th' : 'td';
+    const isVisualRefresh = useVisualRefresh();
 
     nativeAttributes = { ...nativeAttributes, ...getTableCellRoleProps({ tableRole, isRowHeader, colIndex }) };
 
diff --git a/src/table/header-cell/index.tsx b/src/table/header-cell/index.tsx
index 479d957395..a2be29c72e 100644
--- a/src/table/header-cell/index.tsx
+++ b/src/table/header-cell/index.tsx
@@ -32,7 +32,10 @@ export interface TableHeaderCellProps<ItemType> {
   sortingDescending?: boolean;
   sortingDisabled?: boolean;
   wrapLines?: boolean;
+  stuck?: boolean;
+  sticky?: boolean;
   hidden?: boolean;
+  stripedRows?: boolean;
   onClick(detail: TableProps.SortingState<any>): void;
   onResizeFinish: () => void;
   colIndex: number;
@@ -47,6 +50,7 @@ export interface TableHeaderCellProps<ItemType> {
   resizerRoleDescription?: string;
   isExpandable?: boolean;
   hasDynamicContent?: boolean;
+  variant: TableProps.Variant;
 }
 
 export function TableHeaderCell<ItemType>({
@@ -59,7 +63,10 @@ export function TableHeaderCell<ItemType>({
   sortingDisabled,
   wrapLines,
   focusedComponent,
+  stuck,
+  sticky,
   hidden,
+  stripedRows,
   onClick,
   colIndex,
   updateColumn,
@@ -73,6 +80,7 @@ export function TableHeaderCell<ItemType>({
   resizerRoleDescription,
   isExpandable,
   hasDynamicContent,
+  variant,
 }: TableHeaderCellProps<ItemType>) {
   const i18n = useInternalI18n('table');
   const sortable = !!column.sortingComparator || !!column.sortingField;
@@ -120,11 +128,15 @@ export function TableHeaderCell<ItemType>({
       sortingStatus={sortingStatus}
       sortingDisabled={sortingDisabled}
       focusedComponent={focusedComponent}
+      stuck={stuck}
+      sticky={sticky}
       hidden={hidden}
+      stripedRows={stripedRows}
       colIndex={colIndex}
       columnId={columnId}
       stickyState={stickyState}
       tableRole={tableRole}
+      variant={variant}
       {...(sortingDisabled
         ? {}
         : getAnalyticsMetadataAttribute({
diff --git a/src/table/header-cell/th-element.tsx b/src/table/header-cell/th-element.tsx
index dbbf8aedee..ddbc0dcefe 100644
--- a/src/table/header-cell/th-element.tsx
+++ b/src/table/header-cell/th-element.tsx
@@ -7,6 +7,8 @@ import { copyAnalyticsMetadataAttribute } from '@cloudscape-design/component-too
 
 import { useSingleTabStopNavigation } from '../../internal/context/single-tab-stop-navigation-context';
 import { useMergeRefs } from '../../internal/hooks/use-merge-refs';
+import { useVisualRefresh } from '../../internal/hooks/use-visual-mode';
+import { TableProps } from '../interfaces';
 import { StickyColumnsModel, useStickyCellStyles } from '../sticky-columns';
 import { getTableColHeaderRoleProps, TableRole } from '../table-role';
 import { getStickyClassNames } from '../utils';
@@ -14,19 +16,23 @@ import { SortingStatus } from './utils';
 
 import styles from './styles.css.js';
 
-interface TableThElementProps {
+export interface TableThElementProps {
   className?: string;
   style?: React.CSSProperties;
   sortingStatus?: SortingStatus;
   sortingDisabled?: boolean;
   focusedComponent?: null | string;
+  stuck?: boolean;
+  sticky?: boolean;
   hidden?: boolean;
+  stripedRows?: boolean;
   colIndex: number;
   columnId: PropertyKey;
   stickyState: StickyColumnsModel;
   cellRef?: React.RefCallback<HTMLElement> | null;
   tableRole: TableRole;
   children: React.ReactNode;
+  variant: TableProps.Variant;
 }
 
 export function TableThElement({
@@ -35,15 +41,21 @@ export function TableThElement({
   sortingStatus,
   sortingDisabled,
   focusedComponent,
+  stuck,
+  sticky,
   hidden,
+  stripedRows,
   colIndex,
   columnId,
   stickyState,
   cellRef,
   tableRole,
   children,
+  variant,
   ...props
 }: TableThElementProps) {
+  const isVisualRefresh = useVisualRefresh();
+
   const stickyStyles = useStickyCellStyles({
     stickyColumns: stickyState,
     columnId,
@@ -58,6 +70,12 @@ export function TableThElement({
     <th
       data-focus-id={`header-${String(columnId)}`}
       className={clsx(
+        styles['header-cell'],
+        styles[`header-cell-variant-${variant}`],
+        sticky && styles['header-cell-sticky'],
+        stuck && styles['header-cell-stuck'],
+        stripedRows && styles['has-striped-rows'],
+        isVisualRefresh && styles['is-visual-refresh'],
         className,
         {
           [styles['header-cell-fake-focus']]: focusedComponent === `header-${String(columnId)}`,
diff --git a/src/table/internal.tsx b/src/table/internal.tsx
index 200eb30f9b..779eb22307 100644
--- a/src/table/internal.tsx
+++ b/src/table/internal.tsx
@@ -32,15 +32,15 @@ import { SomeRequired } from '../internal/types';
 import InternalLiveRegion from '../live-region/internal';
 import { GeneratedAnalyticsMetadataTableComponent } from './analytics-metadata/interfaces';
 import { TableBodyCell } from './body-cell';
-import { TableTdElement } from './body-cell/td-element';
 import { checkColumnWidths } from './column-widths-utils';
 import { useExpandableTableProps } from './expandable-rows/expandable-rows-utils';
 import { TableForwardRefType, TableProps, TableRow } from './interfaces';
 import { NoDataCell } from './no-data-cell';
-import { ItemsLoader } from './progressive-loading/items-loader';
+import { TableLoaderCell } from './progressive-loading/loader-cell';
 import { useProgressiveLoadingProps } from './progressive-loading/progressive-loading-utils';
 import { ResizeTracker } from './resizer';
-import { focusMarkers, SelectionControl, useSelection, useSelectionFocusMove } from './selection';
+import { focusMarkers, useSelection, useSelectionFocusMove } from './selection';
+import { TableBodySelectionCell } from './selection/selection-cell';
 import { useStickyColumns } from './sticky-columns';
 import StickyHeader, { StickyHeaderRef } from './sticky-header';
 import { StickyScrollbar } from './sticky-scrollbar';
@@ -550,7 +550,6 @@ const InternalTable = React.forwardRef(
                           });
                           const getTableItemKey = (item: T) => getItemKey(trackBy, item, rowIndex);
                           const sharedCellProps = {
-                            isVisualRefresh,
                             isFirstRow,
                             isLastRow,
                             isSelected: hasSelection && isRowSelected(row),
@@ -584,21 +583,17 @@ const InternalTable = React.forwardRef(
                                 {...rowRoleProps}
                               >
                                 {getItemSelectionProps && (
-                                  <TableTdElement
+                                  <TableBodySelectionCell
                                     {...sharedCellProps}
-                                    className={styles['selection-control']}
-                                    wrapLines={false}
                                     columnId={selectionColumnId}
-                                    colIndex={0}
-                                  >
-                                    <SelectionControl
-                                      onFocusDown={moveFocusDown}
-                                      onFocusUp={moveFocusUp}
-                                      {...getItemSelectionProps(row.item)}
-                                      rowIndex={rowIndex}
-                                      itemKey={`${getTableItemKey(row.item)}`}
-                                    />
-                                  </TableTdElement>
+                                    selectionControlProps={{
+                                      ...getItemSelectionProps(row.item),
+                                      onFocusDown: moveFocusDown,
+                                      onFocusUp: moveFocusUp,
+                                      rowIndex,
+                                      itemKey: `${getTableItemKey(row.item)}`,
+                                    }}
+                                  />
                                 )}
 
                                 {visibleColumnDefinitions.map((column, colIndex) => {
@@ -667,18 +662,10 @@ const InternalTable = React.forwardRef(
                               {...rowRoleProps}
                             >
                               {getItemSelectionProps && (
-                                <TableTdElement
-                                  {...sharedCellProps}
-                                  className={styles['selection-control']}
-                                  wrapLines={false}
-                                  columnId={selectionColumnId}
-                                  colIndex={0}
-                                >
-                                  {null}
-                                </TableTdElement>
+                                <TableBodySelectionCell {...sharedCellProps} columnId={selectionColumnId} />
                               )}
                               {visibleColumnDefinitions.map((column, colIndex) => (
-                                <TableTdElement
+                                <TableLoaderCell
                                   key={getColumnKey(column, colIndex)}
                                   {...sharedCellProps}
                                   wrapLines={false}
@@ -686,18 +673,13 @@ const InternalTable = React.forwardRef(
                                   colIndex={colIndex + colIndexOffset}
                                   isRowHeader={colIndex === 0}
                                   level={row.level}
-                                >
-                                  {colIndex === 0 ? (
-                                    <ItemsLoader
-                                      item={row.item}
-                                      loadingStatus={row.status}
-                                      renderLoaderPending={renderLoaderPending}
-                                      renderLoaderLoading={renderLoaderLoading}
-                                      renderLoaderError={renderLoaderError}
-                                      trackBy={trackBy}
-                                    />
-                                  ) : null}
-                                </TableTdElement>
+                                  item={row.item}
+                                  loadingStatus={row.status}
+                                  renderLoaderPending={renderLoaderPending}
+                                  renderLoaderLoading={renderLoaderLoading}
+                                  renderLoaderError={renderLoaderError}
+                                  trackBy={trackBy}
+                                />
                               ))}
                             </tr>
                           );
diff --git a/src/table/no-data-cell.tsx b/src/table/no-data-cell.tsx
index bc896d5f37..4ef1707902 100644
--- a/src/table/no-data-cell.tsx
+++ b/src/table/no-data-cell.tsx
@@ -10,7 +10,7 @@ import InternalStatusIndicator from '../status-indicator/internal';
 
 import styles from './styles.css.js';
 
-interface NoDataCellProps {
+export interface NoDataCellProps {
   totalColumnsCount: number;
   hasFooter: boolean;
   loading?: boolean;
diff --git a/src/table/progressive-loading/items-loader.tsx b/src/table/progressive-loading/items-loader.tsx
index e8009adb0d..3846429b05 100644
--- a/src/table/progressive-loading/items-loader.tsx
+++ b/src/table/progressive-loading/items-loader.tsx
@@ -10,7 +10,7 @@ import { applyTrackBy } from '../utils';
 
 import styles from './styles.css.js';
 
-interface ItemsLoaderProps<T> {
+export interface ItemsLoaderProps<T> {
   item: null | T;
   loadingStatus: TableProps.LoadingStatus;
   renderLoaderPending?: (detail: TableProps.RenderLoaderDetail<T>) => React.ReactNode;
diff --git a/src/table/progressive-loading/loader-cell.tsx b/src/table/progressive-loading/loader-cell.tsx
new file mode 100644
index 0000000000..d3657ebe48
--- /dev/null
+++ b/src/table/progressive-loading/loader-cell.tsx
@@ -0,0 +1,34 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import React from 'react';
+
+import { TableTdElement, TableTdElementProps } from '../body-cell/td-element';
+import { ItemsLoader, ItemsLoaderProps } from './items-loader';
+
+export interface TableLoaderCellProps<ItemType> extends TableTdElementProps, ItemsLoaderProps<ItemType> {}
+
+export function TableLoaderCell<ItemType>({
+  item,
+  loadingStatus,
+  renderLoaderPending,
+  renderLoaderLoading,
+  renderLoaderError,
+  trackBy,
+  ...props
+}: TableLoaderCellProps<ItemType>) {
+  return (
+    <TableTdElement {...props}>
+      {props.isRowHeader ? (
+        <ItemsLoader
+          item={item}
+          loadingStatus={loadingStatus}
+          renderLoaderPending={renderLoaderPending}
+          renderLoaderLoading={renderLoaderLoading}
+          renderLoaderError={renderLoaderError}
+          trackBy={trackBy}
+        />
+      ) : null}
+    </TableTdElement>
+  );
+}
diff --git a/src/table/selection/selection-cell.tsx b/src/table/selection/selection-cell.tsx
new file mode 100644
index 0000000000..caa19a93d4
--- /dev/null
+++ b/src/table/selection/selection-cell.tsx
@@ -0,0 +1,71 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import React from 'react';
+import clsx from 'clsx';
+
+import ScreenreaderOnly from '../../internal/components/screenreader-only';
+import { useVisualRefresh } from '../../internal/hooks/use-visual-mode';
+import { TableTdElement, TableTdElementProps } from '../body-cell/td-element';
+import { TableThElement, TableThElementProps } from '../header-cell/th-element';
+import { Divider } from '../resizer';
+import { SelectionProps } from './interfaces';
+import { SelectionControl, SelectionControlProps } from './selection-control';
+
+import headerCellStyles from '../header-cell/styles.css.js';
+import styles from '../styles.css.js';
+
+interface TableHeaderSelectionCellProps extends Omit<TableThElementProps, 'children' | 'colIndex'> {
+  focusedComponent?: null | string;
+  singleSelectionHeaderAriaLabel?: string;
+  getSelectAllProps?: () => SelectionProps;
+  onFocusMove: ((sourceElement: HTMLElement, fromIndex: number, direction: -1 | 1) => void) | undefined;
+}
+
+interface TableBodySelectionCellProps extends Omit<TableTdElementProps, 'children' | 'colIndex' | 'wrapLines'> {
+  selectionControlProps?: SelectionControlProps;
+}
+
+export function TableHeaderSelectionCell({
+  focusedComponent,
+  singleSelectionHeaderAriaLabel,
+  getSelectAllProps,
+  onFocusMove,
+  ...props
+}: TableHeaderSelectionCellProps) {
+  const isVisualRefresh = useVisualRefresh();
+  const selectionCellClass = clsx(
+    styles['selection-control'],
+    styles['selection-control-header'],
+    isVisualRefresh && styles['is-visual-refresh']
+  );
+  return (
+    <TableThElement
+      {...props}
+      className={clsx(selectionCellClass, props.hidden && headerCellStyles['header-cell-hidden'])}
+      colIndex={0}
+      focusedComponent={focusedComponent}
+    >
+      {getSelectAllProps ? (
+        <SelectionControl
+          onFocusDown={event => {
+            onFocusMove!(event.target as HTMLElement, -1, +1);
+          }}
+          focusedComponent={focusedComponent}
+          {...getSelectAllProps()}
+          {...(props.sticky ? { tabIndex: -1 } : {})}
+        />
+      ) : (
+        <ScreenreaderOnly>{singleSelectionHeaderAriaLabel}</ScreenreaderOnly>
+      )}
+      <Divider className={styles['resize-divider']} />
+    </TableThElement>
+  );
+}
+
+export function TableBodySelectionCell({ selectionControlProps, ...props }: TableBodySelectionCellProps) {
+  return (
+    <TableTdElement {...props} className={styles['selection-control']} wrapLines={false} colIndex={0}>
+      {selectionControlProps ? <SelectionControl {...selectionControlProps} /> : null}
+    </TableTdElement>
+  );
+}
diff --git a/src/table/thead.tsx b/src/table/thead.tsx
index a9c16bb8d8..6797ab3692 100644
--- a/src/table/thead.tsx
+++ b/src/table/thead.tsx
@@ -6,21 +6,17 @@ import clsx from 'clsx';
 import { findUpUntil } from '@cloudscape-design/component-toolkit/dom';
 import { getAnalyticsMetadataAttribute } from '@cloudscape-design/component-toolkit/internal/analytics-metadata';
 
-import ScreenreaderOnly from '../internal/components/screenreader-only';
 import { fireNonCancelableEvent, NonCancelableEventHandler } from '../internal/events';
-import { useVisualRefresh } from '../internal/hooks/use-visual-mode';
 import { GeneratedAnalyticsMetadataTableSelectAll } from './analytics-metadata/interfaces';
 import { TableHeaderCell } from './header-cell';
-import { TableThElement } from './header-cell/th-element';
 import { TableProps } from './interfaces';
-import { Divider } from './resizer';
-import { focusMarkers, SelectionControl, SelectionProps } from './selection';
+import { focusMarkers, SelectionProps } from './selection';
+import { TableHeaderSelectionCell } from './selection/selection-cell';
 import { StickyColumnsModel } from './sticky-columns';
 import { getTableHeaderRowRoleProps, TableRole } from './table-role';
 import { useColumnWidths } from './use-column-widths';
 import { getColumnKey } from './utils';
 
-import headerCellStyles from './header-cell/styles.css.js';
 import styles from './styles.css.js';
 
 export interface TheadProps {
@@ -82,25 +78,18 @@ const Thead = React.forwardRef(
     }: TheadProps,
     outerRef: React.Ref<HTMLTableRowElement>
   ) => {
-    const isVisualRefresh = useVisualRefresh();
-
-    const headerCellClass = clsx(
-      headerCellStyles['header-cell'],
-      headerCellStyles[`header-cell-variant-${variant}`],
-      sticky && headerCellStyles['header-cell-sticky'],
-      stuck && headerCellStyles['header-cell-stuck'],
-      stripedRows && headerCellStyles['has-striped-rows'],
-      isVisualRefresh && headerCellStyles['is-visual-refresh']
-    );
-
-    const selectionCellClass = clsx(
-      styles['selection-control'],
-      styles['selection-control-header'],
-      isVisualRefresh && styles['is-visual-refresh']
-    );
-
     const { getColumnStyles, columnWidths, updateColumn, setCell } = useColumnWidths();
 
+    const commonCellProps = {
+      stuck,
+      sticky,
+      hidden,
+      stripedRows,
+      tableRole,
+      variant,
+      stickyState,
+    };
+
     return (
       <thead className={clsx(!hidden && styles['thead-active'])}>
         <tr
@@ -116,41 +105,26 @@ const Thead = React.forwardRef(
           onBlur={() => onFocusedComponentChange?.(null)}
         >
           {selectionType ? (
-            <TableThElement
-              className={clsx(headerCellClass, selectionCellClass, hidden && headerCellStyles['header-cell-hidden'])}
-              hidden={hidden}
-              tableRole={tableRole}
-              colIndex={0}
+            <TableHeaderSelectionCell
+              {...commonCellProps}
               focusedComponent={focusedComponent}
               columnId={selectionColumnId}
-              stickyState={stickyState}
+              getSelectAllProps={getSelectAllProps}
+              onFocusMove={onFocusMove}
+              singleSelectionHeaderAriaLabel={singleSelectionHeaderAriaLabel}
               {...getAnalyticsMetadataAttribute({
                 action: 'selectAll',
               } as Partial<GeneratedAnalyticsMetadataTableSelectAll>)}
-            >
-              {getSelectAllProps ? (
-                <SelectionControl
-                  onFocusDown={event => {
-                    onFocusMove!(event.target as HTMLElement, -1, +1);
-                  }}
-                  focusedComponent={focusedComponent}
-                  {...getSelectAllProps()}
-                  {...(sticky ? { tabIndex: -1 } : {})}
-                />
-              ) : (
-                <ScreenreaderOnly>{singleSelectionHeaderAriaLabel}</ScreenreaderOnly>
-              )}
-              <Divider className={styles['resize-divider']} />
-            </TableThElement>
+            />
           ) : null}
 
           {columnDefinitions.map((column, colIndex) => {
             const columnId = getColumnKey(column, colIndex);
             return (
               <TableHeaderCell
+                {...commonCellProps}
                 key={columnId}
                 style={getColumnStyles(sticky, columnId)}
-                className={headerCellClass}
                 tabIndex={sticky ? -1 : 0}
                 focusedComponent={focusedComponent}
                 column={column}
@@ -158,7 +132,6 @@ const Thead = React.forwardRef(
                 sortingDescending={sortingDescending}
                 sortingDisabled={sortingDisabled}
                 wrapLines={wrapLines}
-                hidden={hidden}
                 colIndex={selectionType ? colIndex + 1 : colIndex}
                 columnId={columnId}
                 updateColumn={updateColumn}
@@ -169,7 +142,6 @@ const Thead = React.forwardRef(
                   fireNonCancelableEvent(onSortingChange, detail);
                 }}
                 isEditable={!!column.editConfig}
-                stickyState={stickyState}
                 cellRef={node => setCell(sticky, columnId, node)}
                 tableRole={tableRole}
                 resizerRoleDescription={resizerRoleDescription}