Add possiblity to customize drag data for each row (#8706)

Fixes #8396

client/src/main/java/com/vaadin/client/connectors/grid/GridDragSourceExtensionConnector.java

@@ -66,12 +66,27 @@ public class GridDragSourceExtensionConnector extends
 
             JsonObject rowData = gridConnector.getDataSource().getRow(rowIndex);
 
+            // Set drag data in DataTransfer object
             ((NativeEvent) event).getDataTransfer()
-                    .setData(GridDragSourceExtensionState.DATA_TYPE_ROW_DATA,
-                            rowData.toJson());
+                    .setData(GridDragSourceExtensionState.DATA_TYPE_DRAG_DATA,
+                            getDragData(rowData).toJson());
         }
     }
 
+    /**
+     * Gets drag data from the row data if exists or returns complete row data
+     * otherwise.
+     *
+     * @param row
+     *         Row data.
+     * @return Drag data if present or row data otherwise.
+     */
+    private JsonObject getDragData(JsonObject row) {
+        return row.hasKey(GridDragSourceExtensionState.JSONKEY_DRAG_DATA)
+                ? row.getObject(GridDragSourceExtensionState.JSONKEY_DRAG_DATA)
+                : row;
+    }
+
     @Override
     public void onUnregister() {
         super.onUnregister();

server/src/main/java/com/vaadin/ui/GridDragSourceExtension.java

@@ -15,16 +15,34 @@
  */
 package com.vaadin.ui;
 
+import java.util.Optional;
+
+import com.vaadin.data.provider.DataGenerator;
 import com.vaadin.event.dnd.DragSourceExtension;
+import com.vaadin.server.SerializableFunction;
 import com.vaadin.shared.ui.grid.GridDragSourceExtensionState;
 
+import elemental.json.JsonObject;
+
 /**
  * Makes a Grid's rows draggable for HTML5 drag and drop functionality.
  *
+ * @param <T>
+ *         The Grid bean type.
  * @author Vaadin Ltd.
  * @since
  */
-public class GridDragSourceExtension extends DragSourceExtension<Grid> {
+public class GridDragSourceExtension<T> extends DragSourceExtension<Grid<T>> {
+
+    /**
+     * Drag data generator that appends drag data to each row.
+     */
+    private DataGenerator<T> dragDataGenerator;
+
+    /**
+     * Drag data generator function that is executed for each row.
+     */
+    private SerializableFunction<T, JsonObject> generatorFunction;
 
     /**
      * Extends a Grid and makes it's rows draggable.
@@ -32,8 +50,67 @@ public class GridDragSourceExtension extends DragSourceExtension<Grid> {
      * @param target
      *         Grid to be extended.
      */
-    public GridDragSourceExtension(Grid target) {
+    public GridDragSourceExtension(Grid<T> target) {
         super(target);
+
+        // Create drag data generator
+        dragDataGenerator = this::generateDragData;
+
+        // Add drag data generator to Grid
+        target.addDataGenerator(dragDataGenerator);
+    }
+
+    /**
+     * Drag data generator. Appends drag data to row data json if generator
+     * function is set by the user of this extension.
+     *
+     * @param item
+     *         Row item for data generation.
+     * @param jsonObject
+     *         Row data in json format.
+     */
+    private void generateDragData(T item, JsonObject jsonObject) {
+        Optional.ofNullable(generatorFunction).ifPresent(generator -> jsonObject
+                .put(GridDragSourceExtensionState.JSONKEY_DRAG_DATA,
+                        generator.apply(item)));
+    }
+
+    /**
+     * Sets a generator function for customizing drag data. The function is
+     * executed for each item in the Grid during data generation. Return a
+     * {@link JsonObject} to be appended to the row data.
+     * <p>
+     * Example:
+     * <pre>
+     *     dragSourceExtension.setDragDataGenerator(item -> {
+     *         JsonObject dragData = Json.createObject();
+     *         dragData.put("someKey", item.getValue());
+     *         return dragData;
+     *     });
+     * </pre>
+     *
+     * @param generator
+     *         Function to be executed on row data generation.
+     */
+    public void setDragDataGenerator(
+            SerializableFunction<T, JsonObject> generator) {
+        generatorFunction = generator;
+    }
+
+    /**
+     * Returns the generator function for customizing drag data.
+     *
+     * @return Drag data generator function.
+     */
+    public SerializableFunction<T, JsonObject> getDragDataGenerator() {
+        return generatorFunction;
+    }
+
+    @Override
+    public void remove() {
+        super.remove();
+
+        getParent().removeDataGenerator(dragDataGenerator);
     }
 
     @Override

shared/src/main/java/com/vaadin/shared/ui/grid/GridDragSourceExtensionState.java

@@ -26,8 +26,13 @@ import com.vaadin.shared.ui.dnd.DragSourceState;
 public class GridDragSourceExtensionState extends DragSourceState {
 
     /**
-     * Data type for storing the dragged row's data
+     * Data type for storing the dragged row's data.
      */
-    public static final String DATA_TYPE_ROW_DATA = "grid-row-data";
+    public static final String DATA_TYPE_DRAG_DATA = "grid-drag-data";
+
+    /**
+     * Json key for storing data for a dragged row.
+     */
+    public static final String JSONKEY_DRAG_DATA = "drag-data";
 
 }

uitest/src/main/java/com/vaadin/tests/components/grid/GridDragAndDrop.java

@@ -29,6 +29,9 @@ import com.vaadin.ui.HorizontalLayout;
 import com.vaadin.ui.Label;
 import com.vaadin.ui.Layout;
 
+import elemental.json.Json;
+import elemental.json.JsonObject;
+
 public class GridDragAndDrop extends AbstractTestUIWithLog {
     @Override
     protected void setup(VaadinRequest request) {
@@ -39,8 +42,13 @@ public class GridDragAndDrop extends AbstractTestUIWithLog {
         dragSourceComponent.addColumn(Bean::getId).setCaption("ID");
         dragSourceComponent.addColumn(Bean::getValue).setCaption("Value");
 
-        GridDragSourceExtension dragSource = new GridDragSourceExtension(
+        GridDragSourceExtension<Bean> dragSource = new GridDragSourceExtension<>(
                 dragSourceComponent);
+        dragSource.setDragDataGenerator(bean -> {
+            JsonObject ret = Json.createObject();
+            ret.put("val", bean.getValue());
+            return ret;
+        });
 
         Label dropTargetComponent = new Label("Drop here");
         DropTargetExtension<Label> dropTarget = new DropTargetExtension<>(
@@ -48,7 +56,7 @@ public class GridDragAndDrop extends AbstractTestUIWithLog {
 
         dropTarget.addDropListener(event -> {
             log(event.getTransferData(
-                    GridDragSourceExtensionState.DATA_TYPE_ROW_DATA));
+                    GridDragSourceExtensionState.DATA_TYPE_DRAG_DATA));
         });
 
         Layout layout = new HorizontalLayout();