add fields parameter for update API (#1822)

src/main/java/org/elasticsearch/action/update/TransportUpdateAction.java

@@ -48,6 +48,7 @@ import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.index.engine.DocumentMissingException;
 import org.elasticsearch.index.engine.DocumentSourceMissingException;
 import org.elasticsearch.index.engine.VersionConflictEngineException;
+import org.elasticsearch.index.get.GetField;
 import org.elasticsearch.index.get.GetResult;
 import org.elasticsearch.index.mapper.internal.ParentFieldMapper;
 import org.elasticsearch.index.mapper.internal.RoutingFieldMapper;
@@ -60,12 +61,16 @@ import org.elasticsearch.index.shard.service.IndexShard;
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.ScriptService;
+import org.elasticsearch.search.lookup.SourceLookup;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Map;
 
+import static com.google.common.collect.Maps.newHashMapWithExpectedSize;
+
 /**
  */
 public class TransportUpdateAction extends TransportInstanceSingleOperationAction<UpdateRequest, UpdateResponse> {
@@ -148,6 +153,35 @@ public class TransportUpdateAction extends TransportInstanceSingleOperationActio
         shardOperation(request, listener, 0);
     }
 
+    protected Map<String, GetField> extractFieldsFromSource(final UpdateRequest request, final Map<String, Object> source) {
+        Map<String, GetField> fields = null;
+        if (request.fields() != null && request.fields().length > 0) {
+            SourceLookup sourceLookup = new SourceLookup();
+            sourceLookup.setNextSource(source);
+            for (String field : request.fields()) {
+                Object value = null;
+                if (field.equals("_source")) {
+                    value = source;
+                } else {
+                    value = sourceLookup.extractValue(field);
+                }
+                if (value != null) {
+                    if (fields == null) {
+                        fields = newHashMapWithExpectedSize(2);
+                    }
+                    GetField getField = fields.get(field);
+                    if (getField == null) {
+                        getField = new GetField(field, new ArrayList<Object>(2));
+                        fields.put(field, getField);
+                    }
+                    getField.values().add(value);
+                }
+            }
+        }
+
+        return fields;
+    }
+
     protected void shardOperation(final UpdateRequest request, final ActionListener<UpdateResponse> listener, final int retryCount) throws ElasticSearchException {
         IndexService indexService = indicesService.indexServiceSafe(request.index());
         IndexShard indexShard = indexService.shardSafe(request.shardId());
@@ -207,6 +241,9 @@ public class TransportUpdateAction extends TransportInstanceSingleOperationActio
             }
         }
 
