Dataes 768 add missing query parameters for an update query.

Original PR: #410

src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java

@@ -19,6 +19,7 @@ import static org.elasticsearch.index.query.QueryBuilders.*;
 import static org.springframework.util.CollectionUtils.*;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -36,6 +37,7 @@ import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.search.SearchRequestBuilder;
+import org.elasticsearch.action.support.ActiveShardCount;
 import org.elasticsearch.action.update.UpdateRequest;
 import org.elasticsearch.action.update.UpdateRequestBuilder;
 import org.elasticsearch.client.Client;
@@ -499,6 +501,38 @@ class RequestFactory {
 			updateRequest.fetchSource(query.getFetchSource());
 		}
 
+		if (query.getFetchSourceIncludes() != null || query.getFetchSourceExcludes() != null) {
+			List<String> includes = query.getFetchSourceIncludes() != null ? query.getFetchSourceIncludes()
+					: Collections.emptyList();
+			List<String> excludes = query.getFetchSourceExcludes() != null ? query.getFetchSourceExcludes()
+					: Collections.emptyList();
+			updateRequest.fetchSource(includes.toArray(new String[0]), excludes.toArray(new String[0]));
+		}
+
+		if (query.getIfSeqNo() != null) {
+			updateRequest.setIfSeqNo(query.getIfSeqNo());
+		}
+
+		if (query.getIfPrimaryTerm() != null) {
+			updateRequest.setIfPrimaryTerm(query.getIfPrimaryTerm());
+		}
+
+		if (query.getRefresh() != null) {
+			updateRequest.setRefreshPolicy(query.getRefresh().name().toLowerCase());
+		}
+
+		if (query.getRetryOnConflict() != null) {
+			updateRequest.retryOnConflict(query.getRetryOnConflict());
+		}
+
+		if (query.getTimeout() != null) {
+			updateRequest.timeout(query.getTimeout());
+		}
+
+		if (query.getWaitForActiveShards() != null) {
+			updateRequest.waitForActiveShards(ActiveShardCount.parseString(query.getWaitForActiveShards()));
+		}
+
 		return updateRequest;
 	}
 
@@ -541,6 +575,38 @@ class RequestFactory {
 			updateRequestBuilder.setFetchSource(query.getFetchSource());
 		}
 
+		if (query.getFetchSourceIncludes() != null || query.getFetchSourceExcludes() != null) {
+			List<String> includes = query.getFetchSourceIncludes() != null ? query.getFetchSourceIncludes()
+					: Collections.emptyList();
+			List<String> excludes = query.getFetchSourceExcludes() != null ? query.getFetchSourceExcludes()
+					: Collections.emptyList();
+			updateRequestBuilder.setFetchSource(includes.toArray(new String[0]), excludes.toArray(new String[0]));
+		}
+
+		if (query.getIfSeqNo() != null) {
+			updateRequestBuilder.setIfSeqNo(query.getIfSeqNo());
+		}
+
+		if (query.getIfPrimaryTerm() != null) {
+			updateRequestBuilder.setIfPrimaryTerm(query.getIfPrimaryTerm());
+		}
+
+		if (query.getRefresh() != null) {
+			updateRequestBuilder.setRefreshPolicy(query.getRefresh().name().toLowerCase());
+		}
+
+		if (query.getRetryOnConflict() != null) {
+			updateRequestBuilder.setRetryOnConflict(query.getRetryOnConflict());
+		}
+
+		if (query.getTimeout() != null) {
+			updateRequestBuilder.setTimeout(query.getTimeout());
+		}
+
+		if (query.getWaitForActiveShards() != null) {
+			updateRequestBuilder.setWaitForActiveShards(ActiveShardCount.parseString(query.getWaitForActiveShards()));
+		}
+
 		return updateRequestBuilder;
 	}
 

src/main/java/org/springframework/data/elasticsearch/core/query/UpdateQuery.java

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.elasticsearch.core.query;
 
+import java.util.List;
 import java.util.Map;
 
 import org.springframework.data.elasticsearch.core.document.Document;
