Update Select all -CheckBox from server and partial selections (#17590)

Change-Id: I8f4986455029fc3b997ec5fee8916fa118a487ca

2 files changed, 67 insertions, 32 deletions

diff --git a/client/src/com/vaadin/client/connectors/GridConnector.java b/client/src/com/vaadin/client/connectors/GridConnector.java
index a99b227abe..15acbc0d5a 100644
--- a/client/src/com/vaadin/client/connectors/GridConnector.java
+++ b/client/src/com/vaadin/client/connectors/GridConnector.java
@@ -33,6 +33,7 @@ import com.google.gwt.core.client.Scheduler;
 import com.google.gwt.core.client.Scheduler.ScheduledCommand;
 import com.google.gwt.dom.client.NativeEvent;
 import com.google.gwt.user.client.Timer;
+import com.google.gwt.user.client.ui.CheckBox;
 import com.google.gwt.user.client.ui.Widget;
 import com.vaadin.client.ComponentConnector;
 import com.vaadin.client.ConnectorHierarchyChangeEvent;
@@ -681,6 +682,26 @@ public class GridConnector extends AbstractHasComponentsConnector implements
             public void recalculateColumnWidths() {
                 getWidget().recalculateColumnWidths();
             }
+
+            @Override
+            public void setSelectAll(boolean allSelected) {
+                if (selectionModel instanceof SelectionModel.Multi
+                        && selectionModel.getSelectionColumnRenderer() != null) {
+                    final Widget widget;
+                    try {
+                        HeaderRow defaultHeaderRow = getWidget()
+                                .getDefaultHeaderRow();
+                        if (defaultHeaderRow != null) {
+                            widget = defaultHeaderRow.getCell(
+                                    getWidget().getColumn(0)).getWidget();
+                            ((CheckBox) widget).setValue(allSelected, false);
+                        }
+                    } catch (Exception e) {
+                        getLogger().warning(
+                                "Problems finding select all checkbox.");
+                    }
+                }
+            }
         });
 
         getWidget().addSelectionHandler(internalSelectionChangeHandler);
diff --git a/client/src/com/vaadin/client/widgets/Grid.java b/client/src/com/vaadin/client/widgets/Grid.java
index 377943ed61..8db6e9a55a 100644
--- a/client/src/com/vaadin/client/widgets/Grid.java
+++ b/client/src/com/vaadin/client/widgets/Grid.java
@@ -2377,6 +2377,7 @@ public class Grid<T> extends ResizeComposite implements
 
         private boolean initDone = false;
         private boolean selected = false;
+        private CheckBox selectAllCheckBox;
 
         SelectionColumn(final Renderer<Boolean> selectColumnRenderer) {
             super(selectColumnRenderer);
@@ -2401,41 +2402,55 @@ public class Grid<T> extends ResizeComposite implements
              * exist.
              */
             final SelectionModel.Multi<T> model = (Multi<T>) getSelectionModel();
-            final CheckBox checkBox = GWT.create(CheckBox.class);
-            checkBox.addValueChangeHandler(new ValueChangeHandler<Boolean>() {
 
-                @Override
-                public void onValueChange(ValueChangeEvent<Boolean> event) {
-                    if (event.getValue()) {
-                        fireEvent(new SelectAllEvent<T>(model));
-                        selected = true;
-                    } else {
-                        model.deselectAll();
-                        selected = false;
-                    }
-                }
-            });
-            checkBox.setValue(selected);
-            selectionCell.setWidget(checkBox);
-            // Select all with space when "select all" cell is active
-            addHeaderKeyUpHandler(new HeaderKeyUpHandler() {
-                @Override
-                public void onKeyUp(GridKeyUpEvent event) {
-                    if (event.getNativeKeyCode() != KeyCodes.KEY_SPACE) {
-                        return;
-                    }
-                    HeaderRow targetHeaderRow = getHeader().getRow(
-                            event.getFocusedCell().getRowIndex());
-                    if (!targetHeaderRow.isDefault()) {
-                        return;
+            if (selectAllCheckBox == null) {
+                selectAllCheckBox = GWT.create(CheckBox.class);
+                selectAllCheckBox
+                        .addValueChangeHandler(new ValueChangeHandler<Boolean>() {
+
+                            @Override
+                            public void onValueChange(
+                                    ValueChangeEvent<Boolean> event) {
+                                if (event.getValue()) {
+                                    fireEvent(new SelectAllEvent<T>(model));
+                                    selected = true;
+                                } else {
+                                    model.deselectAll();
+                                    selected = false;
+                                }
+                            }
+                        });
+                selectAllCheckBox.setValue(selected);
+
+                // Select all with space when "select all" cell is active
+                addHeaderKeyUpHandler(new HeaderKeyUpHandler() {
+                    @Override
+                    public void onKeyUp(GridKeyUpEvent event) {
+                        if (event.getNativeKeyCode() != KeyCodes.KEY_SPACE) {
+                            return;
+                        }
+                        HeaderRow targetHeaderRow = getHeader().getRow(
+                                event.getFocusedCell().getRowIndex());
+                        if (!targetHeaderRow.isDefault()) {
+                            return;
+                        }
+                        if (event.getFocusedCell().getColumn() == SelectionColumn.this) {
+                            // Send events to ensure state is updated
+                            selectAllCheckBox.setValue(
+                                    !selectAllCheckBox.getValue(), true);
+                        }
                     }
-                    if (event.getFocusedCell().getColumn() == SelectionColumn.this) {
-                        // Send events to ensure row selection state is
-                        // updated
-                        checkBox.setValue(!checkBox.getValue(), true);
+                });
+            } else {
+                for (HeaderRow row : header.getRows()) {
+                    if (row.getCell(this).getType() == GridStaticCellType.WIDGET) {
+                        // Detach from old header.
+                        row.getCell(this).setText("");
                     }
                 }
-            });
+            }
+
+            selectionCell.setWidget(selectAllCheckBox);
         }
 
         @Override
@@ -2497,7 +2512,6 @@ public class Grid<T> extends ResizeComposite implements
             super.setEditable(editable);
             return this;
         }
-
     }
 
     /**