+        // Extract fields from updated source if necessary
+        final Map<String, GetField> fields = extractFieldsFromSource(request, source);
+
         // TODO: external version type, does it make sense here? does not seem like it...
 
         if (operation == null || "index".equals(operation)) {
@@ -222,6 +259,7 @@ public class TransportUpdateAction extends TransportInstanceSingleOperationActio
                 public void onResponse(IndexResponse response) {
                     UpdateResponse update = new UpdateResponse(response.index(), response.type(), response.id(), response.version());
                     update.matches(response.matches());
+                    update.fields(fields);
                     listener.onResponse(update);
                 }
 
@@ -250,6 +288,7 @@ public class TransportUpdateAction extends TransportInstanceSingleOperationActio
                 @Override
                 public void onResponse(DeleteResponse response) {
                     UpdateResponse update = new UpdateResponse(response.index(), response.type(), response.id(), response.version());
+                    update.fields(fields);
                     listener.onResponse(update);
                 }
 
@@ -271,7 +310,9 @@ public class TransportUpdateAction extends TransportInstanceSingleOperationActio
                 }
             });
         } else if ("none".equals(operation)) {
-            listener.onResponse(new UpdateResponse(getResult.index(), getResult.type(), getResult.id(), getResult.version()));
+            UpdateResponse update = new UpdateResponse(getResult.index(), getResult.type(), getResult.id(), getResult.version());
+            update.fields(fields);
+            listener.onResponse(update);
         } else {
             logger.warn("Used update operation [{}] for script [{}], doing nothing...", operation, request.script);
             listener.onResponse(new UpdateResponse(getResult.index(), getResult.type(), getResult.id(), getResult.version()));

src/main/java/org/elasticsearch/action/update/UpdateRequest.java

@@ -49,6 +49,8 @@ public class UpdateRequest extends InstanceShardOperationRequest {
     @Nullable
     Map<String, Object> scriptParams;
 
+    private String[] fields;
+
     int retryOnConflict = 0;
 
     private String percolate;
@@ -230,6 +232,21 @@ public class UpdateRequest extends InstanceShardOperationRequest {
         return this;
     }
 
+    /**
+     * Explicitly specify the fields that will be returned. By default, nothing is returned.
+     */
+    public UpdateRequest fields(String... fields) {
+        this.fields = fields;
+        return this;
+    }
+
+    /**
+     * Get the fields to be returned.
+     */
+    public String[] fields() {
+        return this.fields;
+    }
+
     /**
      * Sets the number of retries of a version conflict occurs because the document was updated between
      * getting it and updating it. Defaults to 1.
@@ -333,6 +350,13 @@ public class UpdateRequest extends InstanceShardOperationRequest {
             percolate = in.readUTF();
         }
         refresh = in.readBoolean();
+        int size = in.readInt();
+        if (size >= 0) {
+            fields = new String[size];
+            for (int i = 0; i < size; i++) {
+                fields[i] = in.readUTF();
+            }
+        }
     }
 
     @Override
@@ -364,5 +388,13 @@ public class UpdateRequest extends InstanceShardOperationRequest {
             out.writeUTF(percolate);
         }
         out.writeBoolean(refresh);
+        if (fields == null) {
+            out.writeInt(-1);
+        } else {
+            out.writeInt(fields.length);
+            for (String field : fields) {
+                out.writeUTF(field);
+            }
+        }
     }
 }

src/main/java/org/elasticsearch/action/update/UpdateRequestBuilder.java

@@ -111,6 +111,14 @@ public class UpdateRequestBuilder extends BaseRequestBuilder<UpdateRequest, Upda
         return this;
     }
 
+    /**
+     * Explicitly specify the fields that will be returned. By default, nothing is returned.
+     */
+    public UpdateRequestBuilder setFields(String... fields) {
+        request.fields(fields);
+        return this;
+    }
+
     /**
      * Sets the number of retries of a version conflict occurs because the document was updated between
      * getting it and updating it. Defaults to 1.

src/main/java/org/elasticsearch/action/update/UpdateResponse.java

@@ -20,13 +20,19 @@
 package org.elasticsearch.action.update;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.index.get.GetField;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+
+import static com.google.common.collect.Maps.newHashMapWithExpectedSize;
+import static org.elasticsearch.index.get.GetField.readGetField;
 
 /**
  */
@@ -42,6 +48,8 @@ public class UpdateResponse implements ActionResponse {
 
     private List<String> matches;
 
+    private Map<String, GetField> fields;
+
     public UpdateResponse() {
 
     }
@@ -123,6 +131,27 @@ public class UpdateResponse implements ActionResponse {
         return this.matches;
     }
 
+    /**
+     * Internal.
+     */
+    public void fields(Map<String, GetField> fields) {
+        this.fields = fields;
+    }
+
+    /**
+     * Returns extracted fields from updated source. <tt>null</tt> if no field was requested.
+     */
+    public Map<String, GetField> fields() {
+        return this.fields;
+    }
+
+    /**
+     * Returns extracted fields from updated source. <tt>null</tt> if no field was requested.
+     */
+    public Map<String, GetField> getFields() {
+        return this.fields;
+    }
+
     /**
      * Internal.
      */
@@ -157,6 +186,16 @@ public class UpdateResponse implements ActionResponse {
                 }
             }
         }
+        int size = in.readVInt();
+        if (size == 0) {
+            fields = ImmutableMap.of();
+        } else {
+            fields = newHashMapWithExpectedSize(size);
+            for (int i = 0; i < size; i++) {
+                GetField field = readGetField(in);
+                fields.put(field.name(), field);
+            }
+        }
     }
 
     @Override
