From 97adc4951683a20510b27a93ad1cafd63cec0c81 Mon Sep 17 00:00:00 2001
From: Wroud <wrouds@gmail.com>
Date: Thu, 16 Apr 2020 16:47:43 +0300
Subject: [PATCH] fix(ag-grid-plugin): restore focus after column selection

dbeaver-corp/dbeaver-web#450
---
 .../modules/RangeSelection/rangeController.ts | 27 ++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/webapp/packages/ag-grid-plugin/src/modules/RangeSelection/rangeController.ts b/webapp/packages/ag-grid-plugin/src/modules/RangeSelection/rangeController.ts
index c7ff742250..4f72573850 100644
--- a/webapp/packages/ag-grid-plugin/src/modules/RangeSelection/rangeController.ts
+++ b/webapp/packages/ag-grid-plugin/src/modules/RangeSelection/rangeController.ts
@@ -72,12 +72,16 @@ export class RangeController implements IRangeController {
     private lastSelectedCell?: CellPosition
     private isDragging = false;
     private isDraggingMultiple = false;
+    private lastFocus: CellPosition | null = null
     private startDraggingCell: CellPosition | null = null
     private endDraggingCell: CellPosition | null = null
     private temporaryRange: TemporarySelectionRange | null = null
     private selection?: TableSelection
+    private gridPanel!: GridPanel;
 
-    registerGridComp(gridPanel: GridPanel): void { }
+    registerGridComp(gridPanel: GridPanel): void {
+      this.gridPanel = gridPanel;
+    }
 
     @PostConstruct
     private init(): void {
@@ -152,6 +156,7 @@ export class RangeController implements IRangeController {
       shadowElement.select();
       document.execCommand('copy');
       document.body.removeChild(shadowElement);
+      this.restoreFocus();
     }
 
     private selectColumn(event: IAgColumnClickEvent) {
@@ -169,6 +174,7 @@ export class RangeController implements IRangeController {
       } else {
         this.selection.selectRange(0, lastRowId, [event.columnIndex], event.isMultiple);
       }
+      this.restoreFocus();
       this.dispatchChangedEvent(false, true);
     }
 
@@ -193,6 +199,7 @@ export class RangeController implements IRangeController {
         return;
       }
       this.lastSelectedCell = cell;
+      this.lastFocus = cell;
       if (!this.gridOptionsWrapper.isEnableRangeSelection()
         || cell.column.getColDef().field === INDEX_COLUMN_DEF.field) {
         return;
@@ -239,6 +246,7 @@ export class RangeController implements IRangeController {
       );
       this.dispatchChangedEvent(false, true);
       this.lastSelectedCell = position;
+      this.lastFocus = position;
     }
 
     getCellRanges(): CellRange[] {
@@ -373,6 +381,23 @@ export class RangeController implements IRangeController {
       this.endDraggingCell = null;
     }
 
+    private restoreFocus() {
+      const scroll = this.gridPanel.getVScrollPosition();
+      const nodes = this.gridApi.getRenderedNodes();
+      const isFocusRendered = nodes.some(node => node.rowIndex === this.lastFocus?.rowIndex);
+
+      if (this.lastFocus && isFocusRendered) {
+        this.gridApi.setFocusedCell(this.lastFocus.rowIndex, this.lastFocus.column);
+      } else {
+        const node = nodes.find(node => node.isPixelInRange(scroll.top));
+        const columns = this.columnController.getAllDisplayedColumns();
+
+        if (nodes.length > 0 && columns.length > 0) {
+          this.gridApi.setFocusedCell((node || nodes[0]).rowIndex, columns[0]);
+        }
+      }
+    }
+
     private updateDraggingSelection(isExtends: boolean) {
       if (!this.selection || !this.startDraggingCell || !this.endDraggingCell) {
         return;