From cb844097017f7304bcb7af6fe53fba8a5b9fb2ab Mon Sep 17 00:00:00 2001
From: Sychev Andrey <44414066+SychevAndrey@users.noreply.github.com>
Date: Fri, 27 Jun 2025 12:23:37 +0200
Subject: [PATCH] dbeaver/pro#5998 Fix:  can't copy the whole table selected
 with the corner cell (#3550)

* dbeaver/pro#5998 fix: selectAll index header cell

* dbeaver/pro#5998 fix: removed OrderButton from tabbing order, add sorting on space/enter when header column cell is in focus

* dbeaver/pro#5998 feat: add onHeaderKeyDown handler to DataGrid and HeaderCellContentRenderer

we lost ability to copy data grid cells' data when one of the header cells is focused, because the header cell renderer was not handling keydown events. This commit utilizing gridSelectedCellCopy.onKeydownHandler inside HeaderCellContentRenderer keydown event handler

* dbeaver/pro#5998 style: change onKeyDown to a named function

* dbeaver/pro#5998 refactor: simplify OrderButton and update sorting logic in HeaderCellContentRenderer

* dbeaver/pro#5998 refactor: rename container class to clickAreaOverlay in TableIndexColumnHeader, remove z-index

---------

Co-authored-by: Daria Marutkina <125263541+dariamarutkina@users.noreply.github.com>
---
 .../@dbeaver/react-data-grid/src/DataGrid.tsx |  3 +-
 .../src/DataGridHeaderCellContext.ts          |  1 +
 .../renderers/HeaderCellContentRenderer.tsx   | 33 +++++++++++--------
 .../src/renderers/OrderButton.tsx             | 14 ++------
 .../src/DataGrid/DataGridTable.tsx            | 14 +++++---
 .../TableIndexColumnHeader.module.css         |  9 ++++-
 .../TableIndexColumnHeader.tsx                | 10 ++++--
 7 files changed, 50 insertions(+), 34 deletions(-)

diff --git a/common-react/@dbeaver/react-data-grid/src/DataGrid.tsx b/common-react/@dbeaver/react-data-grid/src/DataGrid.tsx
index 26f8eadaf6..03afecd9ff 100644
--- a/common-react/@dbeaver/react-data-grid/src/DataGrid.tsx
+++ b/common-react/@dbeaver/react-data-grid/src/DataGrid.tsx
@@ -70,6 +70,7 @@ export const DataGrid = forwardRef<DataGridRef, DataGridProps>(function DataGrid
     onFocus,
     onCellChange,
     onColumnSort,
+    onHeaderKeyDown,
     children,
     className,
     onCellKeyDown,
@@ -151,7 +152,7 @@ export const DataGrid = forwardRef<DataGridRef, DataGridProps>(function DataGrid
     <HeaderDnDContext value={dndHeaderContext}>
       <DataGridRowContext value={{ rowCount, onScrollToBottom }}>
         <DataGridCellContext value={{ cell, cellText, cellElement, cellTooltip, onCellChange }}>
-          <DataGridCellHeaderContext value={{ headerElement, headerText, getHeaderDnD, columnSortable, onColumnSort, columnSortingState }}>
+          <DataGridCellHeaderContext value={{ headerElement, headerText, getHeaderDnD, columnSortable, onColumnSort, columnSortingState, onHeaderKeyDown }}>
             <DataGridBase
               ref={innerGridRef}
               columns={dndHeaderContext.columns}
diff --git a/common-react/@dbeaver/react-data-grid/src/DataGridHeaderCellContext.ts b/common-react/@dbeaver/react-data-grid/src/DataGridHeaderCellContext.ts
index feddaf9c86..c03a5bd415 100644
--- a/common-react/@dbeaver/react-data-grid/src/DataGridHeaderCellContext.ts
+++ b/common-react/@dbeaver/react-data-grid/src/DataGridHeaderCellContext.ts
@@ -14,6 +14,7 @@ export interface IDataGridHeaderCellContext {
   columnSortingState?: IGridReactiveValue<'asc' | 'desc' | undefined | null, [colIdx: number]>;
   columnSortable?: IGridReactiveValue<boolean, [colIdx: number]>;
   onColumnSort?: (colIdx: number, order: 'asc' | 'desc' | null, isMultiple: boolean) => void;
+  onHeaderKeyDown?: (event: React.KeyboardEvent) => void;
 }
 
 export const DataGridCellHeaderContext = createContext<IDataGridHeaderCellContext | null>(null);
diff --git a/common-react/@dbeaver/react-data-grid/src/renderers/HeaderCellContentRenderer.tsx b/common-react/@dbeaver/react-data-grid/src/renderers/HeaderCellContentRenderer.tsx
index 275d038cd8..0e18d26ae0 100644
--- a/common-react/@dbeaver/react-data-grid/src/renderers/HeaderCellContentRenderer.tsx
+++ b/common-react/@dbeaver/react-data-grid/src/renderers/HeaderCellContentRenderer.tsx
@@ -1,4 +1,4 @@
-import { memo, use, useRef } from 'react';
+import { memo, use } from 'react';
 import { useDrag, useDrop, type DnDStoreProvider } from '@dbeaver/react-dnd';
 import { DataGridCellHeaderContext } from '../DataGridHeaderCellContext.js';
 import { useGridReactiveValue } from '../useGridReactiveValue.js';
@@ -9,6 +9,7 @@ interface Props {
   colIdx: number;
   tabIndex?: number;
 }
+
 export const HeaderCellContentRenderer = memo(function HeaderCellContentRenderer({ colIdx, tabIndex }: Props) {
   const dndHeaderContext = use(HeaderDnDContext);
   const cellHeaderContext = use(DataGridCellHeaderContext);
@@ -16,11 +17,9 @@ export const HeaderCellContentRenderer = memo(function HeaderCellContentRenderer
   const getHeaderText = useGridReactiveValue(headerElement ? undefined : cellHeaderContext?.headerText, colIdx);
   const isColumnSortable = useGridReactiveValue(cellHeaderContext?.columnSortable, colIdx);
   const onColumnSort = cellHeaderContext?.onColumnSort;
+  const onHeaderKeyDown = cellHeaderContext?.onHeaderKeyDown;
   const sortingState = useGridReactiveValue(cellHeaderContext?.columnSortingState, colIdx);
 
-  const orderButtonRef = useRef<HTMLButtonElement>(null);
-  const containerRef = useRef<HTMLDivElement>(null);
-
   const draggable = dndHeaderContext?.getCanDrag?.(colIdx) ?? false;
   const drag = useDrag({
     draggable,
@@ -55,26 +54,32 @@ export const HeaderCellContentRenderer = memo(function HeaderCellContentRenderer
     },
   });
 
-  const handleKeyDown = (e: React.KeyboardEvent) => {
-    if (e.key === 'Tab' && !e.shiftKey && isColumnSortable && onColumnSort && orderButtonRef.current !== document.activeElement) {
-      e.preventDefault();
-      e.stopPropagation();
-      containerRef.current?.parentElement?.setAttribute('aria-selected', 'false');
-      orderButtonRef.current?.focus();
+  function handleSort(e: React.MouseEvent<HTMLElement> | React.KeyboardEvent<HTMLElement>) {
+    if (!onColumnSort) return;
+
+    const nextSortState = sortingState === 'asc' ? 'desc' : sortingState === 'desc' ? null : 'asc';
+    onColumnSort(colIdx, nextSortState, e.ctrlKey || e.metaKey);
+  }
+
+  function onKeyDown(event: React.KeyboardEvent<HTMLElement>) {
+    onHeaderKeyDown?.(event);
+
+    if ((event.key === 'Enter' || event.key === ' ') && isColumnSortable && onColumnSort) {
+      event.preventDefault();
+      handleSort(event);
     }
-  };
+  }
 
   return (
     <div
-      ref={containerRef}
+      onKeyDown={onKeyDown}
       tabIndex={tabIndex}
-      onKeyDown={handleKeyDown}
       className="tw:w-full tw:h-full tw:content-center tw:flex tw:items-center tw:justify-between tw:gap-1 tw:outline-none tw:group"
       {...drag.props}
       {...drop.props}
     >
       <span className="tw:overflow-hidden tw:text-ellipsis">{headerElement ?? getHeaderText ?? ''}</span>
-      {isColumnSortable && onColumnSort && <OrderButton ref={orderButtonRef} colIdx={colIdx} sortState={sortingState} onSort={onColumnSort} />}
+      {isColumnSortable && onColumnSort && <OrderButton sortState={sortingState} onClick={handleSort} />}
     </div>
   );
 });
diff --git a/common-react/@dbeaver/react-data-grid/src/renderers/OrderButton.tsx b/common-react/@dbeaver/react-data-grid/src/renderers/OrderButton.tsx
index 272866d11e..65bac1c825 100644
--- a/common-react/@dbeaver/react-data-grid/src/renderers/OrderButton.tsx
+++ b/common-react/@dbeaver/react-data-grid/src/renderers/OrderButton.tsx
@@ -1,19 +1,12 @@
 import { clsx, IconButton, Icon } from '@dbeaver/ui-kit';
 
 interface OrderButtonProps {
-  colIdx: number;
   sortState?: 'asc' | 'desc' | null;
-  onSort: (attributePosition: number, order: 'asc' | 'desc' | null, isMultiple: boolean) => void;
+  onClick: (e: React.MouseEvent<HTMLElement>) => void;
   tabIndex?: number;
-  ref: React.Ref<HTMLButtonElement>;
 }
 
-export function OrderButton({ colIdx, sortState, onSort, tabIndex, ref }: OrderButtonProps) {
-  function handleClick(e: React.MouseEvent<HTMLButtonElement>) {
-    const nextSortState = sortState === 'asc' ? 'desc' : sortState === 'desc' ? null : 'asc';
-    onSort(colIdx, nextSortState, e.ctrlKey || e.metaKey);
-  }
-
+export function OrderButton({ sortState, onClick, tabIndex }: OrderButtonProps) {
   const svgSortAsc = 'sort-asc';
   const svgSortDesc = 'sort-desc';
   const svgSortUnknown = 'sort-unknown';
@@ -24,9 +17,8 @@ export function OrderButton({ colIdx, sortState, onSort, tabIndex, ref }: OrderB
     <IconButton
       variant="secondary"
       size="small"
-      onClick={handleClick}
+      onClick={onClick}
       tabIndex={tabIndex}
-      ref={ref}
       title="Sort by column"
       aria-label="Sort by column"
       className={clsx(
diff --git a/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/DataGridTable.tsx b/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/DataGridTable.tsx
index d49b833b82..581ecaec7f 100644
--- a/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/DataGridTable.tsx
+++ b/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/DataGridTable.tsx
@@ -84,11 +84,14 @@ export const DataGridTable = observer<IDataPresentationProps>(function DataGridT
   const tableData = useTableData(model as unknown as IDatabaseDataModel<ResultSetDataSource>, resultIndex, dataGridDivRef);
   const gridSelectionContext = useGridSelectionContext(tableData, selectionAction);
 
-  function restoreFocus() {
-    const gridDiv = gridContainerRef.current;
-    const focusSink = gridDiv?.querySelector<HTMLDivElement>('[aria-selected="true"]');
-    focusSink?.focus();
-  }
+  const restoreFocus = useCallback(
+    function restoreFocus() {
+      const gridDiv = gridContainerRef.current;
+      const focusSink = gridDiv?.querySelector<HTMLDivElement>('[aria-selected="true"]');
+      focusSink?.focus();
+    },
+    [gridContainerRef],
+  );
 
   function isGridInFocus(): boolean {
     const gridDiv = gridContainerRef.current;
@@ -544,6 +547,7 @@ export const DataGridTable = observer<IDataPresentationProps>(function DataGridT
               onColumnSort={handleSort}
               onCellChange={handleCellChange}
               onCellKeyDown={handleCellKeyDown}
+              onHeaderKeyDown={gridSelectedCellCopy.onKeydownHandler}
             />
           </div>
         </TableDataContext.Provider>
diff --git a/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/TableColumnHeader/TableIndexColumnHeader.module.css b/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/TableColumnHeader/TableIndexColumnHeader.module.css
index 156d4c9d66..c60f08c603 100644
--- a/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/TableColumnHeader/TableIndexColumnHeader.module.css
+++ b/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/TableColumnHeader/TableIndexColumnHeader.module.css
@@ -6,9 +6,16 @@
  * you may not use this file except in compliance with the License.
  */
 
-.container {
+.clickAreaOverlay {
   cursor: pointer;
+  position: absolute;
+
+  top: 0;
+  left: 0;
+  bottom: 0;
+  right: 0;
 }
+
 .iconOrImage {
   cursor: auto;
   width: 10px;
diff --git a/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/TableColumnHeader/TableIndexColumnHeader.tsx b/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/TableColumnHeader/TableIndexColumnHeader.tsx
index 3fcc2aebab..60574fcb8f 100644
--- a/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/TableColumnHeader/TableIndexColumnHeader.tsx
+++ b/webapp/packages/plugin-data-spreadsheet-new/src/DataGrid/TableColumnHeader/TableIndexColumnHeader.tsx
@@ -36,10 +36,16 @@ export const TableIndexColumnHeader = observer(function TableIndexColumnHeader()
   }
 
   return (
-    <div title={translate('data_grid_table_index_column_tooltip')} className={s(styles, { container: true })} onClick={handleClick}>
+    <>
       {readonly && (
         <IconOrImage title={translate('data_grid_table_readonly_tooltip')} icon="/icons/lock.png" className={s(styles, { iconOrImage: true })} />
       )}
-    </div>
+      <div
+        role="button"
+        title={translate('data_grid_table_index_column_tooltip')}
+        className={s(styles, { clickAreaOverlay: true })}
+        onClick={handleClick}
+      />
+    </>
   );
 });