@@ -174,5 +213,13 @@ public class UpdateResponse implements ActionResponse {
                 out.writeUTF(match);
             }
         }
+        if (fields == null) {
+            out.writeVInt(0);
+        } else {
+            out.writeVInt(fields.size());
+            for (GetField field : fields.values()) {
+                field.writeTo(out);
+            }
+        }
     }
 }

src/main/java/org/elasticsearch/rest/action/update/RestUpdateAction.java

@@ -25,12 +25,14 @@ import org.elasticsearch.action.support.replication.ReplicationType;
 import org.elasticsearch.action.update.UpdateRequest;
 import org.elasticsearch.action.update.UpdateResponse;
 import org.elasticsearch.client.Client;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentBuilderString;
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.index.get.GetField;
 import org.elasticsearch.rest.*;
 import org.elasticsearch.rest.action.support.RestXContentBuilder;
 
@@ -76,6 +78,13 @@ public class RestUpdateAction extends BaseRestHandler {
                 updateRequest.addScriptParam(entry.getKey().substring(3), entry.getValue());
             }
         }
+        String sField = request.param("fields");
+        if (sField != null) {
+            String[] sFields = Strings.splitStringByCommaToArray(sField);
+            if (sFields != null) {
+                updateRequest.fields(sFields);
+            }
+        }
         updateRequest.retryOnConflict(request.paramAsInt("retry_on_conflict", updateRequest.retryOnConflict()));
 
         // see if we have it in the body
@@ -116,6 +125,34 @@ public class RestUpdateAction extends BaseRestHandler {
                             .field(Fields._TYPE, response.type())
                             .field(Fields._ID, response.id())
                             .field(Fields._VERSION, response.version());
+
+                    if (response.fields() != null) {
+                        Map<String, GetField> fields = response.fields();
+                        GetField sourceField = fields.get("_source");
+                        if (sourceField != null) {
+                            builder.field(Fields._SOURCE, sourceField.values().get(0));
+                            fields.remove("_source");
+                        }
+                        if (fields.size() > 0) {
+                            builder.startObject(Fields.FIELDS);
+                            for (GetField field : fields.values()) {
+                                if (field.values().isEmpty()) {
+                                    continue;
+                                }
+                                if (field.values().size() == 1) {
+                                    builder.field(field.name(), field.values().get(0));
+                                } else {
+                                    builder.field(field.name());
+                                    builder.startArray();
+                                    for (Object value : field.values()) {
+                                        builder.value(value);
+                                    }
+                                    builder.endArray();
+                                }
+                            }
+                            builder.endObject();
+                        }
+                    }
                     if (response.matches() != null) {
                         builder.startArray(Fields.MATCHES);
                         for (String match : response.matches()) {
@@ -151,6 +188,8 @@ public class RestUpdateAction extends BaseRestHandler {
         static final XContentBuilderString _TYPE = new XContentBuilderString("_type");
         static final XContentBuilderString _ID = new XContentBuilderString("_id");
         static final XContentBuilderString _VERSION = new XContentBuilderString("_version");
+        static final XContentBuilderString _SOURCE = new XContentBuilderString("_source");
         static final XContentBuilderString MATCHES = new XContentBuilderString("matches");
+        static final XContentBuilderString FIELDS = new XContentBuilderString("fields");
     }
 }

src/test/java/org/elasticsearch/test/integration/update/UpdateTests.java

@@ -169,5 +169,10 @@ public class UpdateTests extends AbstractNodesTests {
         getResponse = client.prepareGet("test", "type1", "3").setFields("_timestamp").execute().actionGet();
         long timestamp = ((Number) getResponse.field("_timestamp").value()).longValue();
         assertThat(timestamp, equalTo(1258294332000L));
+
+        // check fields parameter
+        client.prepareIndex("test", "type1", "1").setSource("field", 1).execute().actionGet();
+        updateResponse = client.prepareUpdate("test", "type1", "1").setScript("ctx._source.field += 1").setFields("_source", "field").execute().actionGet();
+        assertThat(updateResponse.fields().size(), equalTo(2));
     }
 }