From c06db86ae22f0db1ffb5f8da04172c9aa71c8ecb Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Tue, 11 Jan 2022 10:27:58 -0800
Subject: [PATCH 01/10] [#5517] Add `isFocusedCellInView` state to track
 visibility

- and use this to determine whether or not the entire grid should have a tabindex/be focusable, rather than a simple 'hasFocus' check which doesn't account for virtualization

- Add new mockFocusContext for easier testing for various tests that need to set a focus context

- EuiDataGridCell
  - DRY out an `this.isFocusedCell` method
  - Call `setIsFocusedCellInView` on mount and on unmount (`setFocusedCell` handles setting true/false on new cell focus)
  - Write new unit tests for componentWillUnmount
  - Clean up unnecessary uncovered function fallback for this.unsubscribeCell - we're already checking for a falsy value before calling it
---
 .../datagrid/body/data_grid_cell.test.tsx     | 51 +++++++++++++++----
 .../datagrid/body/data_grid_cell.tsx          | 19 +++++--
 .../data_grid_header_cell_wrapper.test.tsx    | 11 ++--
 src/components/datagrid/data_grid_types.ts    |  1 +
 .../datagrid/utils/__mocks__/focus_context.ts | 14 +++++
 src/components/datagrid/utils/focus.ts        | 16 ++++--
 6 files changed, 86 insertions(+), 26 deletions(-)
 create mode 100644 src/components/datagrid/utils/__mocks__/focus_context.ts

diff --git a/src/components/datagrid/body/data_grid_cell.test.tsx b/src/components/datagrid/body/data_grid_cell.test.tsx
index 5f96eca8111..e74c9efd830 100644
--- a/src/components/datagrid/body/data_grid_cell.test.tsx
+++ b/src/components/datagrid/body/data_grid_cell.test.tsx
@@ -10,6 +10,7 @@ import React from 'react';
 import { mount, ReactWrapper } from 'enzyme';
 import { keys } from '../../../services';
 import { mockRowHeightUtils } from '../utils/__mocks__/row_heights';
+import { mockFocusContext } from '../utils/__mocks__/focus_context';
 import { DataGridFocusContext } from '../utils/focus';
 
 import { EuiDataGridCell } from './data_grid_cell';
@@ -166,27 +167,21 @@ describe('EuiDataGridCell', () => {
   });
 
   describe('componentDidMount', () => {
-    const focusContext = {
-      focusedCell: undefined,
-      onFocusUpdate: jest.fn(),
-      setFocusedCell: jest.fn(),
-    };
-
     it('creates an onFocusUpdate subscription', () => {
       mount(
-        <DataGridFocusContext.Provider value={focusContext}>
+        <DataGridFocusContext.Provider value={mockFocusContext}>
           <EuiDataGridCell {...requiredProps} />
         </DataGridFocusContext.Provider>
       );
 
-      expect(focusContext.onFocusUpdate).toHaveBeenCalled();
+      expect(mockFocusContext.onFocusUpdate).toHaveBeenCalled();
     });
 
     it('mounts the cell with focus state if the current cell should be focused', () => {
       const focusSpy = jest.spyOn(HTMLElement.prototype, 'focus');
       const component = mount(
         <DataGridFocusContext.Provider
-          value={{ ...focusContext, focusedCell: [3, 3] }}
+          value={{ ...mockFocusContext, focusedCell: [3, 3] }}
         >
           <EuiDataGridCell
             {...requiredProps}
@@ -198,6 +193,44 @@ describe('EuiDataGridCell', () => {
 
       expect((component.instance().state as any).isFocused).toEqual(true);
       expect(focusSpy).toHaveBeenCalledWith({ preventScroll: true });
+      expect(mockFocusContext.setIsFocusedCellInView).toHaveBeenCalledWith(
+        true
+      );
+    });
+  });
+
+  describe('componentWillUnmount', () => {
+    it('unsubscribes from the onFocusUpdate subscription', () => {
+      const unsubscribeCellMock = jest.fn();
+      mockFocusContext.onFocusUpdate.mockReturnValueOnce(unsubscribeCellMock);
+
+      const component = mount(
+        <DataGridFocusContext.Provider value={mockFocusContext}>
+          <EuiDataGridCell {...requiredProps} />
+        </DataGridFocusContext.Provider>
+      );
+      component.unmount();
+
+      expect(unsubscribeCellMock).toHaveBeenCalled();
+    });
+
+    it('sets isFocusedCellInView to false if the current cell is focused and unmounting due to being scrolled out of view', () => {
+      const component = mount(
+        <DataGridFocusContext.Provider
+          value={{ ...mockFocusContext, focusedCell: [3, 3] }}
+        >
+          <EuiDataGridCell
+            {...requiredProps}
+            colIndex={3}
+            visibleRowIndex={3}
+          />
+        </DataGridFocusContext.Provider>
+      );
+      component.unmount();
+
+      expect(mockFocusContext.setIsFocusedCellInView).toHaveBeenCalledWith(
+        false
+      );
     });
   });
 
diff --git a/src/components/datagrid/body/data_grid_cell.tsx b/src/components/datagrid/body/data_grid_cell.tsx
index 26ab518a0a4..2f200cf1684 100644
--- a/src/components/datagrid/body/data_grid_cell.tsx
+++ b/src/components/datagrid/body/data_grid_cell.tsx
@@ -124,7 +124,7 @@ export class EuiDataGridCell extends Component<
     enableInteractions: false,
     disableCellTabIndex: false,
   };
-  unsubscribeCell?: Function = () => {};
+  unsubscribeCell?: Function;
   focusTimeout: number | undefined;
   style = null;
 
@@ -234,16 +234,21 @@ export class EuiDataGridCell extends Component<
 
     // Account for virtualization - when a cell unmounts when scrolled out of view
     // and then remounts when scrolled back into view, it should retain focus state
-    if (
-      this.context.focusedCell?.[0] === colIndex &&
-      this.context.focusedCell?.[1] === visibleRowIndex
-    ) {
+    if (this.isFocusedCell()) {
       // The second flag sets preventScroll: true as a focus option, which prevents
       // hijacking the user's scroll behavior when the cell re-mounts on scroll
       this.onFocusUpdate(true, true);
+      this.context.setIsFocusedCellInView(true);
     }
   }
 
+  isFocusedCell = () => {
+    return (
+      this.context.focusedCell?.[0] === this.props.colIndex &&
+      this.context.focusedCell?.[1] === this.props.visibleRowIndex
+    );
+  };
+
   onFocusUpdate = (isFocused: boolean, preventScroll = false) => {
     this.setState({ isFocused }, () => {
       if (isFocused) {
@@ -257,6 +262,10 @@ export class EuiDataGridCell extends Component<
     if (this.unsubscribeCell) {
       this.unsubscribeCell();
     }
+
+    if (this.isFocusedCell()) {
+      this.context.setIsFocusedCellInView(false);
+    }
   }
 
   componentDidUpdate(prevProps: EuiDataGridCellProps) {
diff --git a/src/components/datagrid/body/header/data_grid_header_cell_wrapper.test.tsx b/src/components/datagrid/body/header/data_grid_header_cell_wrapper.test.tsx
index 858bf32c97d..908593766de 100644
--- a/src/components/datagrid/body/header/data_grid_header_cell_wrapper.test.tsx
+++ b/src/components/datagrid/body/header/data_grid_header_cell_wrapper.test.tsx
@@ -12,6 +12,7 @@ import { mount } from 'enzyme';
 
 import { keys } from '../../../../services';
 import { DataGridFocusContext } from '../../utils/focus';
+import { mockFocusContext } from '../../utils/__mocks__/focus_context';
 
 import { EuiDataGridHeaderCellWrapper } from './data_grid_header_cell_wrapper';
 
@@ -23,16 +24,12 @@ describe('EuiDataGridHeaderCellWrapper', () => {
     children: <button />,
   };
 
-  const focusContext = {
-    setFocusedCell: jest.fn(),
-    onFocusUpdate: jest.fn(),
-  };
   const mountWithContext = (props = {}, isFocused = true) => {
-    (focusContext.onFocusUpdate as jest.Mock).mockImplementation(
+    (mockFocusContext.onFocusUpdate as jest.Mock).mockImplementation(
       (_, callback) => callback(isFocused) // allows us to mock isFocused state
     );
     return mount(
-      <DataGridFocusContext.Provider value={focusContext}>
+      <DataGridFocusContext.Provider value={mockFocusContext}>
         <EuiDataGridHeaderCellWrapper {...requiredProps} {...props} />
       </DataGridFocusContext.Provider>
     );
@@ -186,7 +183,7 @@ describe('EuiDataGridHeaderCellWrapper', () => {
               act(() => {
                 headerCell.dispatchEvent(new FocusEvent('focusin'));
               });
-              expect(focusContext.setFocusedCell).toHaveBeenCalled();
+              expect(mockFocusContext.setFocusedCell).toHaveBeenCalled();
             });
 
             it('re-enables and focuses cell interactives when already focused', () => {
diff --git a/src/components/datagrid/data_grid_types.ts b/src/components/datagrid/data_grid_types.ts
index 4e06beba1c0..29a5b03ea06 100644
--- a/src/components/datagrid/data_grid_types.ts
+++ b/src/components/datagrid/data_grid_types.ts
@@ -183,6 +183,7 @@ export interface DataGridSortingContextShape {
 export interface DataGridFocusContextShape {
   focusedCell?: EuiDataGridFocusedCell;
   setFocusedCell: (cell: EuiDataGridFocusedCell) => void;
+  setIsFocusedCellInView: (isFocusedCellInView: boolean) => void;
   onFocusUpdate: (
     cell: EuiDataGridFocusedCell,
     updateFocus: Function
diff --git a/src/components/datagrid/utils/__mocks__/focus_context.ts b/src/components/datagrid/utils/__mocks__/focus_context.ts
new file mode 100644
index 00000000000..41783a079f1
--- /dev/null
+++ b/src/components/datagrid/utils/__mocks__/focus_context.ts
@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+export const mockFocusContext = {
+  focusedCell: undefined,
+  onFocusUpdate: jest.fn(),
+  setFocusedCell: jest.fn(),
+  setIsFocusedCellInView: jest.fn(),
+};
diff --git a/src/components/datagrid/utils/focus.ts b/src/components/datagrid/utils/focus.ts
index 1fd7b6c6b98..c8b54d8e532 100644
--- a/src/components/datagrid/utils/focus.ts
+++ b/src/components/datagrid/utils/focus.ts
@@ -28,6 +28,7 @@ import {
 export const DataGridFocusContext = createContext<DataGridFocusContextShape>({
   focusedCell: undefined,
   setFocusedCell: () => {},
+  setIsFocusedCellInView: () => {},
   onFocusUpdate: () => () => {},
 });
 
@@ -54,10 +55,16 @@ export const useFocus = (
   );
 
   // Current focused cell
-  const [focusedCell, setFocusedCell] = useState<
+  const [isFocusedCellInView, setIsFocusedCellInView] = useState(false);
+  const [focusedCell, _setFocusedCell] = useState<
     EuiDataGridFocusedCell | undefined
   >(undefined);
 
+  const setFocusedCell = useCallback((focusedCell: EuiDataGridFocusedCell) => {
+    _setFocusedCell(focusedCell);
+    setIsFocusedCellInView(true); // The cell must be in view to focus it
+  }, []);
+
   const previousCell = useRef<EuiDataGridFocusedCell | undefined>(undefined);
   useEffect(() => {
     if (previousCell.current) {
@@ -74,11 +81,9 @@ export const useFocus = (
     }
   }, [cellsUpdateFocus, focusedCell]);
 
-  const hasHadFocus = useMemo(() => focusedCell != null, [focusedCell]);
-
   const focusProps = useMemo<FocusProps>(
     () =>
-      hasHadFocus
+      isFocusedCellInView
         ? {
             // FireFox allows tabbing to a div that is scrollable, while Chrome does not
             tabIndex: -1,
@@ -95,13 +100,14 @@ export const useFocus = (
               }
             },
           },
-    [hasHadFocus, setFocusedCell, headerIsInteractive]
+    [isFocusedCellInView, setFocusedCell, headerIsInteractive]
   );
 
   return {
     onFocusUpdate,
     focusedCell,
     setFocusedCell,
+    setIsFocusedCellInView,
     focusProps,
   };
 };

From 7f031ac496e7cb11337457cac5e524310b6ec886 Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Tue, 11 Jan 2022 11:47:48 -0800
Subject: [PATCH 02/10] [REVERT] Update demo to test noninteractive headers
 with scrolling content

- see http://localhost:8030/#/tabular-content/data-grid-focus
---
 src-docs/src/views/datagrid/focus.js | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src-docs/src/views/datagrid/focus.js b/src-docs/src/views/datagrid/focus.js
index 5637812fc1c..1c57493ae9d 100644
--- a/src-docs/src/views/datagrid/focus.js
+++ b/src-docs/src/views/datagrid/focus.js
@@ -17,7 +17,7 @@ import {
 
 const data = [];
 
-for (let i = 0; i < 10; i++) {
+for (let i = 0; i < 25; i++) {
   data.push([
     <span>{fake('{{name.firstName}}')}</span>,
     <span>{fake('{{name.firstName}}')}</span>,
@@ -214,7 +214,7 @@ export default () => {
   );
 
   const [pagination, setPagination] = useState({
-    pageSize: 4,
+    pageSize: 20,
     pageIndex: 0,
     pageSizeOptions: [4],
   });
@@ -250,6 +250,7 @@ export default () => {
           onChangeItemsPerPage,
           onChangePage,
         }}
+        height={300}
       />
     </>
   );

From 2f44ab92f474c1f205757d0e86c8e7132ca36610 Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Tue, 11 Jan 2022 11:45:03 -0800
Subject: [PATCH 03/10] [#5517] Add `onItemsRendered` ref to store currently
 visible/virtualized refs

- see https://react-window.vercel.app/#/api/FixedSizeGrid for `onItemsRendered` API
- this struck me as the quickest/easiest way to determine which virtualized cells are in view
- I opted to save this as a ref instead of as state as I didn't see the need to cause a rerender

- `focusFirstVisibleInteractiveCell` was split out to its own fn as it will shortly be used elsewhere to fix another focus bug in the grid
---
 .../datagrid/body/data_grid_body.test.tsx     |  1 +
 .../datagrid/body/data_grid_body.tsx          |  4 +++
 src/components/datagrid/data_grid.tsx         | 14 ++++++---
 src/components/datagrid/data_grid_types.ts    |  7 ++++-
 src/components/datagrid/utils/focus.ts        | 31 +++++++++++++++----
 5 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/src/components/datagrid/body/data_grid_body.test.tsx b/src/components/datagrid/body/data_grid_body.test.tsx
index 4a0a84a341c..a6fe83ddf1a 100644
--- a/src/components/datagrid/body/data_grid_body.test.tsx
+++ b/src/components/datagrid/body/data_grid_body.test.tsx
@@ -52,6 +52,7 @@ describe('EuiDataGridBody', () => {
         resetAfterRowIndex: jest.fn(),
       } as any,
     },
+    gridItemsRendered: {} as any,
     wrapperRef: { current: document.createElement('div') },
   };
 
diff --git a/src/components/datagrid/body/data_grid_body.tsx b/src/components/datagrid/body/data_grid_body.tsx
index 15c6488cf79..d640ddb1fa0 100644
--- a/src/components/datagrid/body/data_grid_body.tsx
+++ b/src/components/datagrid/body/data_grid_body.tsx
@@ -243,6 +243,7 @@ export const EuiDataGridBody: FunctionComponent<EuiDataGridBodyProps> = (
     gridStyles,
     gridWidth,
     gridRef,
+    gridItemsRendered,
     wrapperRef,
   } = props;
 
@@ -442,6 +443,9 @@ export const EuiDataGridBody: FunctionComponent<EuiDataGridBodyProps> = (
       <Grid
         {...(virtualizationOptions ? virtualizationOptions : {})}
         ref={gridRef}
+        onItemsRendered={(itemsRendered) => {
+          gridItemsRendered.current = itemsRendered;
+        }}
         innerElementType={InnerElement}
         outerRef={outerGridRef}
         innerRef={innerGridRef}
diff --git a/src/components/datagrid/data_grid.tsx b/src/components/datagrid/data_grid.tsx
index 46096ad9aaa..8dfb720bc80 100644
--- a/src/components/datagrid/data_grid.tsx
+++ b/src/components/datagrid/data_grid.tsx
@@ -14,7 +14,10 @@ import React, {
   useRef,
   useState,
 } from 'react';
-import { VariableSizeGrid as Grid } from 'react-window';
+import {
+  VariableSizeGrid as Grid,
+  GridOnItemsRenderedProps,
+} from 'react-window';
 import { useGeneratedHtmlId, keys } from '../../services';
 import { EuiFocusTrap } from '../focus_trap';
 import { EuiI18n, useEuiI18n } from '../i18n';
@@ -169,6 +172,7 @@ export const EuiDataGrid: FunctionComponent<EuiDataGridProps> = (props) => {
   // Imperative handler passed back by react-window - we're setting this at
   // the top datagrid level to make passing it to other children & utils easier
   const gridRef = useRef<Grid | null>(null);
+  const gridItemsRendered = useRef<GridOnItemsRenderedProps | null>(null);
 
   /**
    * Display
@@ -254,9 +258,10 @@ export const EuiDataGrid: FunctionComponent<EuiDataGridProps> = (props) => {
   const { headerIsInteractive, handleHeaderMutation } = useHeaderIsInteractive(
     contentRef.current
   );
-  const { focusProps: wrappingDivFocusProps, ...focusContext } = useFocus(
-    headerIsInteractive
-  );
+  const { focusProps: wrappingDivFocusProps, ...focusContext } = useFocus({
+    headerIsInteractive,
+    gridItemsRendered,
+  });
 
   /**
    * Toolbar & full-screen
@@ -437,6 +442,7 @@ export const EuiDataGrid: FunctionComponent<EuiDataGridProps> = (props) => {
                 gridStyles={gridStyles}
                 gridWidth={gridWidth}
                 gridRef={gridRef}
+                gridItemsRendered={gridItemsRendered}
                 wrapperRef={contentRef}
               />
             </div>
diff --git a/src/components/datagrid/data_grid_types.ts b/src/components/datagrid/data_grid_types.ts
index 29a5b03ea06..6162b44ba78 100644
--- a/src/components/datagrid/data_grid_types.ts
+++ b/src/components/datagrid/data_grid_types.ts
@@ -19,7 +19,11 @@ import {
   SetStateAction,
   MutableRefObject,
 } from 'react';
-import { VariableSizeGridProps, VariableSizeGrid as Grid } from 'react-window';
+import {
+  VariableSizeGridProps,
+  VariableSizeGrid as Grid,
+  GridOnItemsRenderedProps,
+} from 'react-window';
 import { EuiListGroupItemProps } from '../list_group';
 import { EuiButtonEmpty, EuiButtonIcon } from '../button';
 import { ExclusiveUnion, CommonProps, OneOf } from '../common';
@@ -347,6 +351,7 @@ export interface EuiDataGridBodyProps {
   gridStyles: EuiDataGridStyle;
   gridWidth: number;
   gridRef: MutableRefObject<Grid | null>;
+  gridItemsRendered: MutableRefObject<GridOnItemsRenderedProps | null>;
   wrapperRef: MutableRefObject<HTMLDivElement | null>;
 }
 export interface EuiDataGridCellValueElementProps {
diff --git a/src/components/datagrid/utils/focus.ts b/src/components/datagrid/utils/focus.ts
index c8b54d8e532..d570b5679a1 100644
--- a/src/components/datagrid/utils/focus.ts
+++ b/src/components/datagrid/utils/focus.ts
@@ -16,7 +16,9 @@ import {
   useState,
   HTMLAttributes,
   KeyboardEvent,
+  MutableRefObject,
 } from 'react';
+import { GridOnItemsRenderedProps } from 'react-window';
 import tabbable from 'tabbable';
 import { keys } from '../../../services';
 import {
@@ -37,9 +39,13 @@ type FocusProps = Pick<HTMLAttributes<HTMLDivElement>, 'tabIndex' | 'onFocus'>;
 /**
  * Main focus context and overarching focus state management
  */
-export const useFocus = (
-  headerIsInteractive: boolean
-): DataGridFocusContextShape & { focusProps: FocusProps } => {
+export const useFocus = ({
+  headerIsInteractive,
+  gridItemsRendered,
+}: {
+  headerIsInteractive: boolean;
+  gridItemsRendered: MutableRefObject<GridOnItemsRenderedProps | null>;
+}): DataGridFocusContextShape & { focusProps: FocusProps } => {
   // Maintain a map of focus cell state callbacks
   const cellsUpdateFocus = useRef<Map<string, Function>>(new Map());
 
@@ -81,6 +87,19 @@ export const useFocus = (
     }
   }, [cellsUpdateFocus, focusedCell]);
 
+  const focusFirstVisibleInteractiveCell = useCallback(() => {
+    const {
+      visibleColumnStartIndex,
+      visibleRowStartIndex,
+    } = gridItemsRendered.current!;
+
+    setFocusedCell(
+      headerIsInteractive
+        ? [0, -1]
+        : [visibleColumnStartIndex, visibleRowStartIndex]
+    );
+  }, [setFocusedCell, headerIsInteractive, gridItemsRendered]);
+
   const focusProps = useMemo<FocusProps>(
     () =>
       isFocusedCellInView
@@ -94,13 +113,13 @@ export const useFocus = (
               // if e.target (the source element of the `focus event`
               // matches e.currentTarget (always the div with this onFocus listener)
               // then the user has focused directly on the data grid wrapper (almost definitely by tabbing)
-              // so shift focus to the first interactive cell within the grid
+              // so shift focus to the first visible and interactive cell within the grid
               if (e.target === e.currentTarget) {
-                setFocusedCell(headerIsInteractive ? [0, -1] : [0, 0]);
+                focusFirstVisibleInteractiveCell();
               }
             },
           },
-    [isFocusedCellInView, setFocusedCell, headerIsInteractive]
+    [isFocusedCellInView, focusFirstVisibleInteractiveCell]
   );
 
   return {

From d29d9381fe31366d6edbd6795e91b93e9fcc3f19 Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Tue, 11 Jan 2022 11:59:59 -0800
Subject: [PATCH 04/10] [#4923] Fix focus returning to document body when
 hiding a column via header

---
 .../datagrid/body/header/column_actions.test.tsx  |  5 ++++-
 .../datagrid/body/header/column_actions.tsx       | 15 +++++++++++++--
 .../body/header/data_grid_header_cell.tsx         |  4 ++++
 src/components/datagrid/data_grid_types.ts        |  1 +
 .../datagrid/utils/__mocks__/focus_context.ts     |  1 +
 src/components/datagrid/utils/focus.ts            |  2 ++
 6 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/src/components/datagrid/body/header/column_actions.test.tsx b/src/components/datagrid/body/header/column_actions.test.tsx
index 1a0d6e740d9..058353544bb 100644
--- a/src/components/datagrid/body/header/column_actions.test.tsx
+++ b/src/components/datagrid/body/header/column_actions.test.tsx
@@ -19,6 +19,7 @@ import {
 
 describe('getColumnActions', () => {
   const setVisibleColumns = jest.fn();
+  const focusFirstVisibleInteractiveCell = jest.fn();
   const setIsPopoverOpen = jest.fn();
   const switchColumnPos = jest.fn();
 
@@ -28,6 +29,7 @@ describe('getColumnActions', () => {
     schema: {},
     schemaDetectors,
     setVisibleColumns,
+    focusFirstVisibleInteractiveCell,
     setIsPopoverOpen,
     sorting: undefined,
     switchColumnPos,
@@ -104,9 +106,10 @@ describe('getColumnActions', () => {
           `);
       });
 
-      it('sets column visibility on click', () => {
+      it('hides the current column on click and refocuses into the grid', () => {
         callActionOnClick(hideColumn);
         expect(setVisibleColumns).toHaveBeenCalledWith(['A', 'C']);
+        expect(focusFirstVisibleInteractiveCell).toHaveBeenCalled();
       });
     });
 
diff --git a/src/components/datagrid/body/header/column_actions.tsx b/src/components/datagrid/body/header/column_actions.tsx
index fcfc8601bc0..2a557bdcd0a 100644
--- a/src/components/datagrid/body/header/column_actions.tsx
+++ b/src/components/datagrid/body/header/column_actions.tsx
@@ -13,6 +13,7 @@ import {
   EuiDataGridSchema,
   EuiDataGridSchemaDetector,
   EuiDataGridSorting,
+  DataGridFocusContextShape,
 } from '../../data_grid_types';
 import { EuiI18n } from '../../../i18n';
 import { EuiListGroupItemProps } from '../../../list_group';
@@ -28,6 +29,7 @@ interface GetColumnActions {
   schema: EuiDataGridSchema;
   schemaDetectors: EuiDataGridSchemaDetector[];
   setVisibleColumns: (columnId: string[]) => void;
+  focusFirstVisibleInteractiveCell: DataGridFocusContextShape['focusFirstVisibleInteractiveCell'];
   setIsPopoverOpen: (value: boolean) => void;
   sorting: EuiDataGridSorting | undefined;
   switchColumnPos: (colFromId: string, colToId: string) => void;
@@ -39,6 +41,7 @@ export const getColumnActions = ({
   schema,
   schemaDetectors,
   setVisibleColumns,
+  focusFirstVisibleInteractiveCell,
   setIsPopoverOpen,
   sorting,
   switchColumnPos,
@@ -52,6 +55,7 @@ export const getColumnActions = ({
       column,
       columns,
       setVisibleColumns,
+      focusFirstVisibleInteractiveCell,
     }),
     ...getSortColumnActions({
       column,
@@ -85,20 +89,27 @@ export const getColumnActions = ({
  */
 type HideColumnAction = Pick<
   GetColumnActions,
-  'column' | 'columns' | 'setVisibleColumns'
+  | 'column'
+  | 'columns'
+  | 'setVisibleColumns'
+  | 'focusFirstVisibleInteractiveCell'
 >;
 
 export const getHideColumnAction = ({
   column,
   columns,
   setVisibleColumns,
+  focusFirstVisibleInteractiveCell,
 }: HideColumnAction): EuiListGroupItemProps[] => {
   const items = [];
 
-  const onClickHideColumn = () =>
+  const onClickHideColumn = () => {
     setVisibleColumns(
       columns.filter((col) => col.id !== column.id).map((col) => col.id)
     );
+    // Since we hid the current column, we need to manually set focus back onto the grid
+    focusFirstVisibleInteractiveCell();
+  };
 
   const action = {
     label: (
diff --git a/src/components/datagrid/body/header/data_grid_header_cell.tsx b/src/components/datagrid/body/header/data_grid_header_cell.tsx
index 5c9f937ad6a..d7cbb919618 100644
--- a/src/components/datagrid/body/header/data_grid_header_cell.tsx
+++ b/src/components/datagrid/body/header/data_grid_header_cell.tsx
@@ -21,6 +21,7 @@ import { EuiIcon } from '../../../icon';
 import { EuiListGroup } from '../../../list_group';
 import { EuiPopover } from '../../../popover';
 import { DataGridSortingContext } from '../../utils/sorting';
+import { DataGridFocusContext } from '../../utils/focus';
 import { EuiDataGridHeaderCellProps } from '../../data_grid_types';
 
 import { getColumnActions } from './column_actions';
@@ -58,6 +59,8 @@ export const EuiDataGridHeaderCell: FunctionComponent<EuiDataGridHeaderCellProps
   } = {};
   const screenReaderId = useGeneratedHtmlId();
 
+  const { focusFirstVisibleInteractiveCell } = useContext(DataGridFocusContext);
+
   const { sorting } = useContext(DataGridSortingContext);
   let sortString;
   if (sorting) {
@@ -92,6 +95,7 @@ export const EuiDataGridHeaderCell: FunctionComponent<EuiDataGridHeaderCellProps
     schema,
     schemaDetectors,
     setVisibleColumns,
+    focusFirstVisibleInteractiveCell,
     setIsPopoverOpen,
     sorting,
     switchColumnPos,
diff --git a/src/components/datagrid/data_grid_types.ts b/src/components/datagrid/data_grid_types.ts
index 6162b44ba78..9100ab2fc93 100644
--- a/src/components/datagrid/data_grid_types.ts
+++ b/src/components/datagrid/data_grid_types.ts
@@ -192,6 +192,7 @@ export interface DataGridFocusContextShape {
     cell: EuiDataGridFocusedCell,
     updateFocus: Function
   ) => () => void;
+  focusFirstVisibleInteractiveCell: () => void;
 }
 
 export type CommonGridProps = CommonProps &
diff --git a/src/components/datagrid/utils/__mocks__/focus_context.ts b/src/components/datagrid/utils/__mocks__/focus_context.ts
index 41783a079f1..65b682a18bd 100644
--- a/src/components/datagrid/utils/__mocks__/focus_context.ts
+++ b/src/components/datagrid/utils/__mocks__/focus_context.ts
@@ -11,4 +11,5 @@ export const mockFocusContext = {
   onFocusUpdate: jest.fn(),
   setFocusedCell: jest.fn(),
   setIsFocusedCellInView: jest.fn(),
+  focusFirstVisibleInteractiveCell: jest.fn(),
 };
diff --git a/src/components/datagrid/utils/focus.ts b/src/components/datagrid/utils/focus.ts
index d570b5679a1..c204525e1c8 100644
--- a/src/components/datagrid/utils/focus.ts
+++ b/src/components/datagrid/utils/focus.ts
@@ -32,6 +32,7 @@ export const DataGridFocusContext = createContext<DataGridFocusContextShape>({
   setFocusedCell: () => {},
   setIsFocusedCellInView: () => {},
   onFocusUpdate: () => () => {},
+  focusFirstVisibleInteractiveCell: () => {},
 });
 
 type FocusProps = Pick<HTMLAttributes<HTMLDivElement>, 'tabIndex' | 'onFocus'>;
@@ -127,6 +128,7 @@ export const useFocus = ({
     focusedCell,
     setFocusedCell,
     setIsFocusedCellInView,
+    focusFirstVisibleInteractiveCell,
     focusProps,
   };
 };

From 8a035de8d4ddb75d45feaacab181f55b1d7a57d7 Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Tue, 11 Jan 2022 12:05:45 -0800
Subject: [PATCH 05/10] [#5476] Fix keyboard navigation after hiding a column
 via header actions

---
 .../datagrid/body/header/column_actions.test.tsx          | 6 +++++-
 src/components/datagrid/body/header/column_actions.tsx    | 8 +++++++-
 .../datagrid/body/header/data_grid_header_cell.tsx        | 5 ++++-
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/components/datagrid/body/header/column_actions.test.tsx b/src/components/datagrid/body/header/column_actions.test.tsx
index 058353544bb..15b39f51e63 100644
--- a/src/components/datagrid/body/header/column_actions.test.tsx
+++ b/src/components/datagrid/body/header/column_actions.test.tsx
@@ -22,6 +22,7 @@ describe('getColumnActions', () => {
   const focusFirstVisibleInteractiveCell = jest.fn();
   const setIsPopoverOpen = jest.fn();
   const switchColumnPos = jest.fn();
+  const setFocusedCell = jest.fn();
 
   const testArgs = {
     column: { id: 'B' },
@@ -33,6 +34,7 @@ describe('getColumnActions', () => {
     setIsPopoverOpen,
     sorting: undefined,
     switchColumnPos,
+    setFocusedCell,
   };
 
   // DRY test helper
@@ -185,12 +187,14 @@ describe('getColumnActions', () => {
           `);
       });
 
-      it('calls switchColumnPos on click', () => {
+      it('calls switchColumnPos and updates the focused cell column index on click', () => {
         callActionOnClick(moveLeft);
         expect(switchColumnPos).toHaveBeenCalledWith('B', 'A');
+        expect(setFocusedCell).toHaveBeenLastCalledWith([0, -1]);
 
         callActionOnClick(moveRight);
         expect(switchColumnPos).toHaveBeenCalledWith('B', 'C');
+        expect(setFocusedCell).toHaveBeenLastCalledWith([2, -1]);
       });
     });
 
diff --git a/src/components/datagrid/body/header/column_actions.tsx b/src/components/datagrid/body/header/column_actions.tsx
index 2a557bdcd0a..3f594ac0703 100644
--- a/src/components/datagrid/body/header/column_actions.tsx
+++ b/src/components/datagrid/body/header/column_actions.tsx
@@ -33,6 +33,7 @@ interface GetColumnActions {
   setIsPopoverOpen: (value: boolean) => void;
   sorting: EuiDataGridSorting | undefined;
   switchColumnPos: (colFromId: string, colToId: string) => void;
+  setFocusedCell: DataGridFocusContextShape['setFocusedCell'];
 }
 
 export const getColumnActions = ({
@@ -45,6 +46,7 @@ export const getColumnActions = ({
   setIsPopoverOpen,
   sorting,
   switchColumnPos,
+  setFocusedCell,
 }: GetColumnActions): EuiListGroupItemProps[] => {
   if (column.actions === false) {
     return [];
@@ -67,6 +69,7 @@ export const getColumnActions = ({
       column,
       columns,
       switchColumnPos,
+      setFocusedCell,
     }),
     ...(column.actions?.additional || []),
   ];
@@ -133,13 +136,14 @@ export const getHideColumnAction = ({
  */
 type MoveColumnActions = Pick<
   GetColumnActions,
-  'column' | 'columns' | 'switchColumnPos'
+  'column' | 'columns' | 'switchColumnPos' | 'setFocusedCell'
 >;
 
 const getMoveColumnActions = ({
   column,
   columns,
   switchColumnPos,
+  setFocusedCell,
 }: MoveColumnActions): EuiListGroupItemProps[] => {
   const items = [];
 
@@ -150,6 +154,7 @@ const getMoveColumnActions = ({
       const targetCol = columns[colIdx - 1];
       if (targetCol) {
         switchColumnPos(column.id, targetCol.id);
+        setFocusedCell([colIdx - 1, -1]);
       }
     };
     const action = {
@@ -169,6 +174,7 @@ const getMoveColumnActions = ({
       const targetCol = columns[colIdx + 1];
       if (targetCol) {
         switchColumnPos(column.id, targetCol.id);
+        setFocusedCell([colIdx + 1, -1]);
       }
     };
     const action = {
diff --git a/src/components/datagrid/body/header/data_grid_header_cell.tsx b/src/components/datagrid/body/header/data_grid_header_cell.tsx
index d7cbb919618..880506e00df 100644
--- a/src/components/datagrid/body/header/data_grid_header_cell.tsx
+++ b/src/components/datagrid/body/header/data_grid_header_cell.tsx
@@ -59,7 +59,9 @@ export const EuiDataGridHeaderCell: FunctionComponent<EuiDataGridHeaderCellProps
   } = {};
   const screenReaderId = useGeneratedHtmlId();
 
-  const { focusFirstVisibleInteractiveCell } = useContext(DataGridFocusContext);
+  const { setFocusedCell, focusFirstVisibleInteractiveCell } = useContext(
+    DataGridFocusContext
+  );
 
   const { sorting } = useContext(DataGridSortingContext);
   let sortString;
@@ -99,6 +101,7 @@ export const EuiDataGridHeaderCell: FunctionComponent<EuiDataGridHeaderCellProps
     setIsPopoverOpen,
     sorting,
     switchColumnPos,
+    setFocusedCell,
   });
 
   const showColumnActions = columnActions && columnActions.length > 0;

From 129febdd3880d93cbe2ce162c648cccd03a1f975 Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Tue, 11 Jan 2022 12:30:09 -0800
Subject: [PATCH 06/10] Add changelog entry

---
 CHANGELOG.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 47d55dd0fe8..9c518d9ace1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 ## [`main`](https://github.com/elastic/eui/tree/main)
 
+**Bug fixes**
+
+- Fixed multiple bugs with `EuiDataGrid` keyboard focus restoration ([#5530](https://github.com/elastic/eui/pull/5530))
+
 **Breaking changes**
 
 - Changed `EuiSearchBar` to preserve phrases with leading and trailing spaces, instead of dropping surrounding whitespace ([#5514](https://github.com/elastic/eui/pull/5514))

From 0fce5495504ebe6ffb686fe180b617642c382fd3 Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Tue, 11 Jan 2022 13:17:37 -0800
Subject: [PATCH 07/10] Revert "[REVERT] Update demo to test noninteractive
 headers with scrolling content"

This reverts commit 7f031ac496e7cb11337457cac5e524310b6ec886.
---
 src-docs/src/views/datagrid/focus.js | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src-docs/src/views/datagrid/focus.js b/src-docs/src/views/datagrid/focus.js
index 1c57493ae9d..5637812fc1c 100644
--- a/src-docs/src/views/datagrid/focus.js
+++ b/src-docs/src/views/datagrid/focus.js
@@ -17,7 +17,7 @@ import {
 
 const data = [];
 
-for (let i = 0; i < 25; i++) {
+for (let i = 0; i < 10; i++) {
   data.push([
     <span>{fake('{{name.firstName}}')}</span>,
     <span>{fake('{{name.firstName}}')}</span>,
@@ -214,7 +214,7 @@ export default () => {
   );
 
   const [pagination, setPagination] = useState({
-    pageSize: 20,
+    pageSize: 4,
     pageIndex: 0,
     pageSizeOptions: [4],
   });
@@ -250,7 +250,6 @@ export default () => {
           onChangeItemsPerPage,
           onChangePage,
         }}
-        height={300}
       />
     </>
   );

From 4fad086399850096a5b4b9c6f55bbac9c5d722f0 Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Tue, 11 Jan 2022 13:23:28 -0800
Subject: [PATCH 08/10] [PR feedback] Add explanatory comment in unit test

---
 src/components/datagrid/body/header/column_actions.test.tsx | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/components/datagrid/body/header/column_actions.test.tsx b/src/components/datagrid/body/header/column_actions.test.tsx
index 15b39f51e63..ce8e7e9797a 100644
--- a/src/components/datagrid/body/header/column_actions.test.tsx
+++ b/src/components/datagrid/body/header/column_actions.test.tsx
@@ -195,6 +195,10 @@ describe('getColumnActions', () => {
         callActionOnClick(moveRight);
         expect(switchColumnPos).toHaveBeenCalledWith('B', 'C');
         expect(setFocusedCell).toHaveBeenLastCalledWith([2, -1]);
+
+        // Quick note about the -1 row index above: `-1` is the index we set for our
+        // sticky header row, and these column actions can only ever be called by an
+        // interactive header cell, so we can safely and statically set the -1 index
       });
     });
 

From 9a6953487d03398cc579ad5baf2bfe9f56fc43d6 Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Tue, 11 Jan 2022 14:38:09 -0800
Subject: [PATCH 09/10] [PR feedback] Account for empty grids

- gridItemsRendered will be falsey if rowCount is 0 and will cause an error on tab, so we should opt to simply leave focus on the grid body wrapper
---
 src/components/datagrid/utils/focus.ts | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/components/datagrid/utils/focus.ts b/src/components/datagrid/utils/focus.ts
index c204525e1c8..09cf768a41d 100644
--- a/src/components/datagrid/utils/focus.ts
+++ b/src/components/datagrid/utils/focus.ts
@@ -89,16 +89,20 @@ export const useFocus = ({
   }, [cellsUpdateFocus, focusedCell]);
 
   const focusFirstVisibleInteractiveCell = useCallback(() => {
-    const {
-      visibleColumnStartIndex,
-      visibleRowStartIndex,
-    } = gridItemsRendered.current!;
+    if (headerIsInteractive) {
+      // The header (rowIndex -1) is sticky and will always be in view
+      setFocusedCell([0, -1]);
+    } else if (gridItemsRendered.current) {
+      const {
+        visibleColumnStartIndex,
+        visibleRowStartIndex,
+      } = gridItemsRendered.current;
 
-    setFocusedCell(
-      headerIsInteractive
-        ? [0, -1]
-        : [visibleColumnStartIndex, visibleRowStartIndex]
-    );
+      setFocusedCell([visibleColumnStartIndex, visibleRowStartIndex]);
+    } else {
+      // If the header is non-interactive and there are no rendered cells,
+      // there's nothing to do - we might as well leave focus on the grid body wrapper
+    }
   }, [setFocusedCell, headerIsInteractive, gridItemsRendered]);
 
   const focusProps = useMemo<FocusProps>(

From 694621b5e2efb69dce475ddb0bed36f2e2a3e40b Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Wed, 12 Jan 2022 09:16:08 -0800
Subject: [PATCH 10/10] [bug] Handle arrow keydown behavior on sticky header
 when row 0 is not in view

---
 src/components/datagrid/data_grid.tsx  |  2 ++
 src/components/datagrid/utils/focus.ts | 11 ++++++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/components/datagrid/data_grid.tsx b/src/components/datagrid/data_grid.tsx
index 8dfb720bc80..09b7c82fef6 100644
--- a/src/components/datagrid/data_grid.tsx
+++ b/src/components/datagrid/data_grid.tsx
@@ -403,6 +403,8 @@ export const EuiDataGrid: FunctionComponent<EuiDataGridProps> = (props) => {
                 gridElement: contentRef.current,
                 visibleColCount,
                 visibleRowCount,
+                visibleRowStartIndex:
+                  gridItemsRendered.current?.visibleRowStartIndex || 0,
                 rowCount,
                 pagination,
                 hasFooter: !!renderFooterCellValue,
diff --git a/src/components/datagrid/utils/focus.ts b/src/components/datagrid/utils/focus.ts
index 09cf768a41d..f42af4ebf9d 100644
--- a/src/components/datagrid/utils/focus.ts
+++ b/src/components/datagrid/utils/focus.ts
@@ -156,6 +156,7 @@ export const createKeyDownHandler = ({
   gridElement,
   visibleColCount,
   visibleRowCount,
+  visibleRowStartIndex,
   rowCount,
   pagination,
   hasFooter,
@@ -165,6 +166,7 @@ export const createKeyDownHandler = ({
   gridElement: HTMLDivElement | null;
   visibleColCount: number;
   visibleRowCount: number;
+  visibleRowStartIndex: number;
   rowCount: EuiDataGridProps['rowCount'];
   pagination: EuiDataGridProps['pagination'];
   hasFooter: boolean;
@@ -187,7 +189,14 @@ export const createKeyDownHandler = ({
     if (key === keys.ARROW_DOWN) {
       event.preventDefault();
       if (hasFooter ? y < visibleRowCount : y < visibleRowCount - 1) {
-        setFocusedCell([x, y + 1]);
+        if (y === -1) {
+          // The header is sticky, so on scrolling virtualized grids, row 0 will not
+          // always be rendered to navigate down to. We need to account for this by
+          // sending the down arrow to the first visible/virtualized row instead
+          setFocusedCell([x, visibleRowStartIndex]);
+        } else {
+          setFocusedCell([x, y + 1]);
+        }
       }
     } else if (key === keys.ARROW_LEFT) {
       event.preventDefault();