[EuiBasicTable][EuiInMemoryTable] Support controlled selection API

src/components/basic_table/basic_table.test.tsx

@@ -7,7 +7,7 @@
  */
 
 import React from 'react';
-import { render } from '../../test/rtl';
+import { render, screen } from '../../test/rtl';
 import { requiredProps } from '../../test';
 import { shouldRenderCustomStyles } from '../../test/internal';
 
@@ -391,24 +391,74 @@ describe('EuiBasicTable', () => {
     expect(container.querySelector('[aria-sort]')).toBeTruthy();
   });
 
-  test('with initial selection', () => {
-    const props: EuiBasicTableProps<BasicItem> = {
-      items: basicItems,
-      columns: basicColumns,
-      itemId: 'id',
-      selection: {
-        onSelectionChange: () => {},
-        initialSelected: [basicItems[0]],
-      },
-    };
-    const { getByTestSubject } = render(<EuiBasicTable {...props} />);
+  describe('selection', () => {
+    const getCheckboxAt = (index: number) =>
+      screen.getByTestSubject(`checkboxSelectRow-${index}`) as HTMLInputElement;
 
-    expect(
-      (getByTestSubject('checkboxSelectRow-1') as HTMLInputElement).checked
-    ).toBeTruthy();
-    expect(
-      (getByTestSubject('checkboxSelectRow-2') as HTMLInputElement).checked
-    ).toBeFalsy();
+    test('initialSelected', () => {
+      const props: EuiBasicTableProps<BasicItem> = {
+        items: basicItems,
+        columns: basicColumns,
+        itemId: 'id',
+        selection: {
+          onSelectionChange: () => {},
+          initialSelected: [basicItems[0]],
+        },
+      };
+      render(<EuiBasicTable {...props} />);
+
+      expect(getCheckboxAt(1).checked).toBeTruthy();
+      expect(getCheckboxAt(2).checked).toBeFalsy();
+    });
+
+    test('selected', () => {
+      const props: EuiBasicTableProps<BasicItem> = {
+        items: basicItems,
+        columns: basicColumns,
+        itemId: 'id',
+        selection: {
+          onSelectionChange: () => {},
+          selected: [basicItems[1]],
+        },
+      };
+      render(<EuiBasicTable {...props} />);
+
+      expect(getCheckboxAt(1).checked).toBeFalsy();
+      expect(getCheckboxAt(2).checked).toBeTruthy();
+    });
+
+    it('ignores initialSelected if selected is passed', () => {
+      const props: EuiBasicTableProps<BasicItem> = {
+        items: basicItems,
+        columns: basicColumns,
+        itemId: 'id',
+        selection: {
+          onSelectionChange: () => {},
+          selected: [],
+        },
+      };
+      render(<EuiBasicTable {...props} />);
+
+      expect(getCheckboxAt(1).checked).toBeFalsy();
+      expect(getCheckboxAt(2).checked).toBeFalsy();
+    });
+
+    it("checks for selections that don't exist within `items`", () => {
+      const onSelectionChange = jest.fn();
+      const props: EuiBasicTableProps<BasicItem> = {
+        items: basicItems,
+        columns: basicColumns,
+        itemId: 'id',
+        selection: {
+          onSelectionChange: onSelectionChange,
+          selected: [{ id: 'notvalid', name: 'notInItems' }],
+        },
+      };
+      const { container } = render(<EuiBasicTable {...props} />);
+
+      expect(onSelectionChange).toHaveBeenCalledWith([]);
+      expect(container.querySelectorAll('[checked]')).toHaveLength(0);
+    });
   });
 
   test('footers', () => {

src/components/basic_table/basic_table.tsx

@@ -340,20 +340,28 @@ export class EuiBasicTable<T = any> extends Component<
       return { selection: [] };
     }
 
-    const { itemId } = nextProps;
-    const selection = prevState.selection.filter(
+    const controlledSelection = nextProps.selection.selected;
+    const unfilteredSelection = controlledSelection ?? prevState.selection;
+
+    // Ensure we're not including selections that aren't in the
+    // current `items` array (affected by pagination)
+    const { itemId, items } = nextProps;
+    const selection = unfilteredSelection.filter(
       (selectedItem: T) =>
-        nextProps.items.findIndex(
+        items.findIndex(
           (item: T) =>
             getItemId(item, itemId) === getItemId(selectedItem, itemId)
         ) !== -1
     );
 
-    if (selection.length !== prevState.selection.length) {
-      if (nextProps.selection.onSelectionChange) {
-        nextProps.selection.onSelectionChange(selection);
-      }
+    // If some selected items were filtered out, update state and callback
+    if (selection.length !== unfilteredSelection.length) {
+      nextProps.selection.onSelectionChange?.(selection);
+      return { selection };
+    }
 
+    // Always update selection state from props if controlled
+    if (controlledSelection) {
       return { selection };
     }
 
@@ -385,18 +393,29 @@ export class EuiBasicTable<T = any> extends Component<
     );
   }
 
+  get isSelectionControlled() {
+    return !!this.props.selection?.selected;
+  }
+
   getInitialSelection() {
+    if (this.isSelectionControlled) return;
+
     if (
       this.props.selection &&
       this.props.selection.initialSelected &&
       !this.state.initialSelectionRendered &&
       this.props.items.length > 0
     ) {
-      this.setState({ selection: this.props.selection.initialSelected });
-      this.setState({ initialSelectionRendered: true });
+      this.setState({
+        selection: this.props.selection.initialSelected,
+        initialSelectionRendered: true,
+      });
     }
   }
 
+  /**
+   * @deprecated Use `selection.selected` instead to declaratively control table selection
+   */
   setSelection(newSelection: T[]) {
     this.changeSelection(newSelection);
   }
@@ -415,13 +434,14 @@ export class EuiBasicTable<T = any> extends Component<
     return criteria;
   }
 
-  changeSelection(selection: T[]) {
-    if (!this.props.selection) {
-      return;
-    }
-    this.setState({ selection });
-    if (this.props.selection.onSelectionChange) {
-      this.props.selection.onSelectionChange(selection);
+  changeSelection(changedSelection: T[]) {
+    const { selection } = this.props;
+    if (!selection) return;
+
+    selection.onSelectionChange?.(changedSelection);
+
+    if (!this.isSelectionControlled) {
+      this.setState({ selection: changedSelection });
     }
   }
 

src/components/basic_table/table_types.ts

@@ -175,16 +175,30 @@ export interface EuiTableSortingType<T> {
 
 export interface EuiTableSelectionType<T> {
   /**
-   * A callback that will be called whenever the item selection changes
+   * A callback that will be called whenever the item selection changes.
+   *
+   * Required if `selected` is passed.
    */
   onSelectionChange?: (selection: T[]) => void;
   /**
    * A callback that is called per item to indicate whether it is selectable
    */
   selectable?: (item: T) => boolean;
   /**
-   * A callback that is called per item to retrieve a message for its selectable state.We display these messages as a tooltip on an unselectable checkbox
+   * A callback that is called per item to retrieve a message for its selectable state.
+   * We display these messages as a tooltip on an unselectable checkbox
    */
   selectableMessage?: (selectable: boolean, item: T) => string;
+  /**
+   * Sets initially selected items. Use for uncontrolled selection behavior (checkbox
+   * will only change from user input, and not from developer control).
+   *
+   * This prop will be ignored if `selected` is passed.
+   */
   initialSelected?: T[];
+  /**
+   * Used for controlled selection behavior, e.g. when you want to programmatically
+   * control which selection checkboxes are checked, and which are not.
+   */
+  selected?: T[];
 }