@@ -40,6 +41,14 @@ public class UpdateQuery {
 	@Nullable private Boolean scriptedUpsert;
 	@Nullable private Boolean docAsUpsert;
 	@Nullable private Boolean fetchSource;
+	@Nullable private List<String> fetchSourceIncludes;
+	@Nullable private List<String> fetchSourceExcludes;
+	@Nullable private Integer ifSeqNo;
+	@Nullable private Integer ifPrimaryTerm;
+	@Nullable private Refresh refresh;
+	@Nullable private Integer retryOnConflict;
+	@Nullable String timeout;
+	@Nullable String waitForActiveShards;
 
 	public static Builder builder(String id) {
 		return new Builder(id);
@@ -47,7 +56,11 @@ public class UpdateQuery {
 
 	private UpdateQuery(String id, @Nullable String script, @Nullable Map<String, Object> params,
 			@Nullable Document document, @Nullable Document upsert, @Nullable String lang, @Nullable String routing,
-			@Nullable Boolean scriptedUpsert, @Nullable Boolean docAsUpsert, @Nullable Boolean fetchSource) {
+			@Nullable Boolean scriptedUpsert, @Nullable Boolean docAsUpsert, @Nullable Boolean fetchSource,
+			@Nullable List<String> fetchSourceIncludes, @Nullable List<String> fetchSourceExcludes, @Nullable Integer ifSeqNo,
+			@Nullable Integer ifPrimaryTerm, @Nullable Refresh refresh, @Nullable Integer retryOnConflict,
+			@Nullable String timeout, @Nullable String waitForActiveShards) {
+
 		this.id = id;
 		this.script = script;
 		this.params = params;
@@ -58,6 +71,14 @@ public class UpdateQuery {
 		this.scriptedUpsert = scriptedUpsert;
 		this.docAsUpsert = docAsUpsert;
 		this.fetchSource = fetchSource;
+		this.fetchSourceIncludes = fetchSourceIncludes;
+		this.fetchSourceExcludes = fetchSourceExcludes;
+		this.ifSeqNo = ifSeqNo;
+		this.ifPrimaryTerm = ifPrimaryTerm;
+		this.refresh = refresh;
+		this.retryOnConflict = retryOnConflict;
+		this.timeout = timeout;
+		this.waitForActiveShards = waitForActiveShards;
 	}
 
 	public String getId() {
@@ -109,6 +130,46 @@ public class UpdateQuery {
 		return fetchSource;
 	}
 
+	@Nullable
+	public List<String> getFetchSourceIncludes() {
+		return fetchSourceIncludes;
+	}
+
+	@Nullable
+	public List<String> getFetchSourceExcludes() {
+		return fetchSourceExcludes;
+	}
+
+	@Nullable
+	public Integer getIfSeqNo() {
+		return ifSeqNo;
+	}
+
+	@Nullable
+	public Integer getIfPrimaryTerm() {
+		return ifPrimaryTerm;
+	}
+
+	@Nullable
+	public Refresh getRefresh() {
+		return refresh;
+	}
+
+	@Nullable
+	public Integer getRetryOnConflict() {
+		return retryOnConflict;
+	}
+
+	@Nullable
+	public String getTimeout() {
+		return timeout;
+	}
+
+	@Nullable
+	public String getWaitForActiveShards() {
+		return waitForActiveShards;
+	}
+
 	public static final class Builder {
 		private String id;
 		@Nullable private String script = null;
@@ -120,6 +181,14 @@ public class UpdateQuery {
 		@Nullable private Boolean scriptedUpsert;
 		@Nullable private Boolean docAsUpsert;
 		@Nullable private Boolean fetchSource;
+		@Nullable private Integer ifSeqNo;
+		@Nullable private Integer ifPrimaryTerm;
+		@Nullable private Refresh refresh;
+		@Nullable private Integer retryOnConflict;
+		@Nullable private String timeout;
+		@Nullable String waitForActiveShards;
+		@Nullable private List<String> fetchSourceIncludes;
+		@Nullable private List<String> fetchSourceExcludes;
 
 		private Builder(String id) {
 			this.id = id;
@@ -170,13 +239,61 @@ public class UpdateQuery {
 			return this;
 		}
 
+		public Builder withIfSeqNo(Integer ifSeqNo) {
+			this.ifSeqNo = ifSeqNo;
+			return this;
+		}
+
+		public Builder withIfPrimaryTerm(Integer ifPrimaryTerm) {
+			this.ifPrimaryTerm = ifPrimaryTerm;
+			return this;
+		}
+
+		public Builder withRefresh(Refresh refresh) {
+			this.refresh = refresh;
+			return this;
+		}
+
+		public Builder withRetryOnConflict(Integer retryOnConflict) {
+			this.retryOnConflict = retryOnConflict;
+			return this;
+		}
+
+		public Builder withTimeout(String timeout) {
+			this.timeout = timeout;
+			return this;
+		}
+
+		public Builder withWaitForActiveShards(String waitForActiveShards) {
+			this.waitForActiveShards = waitForActiveShards;
+			return this;
+		}
+
+		public Builder withFetchSourceIncludes(List<String> fetchSourceIncludes) {
+			this.fetchSourceIncludes = fetchSourceIncludes;
+			return this;
+		}
+
+		public Builder withFetchSourceExcludes(List<String> fetchSourceExcludes) {
+			this.fetchSourceExcludes = fetchSourceExcludes;
+			return this;
+		}
+
 		public UpdateQuery build() {
 
 			if (script == null && document == null) {
 				throw new IllegalArgumentException("either script or document must be set");
 			}
 			return new UpdateQuery(id, script, params, document, upsert, lang, routing, scriptedUpsert, docAsUpsert,
-					fetchSource);
+					fetchSource, fetchSourceIncludes, fetchSourceExcludes, ifSeqNo, ifPrimaryTerm, refresh, retryOnConflict,
+					timeout, waitForActiveShards);
 		}
 	}
+
+	/*
+	 * names will be lowercased on building the query.
+	 */
+	public enum Refresh {
+		True, False, Wait_For
+	}
 }

src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplateTests.java

@@ -21,12 +21,23 @@ import static org.springframework.data.elasticsearch.annotations.FieldType.*;
 
 import lombok.Builder;
 import lombok.Data;
+import lombok.val;
 
+import java.lang.Object;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.elasticsearch.action.support.ActiveShardCount;
+import org.elasticsearch.action.support.WriteRequest;
+import org.elasticsearch.action.update.UpdateRequest;
+import org.elasticsearch.common.unit.TimeValue;
 import org.junit.jupiter.api.Test;
 import org.springframework.data.annotation.Id;
 import org.springframework.data.elasticsearch.UncategorizedElasticsearchException;
 import org.springframework.data.elasticsearch.annotations.Document;
 import org.springframework.data.elasticsearch.annotations.Field;
+import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
 import org.springframework.data.elasticsearch.core.query.UpdateQuery;
 import org.springframework.data.elasticsearch.junit.jupiter.ElasticsearchRestTemplateConfiguration;
 import org.springframework.data.elasticsearch.junit.jupiter.SpringIntegrationTest;
@@ -69,4 +80,43 @@ public class ElasticsearchRestTemplateTests extends ElasticsearchTemplateTests {
 		@Id private String id;
 		@Field(type = Text, store = true, fielddata = true) private String type;
 	}
+
+	@Test // DATAES-768
+	void shouldUseAllOptionsFromUpdateQuery() {
+		Map<String, Object> doc = new HashMap<>();
+		doc.put("id", "1");
+		doc.put("message", "test");
+		org.springframework.data.elasticsearch.core.document.Document document = org.springframework.data.elasticsearch.core.document.Document
+				.from(doc);
+		UpdateQuery updateQuery = UpdateQuery.builder("1") //
+				.withDocument(document) //
+				.withIfSeqNo(42) //
+				.withIfPrimaryTerm(13) //
+				.withScript("script")//
+				.withLang("lang") //
+				.withRefresh(UpdateQuery.Refresh.Wait_For) //
+				.withRetryOnConflict(7) //
+				.withTimeout("4711s") //
+				.withWaitForActiveShards("all") //
+				.withFetchSourceIncludes(Collections.singletonList("incl")) //
+				.withFetchSourceExcludes(Collections.singletonList("excl")) //
+				.build();
+
+		UpdateRequest request = getRequestFactory().updateRequest(updateQuery, IndexCoordinates.of("index"));
+
+		assertThat(request).isNotNull();
+		assertThat(request.ifSeqNo()).isEqualTo(42);
+		assertThat(request.ifPrimaryTerm()).isEqualTo(13);
+		assertThat(request.script().getIdOrCode()).isEqualTo("script");
+		assertThat(request.script().getLang()).isEqualTo("lang");
+		assertThat(request.getRefreshPolicy()).isEqualByComparingTo(WriteRequest.RefreshPolicy.WAIT_UNTIL);
+		assertThat(request.retryOnConflict()).isEqualTo(7);
+		assertThat(request.timeout()).isEqualByComparingTo(TimeValue.parseTimeValue("4711s", "test"));
+		assertThat(request.waitForActiveShards()).isEqualTo(ActiveShardCount.ALL);
+		val fetchSourceContext = request.fetchSource();
+		assertThat(fetchSourceContext).isNotNull();
+		assertThat(fetchSourceContext.includes()).containsExactlyInAnyOrder("incl");
+		assertThat(fetchSourceContext.excludes()).containsExactlyInAnyOrder("excl");
+	}
+
 }

src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java

@@ -28,6 +28,10 @@ import lombok.Data;
 import lombok.EqualsAndHashCode;
 import lombok.NoArgsConstructor;
 
+import java.lang.Double;
+import java.lang.Integer;
+import java.lang.Long;
+import java.lang.Object;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;

src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTransportTemplateTests.java

@@ -20,9 +20,19 @@ import static org.assertj.core.api.Assertions.*;
 import static org.springframework.data.elasticsearch.annotations.FieldType.*;
 
 import lombok.Data;
+import lombok.val;
+
+import java.lang.Object;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 
 import org.elasticsearch.action.search.SearchRequestBuilder;
+import org.elasticsearch.action.support.ActiveShardCount;
+import org.elasticsearch.action.support.WriteRequest;
+import org.elasticsearch.action.update.UpdateRequestBuilder;
 import org.elasticsearch.client.Client;
+import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.index.engine.DocumentMissingException;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -83,6 +93,44 @@ public class ElasticsearchTransportTemplateTests extends ElasticsearchTemplateTe
 		assertThat(searchRequestBuilder.request().source().from()).isEqualTo(30);
 	}
 
+	@Test // DATAES-768
+	void shouldUseAllOptionsFromUpdateQuery() {
+		Map<String, Object> doc = new HashMap<>();
+		doc.put("id", "1");
+		doc.put("message", "test");
+		org.springframework.data.elasticsearch.core.document.Document document = org.springframework.data.elasticsearch.core.document.Document
+				.from(doc);
+		UpdateQuery updateQuery = UpdateQuery.builder("1") //
+				.withDocument(document) //
+				.withIfSeqNo(42) //
+				.withIfPrimaryTerm(13) //
+				.withScript("script")//
+				.withLang("lang") //
+				.withRefresh(UpdateQuery.Refresh.Wait_For) //
+				.withRetryOnConflict(7) //
+				.withTimeout("4711s") //
+				.withWaitForActiveShards("all").withFetchSourceIncludes(Collections.singletonList("incl")) //
+				.withFetchSourceExcludes(Collections.singletonList("excl")) //
+				.build();
+
+		UpdateRequestBuilder request = getRequestFactory().updateRequestBuilderFor(client, updateQuery,
+				IndexCoordinates.of("index"));
+
+		assertThat(request).isNotNull();
+		assertThat(request.request().ifSeqNo()).isEqualTo(42);
+		assertThat(request.request().ifPrimaryTerm()).isEqualTo(13);
+		assertThat(request.request().script().getIdOrCode()).isEqualTo("script");
+		assertThat(request.request().script().getLang()).isEqualTo("lang");
+		assertThat(request.request().getRefreshPolicy()).isEqualByComparingTo(WriteRequest.RefreshPolicy.WAIT_UNTIL);
+		assertThat(request.request().retryOnConflict()).isEqualTo(7);
+		assertThat(request.request().timeout()).isEqualByComparingTo(TimeValue.parseTimeValue("4711s", "test"));
+		assertThat(request.request().waitForActiveShards()).isEqualTo(ActiveShardCount.ALL);
+		val fetchSourceContext = request.request().fetchSource();
+		assertThat(fetchSourceContext).isNotNull();
+		assertThat(fetchSourceContext.includes()).containsExactlyInAnyOrder("incl");
+		assertThat(fetchSourceContext.excludes()).containsExactlyInAnyOrder("excl");
+	}
+
 	@Data
 	@Document(indexName = "test-index-sample-core-transport-template", replicas = 0, refreshInterval = "-1")
 	static class SampleEntity {