DATAES-227 - prepareUpdate does not consider UpdateRequest options.

Original PR: #327

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

@@ -55,6 +55,7 @@ import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.search.SearchScrollRequest;
 import org.elasticsearch.action.update.UpdateRequest;
+import org.elasticsearch.action.update.UpdateRequestBuilder;
 import org.elasticsearch.action.update.UpdateResponse;
 import org.elasticsearch.client.RequestOptions;
 import org.elasticsearch.client.Requests;
@@ -713,26 +714,39 @@ public class ElasticsearchRestTemplate extends AbstractElasticsearchTemplate
 	}
 
 	private UpdateRequest prepareUpdate(UpdateQuery query) {
+
 		String indexName = hasText(query.getIndexName()) ? query.getIndexName()
 				: getPersistentEntityFor(query.getClazz()).getIndexName();
 		String type = hasText(query.getType()) ? query.getType() : getPersistentEntityFor(query.getClazz()).getIndexType();
+
 		Assert.notNull(indexName, "No index defined for Query");
 		Assert.notNull(type, "No type define for Query");
 		Assert.notNull(query.getId(), "No Id define for Query");
-		Assert.notNull(query.getUpdateRequest(), "No IndexRequest define for Query");
+		Assert.notNull(query.getUpdateRequest(), "No UpdateRequest define for Query");
-		UpdateRequest updateRequest = new UpdateRequest(indexName, type, query.getId());
-		updateRequest.routing(query.getUpdateRequest().routing());
 
-		if (query.getUpdateRequest().script() == null) {
+		UpdateRequest queryUpdateRequest = query.getUpdateRequest();
-			// doc
+
-			if (query.DoUpsert()) {
+		UpdateRequest updateRequest = new UpdateRequest(indexName, type, query.getId()) //
-				updateRequest.docAsUpsert(true).doc(query.getUpdateRequest().doc());
+				.routing(queryUpdateRequest.routing()) //
-			} else {
+				.retryOnConflict(queryUpdateRequest.retryOnConflict()) //
-				updateRequest.doc(query.getUpdateRequest().doc());
+				.timeout(queryUpdateRequest.timeout()) //
-			}
+				.waitForActiveShards(queryUpdateRequest.waitForActiveShards()) //
-		} else {
+				.setRefreshPolicy(queryUpdateRequest.getRefreshPolicy()) //
-			// or script
+				.waitForActiveShards(queryUpdateRequest.waitForActiveShards()) //
-			updateRequest.script(query.getUpdateRequest().script());
+				.scriptedUpsert(queryUpdateRequest.scriptedUpsert()) //
+				.docAsUpsert(queryUpdateRequest.docAsUpsert());
+
+		if (query.DoUpsert()) {
+			updateRequest.docAsUpsert(true);
+		}
+		if (queryUpdateRequest.script() != null) {
+			updateRequest.script(queryUpdateRequest.script());
+		}
+		if (queryUpdateRequest.doc() != null) {
+			updateRequest.doc(queryUpdateRequest.doc());
+		}
+		if (queryUpdateRequest.upsertRequest() != null) {
+			updateRequest.upsert(queryUpdateRequest.upsertRequest());
 		}
 
 		return updateRequest;

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

@@ -47,6 +47,7 @@ import org.elasticsearch.action.search.MultiSearchRequest;
 import org.elasticsearch.action.search.MultiSearchResponse;
 import org.elasticsearch.action.search.SearchRequestBuilder;
 import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.action.update.UpdateRequest;
 import org.elasticsearch.action.update.UpdateRequestBuilder;
 import org.elasticsearch.action.update.UpdateResponse;
 import org.elasticsearch.client.Client;
@@ -596,27 +597,40 @@ public class ElasticsearchTemplate extends AbstractElasticsearchTemplate impleme
 	}
 
 	private UpdateRequestBuilder prepareUpdate(UpdateQuery query) {
+
 		String indexName = !StringUtils.isEmpty(query.getIndexName()) ? query.getIndexName()
 				: getPersistentEntityFor(query.getClazz()).getIndexName();
 		String type = !StringUtils.isEmpty(query.getType()) ? query.getType()
 				: getPersistentEntityFor(query.getClazz()).getIndexType();
+
 		Assert.notNull(indexName, "No index defined for Query");
 		Assert.notNull(type, "No type define for Query");
 		Assert.notNull(query.getId(), "No Id define for Query");
-		Assert.notNull(query.getUpdateRequest(), "No IndexRequest define for Query");
+		Assert.notNull(query.getUpdateRequest(), "No UpdateRequest define for Query");
-		UpdateRequestBuilder updateRequestBuilder = client.prepareUpdate(indexName, type, query.getId());
-		updateRequestBuilder.setRouting(query.getUpdateRequest().routing());
 
-		if (query.getUpdateRequest().script() == null) {
+		UpdateRequest queryUpdateRequest = query.getUpdateRequest();
-			// doc
+
-			if (query.DoUpsert()) {
+		UpdateRequestBuilder updateRequestBuilder = client.prepareUpdate(indexName, type, query.getId()) //
-				updateRequestBuilder.setDocAsUpsert(true).setDoc(query.getUpdateRequest().doc());
+				.setRouting(queryUpdateRequest.routing()) //
-			} else {
+				.setRetryOnConflict(queryUpdateRequest.retryOnConflict()) //
-				updateRequestBuilder.setDoc(query.getUpdateRequest().doc());
+				.setTimeout(queryUpdateRequest.timeout()) //
-			}
+				.setWaitForActiveShards(queryUpdateRequest.waitForActiveShards()) //
-		} else {
+				.setRefreshPolicy(queryUpdateRequest.getRefreshPolicy()) //
-			// or script
+				.setWaitForActiveShards(queryUpdateRequest.waitForActiveShards()) //
-			updateRequestBuilder.setScript(query.getUpdateRequest().script());
+				.setScriptedUpsert(queryUpdateRequest.scriptedUpsert()) //
+				.setDocAsUpsert(queryUpdateRequest.docAsUpsert());
+
+		if (query.DoUpsert()) {
+			updateRequestBuilder.setDocAsUpsert(true);
+		}
+		if (queryUpdateRequest.script() != null) {
+			updateRequestBuilder.setScript(queryUpdateRequest.script());
+		}
+		if (queryUpdateRequest.doc() != null) {
+			updateRequestBuilder.setDoc(queryUpdateRequest.doc());
+		}
+		if (queryUpdateRequest.upsertRequest() != null) {
+			updateRequestBuilder.setUpsert(queryUpdateRequest.upsertRequest());
 		}
 
 		return updateRequestBuilder;

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

@@ -16,17 +16,22 @@
 package org.springframework.data.elasticsearch.core;
 
 import static org.apache.commons.lang.RandomStringUtils.*;
+import static org.assertj.core.api.Assertions.*;
 import static org.springframework.data.elasticsearch.annotations.FieldType.*;
 
 import lombok.Builder;
 import lombok.Data;
 
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
 import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.index.IndexRequest;
+import org.elasticsearch.action.update.UpdateRequest;
 import org.elasticsearch.common.xcontent.XContentType;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-
 import org.springframework.data.annotation.Id;
 import org.springframework.data.elasticsearch.annotations.Document;
 import org.springframework.data.elasticsearch.annotations.Field;
@@ -34,6 +39,7 @@ import org.springframework.data.elasticsearch.core.query.UpdateQuery;
 import org.springframework.data.elasticsearch.core.query.UpdateQueryBuilder;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringRunner;
+import org.springframework.test.util.ReflectionTestUtils;
 
 /**
  * @author Rizwan Idrees
@@ -64,6 +70,33 @@ public class ElasticsearchRestTemplateTests extends ElasticsearchTemplateTests {
 		elasticsearchTemplate.update(updateQuery);
 	}
 
+	@Test // DATAES-227
+	@Override
+	public void shouldUseUpsertOnUpdate() throws IOException {
+
+		// given
+		Map<String, Object> doc = new HashMap<>();
+		doc.put("id", "1");
+		doc.put("message", "test");
+
+		UpdateRequest updateRequest = new UpdateRequest() //
+				.doc(doc) //
+				.upsert(doc);
+
+		UpdateQuery updateQuery = new UpdateQueryBuilder() //
+				.withClass(SampleEntity.class) //
+				.withId("1") //
+				.withUpdateRequest(updateRequest).build();
+
+		// when
+		UpdateRequest request = (UpdateRequest) ReflectionTestUtils //
+				.invokeMethod(elasticsearchTemplate, "prepareUpdate", updateQuery);
+
+		// then
+		assertThat(request).isNotNull();
+		assertThat(request.upsertRequest()).isNotNull();
+	}
+
 	@Data
 	@Builder
 	@Document(indexName = "test-index-sample-core-rest-template", type = "test-type", shards = 1, replicas = 0,

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

@@ -46,6 +46,8 @@ import org.elasticsearch.action.get.MultiGetResponse;
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.support.IndicesOptions;
+import org.elasticsearch.action.update.UpdateRequest;
+import org.elasticsearch.action.update.UpdateRequestBuilder;
 import org.elasticsearch.cluster.metadata.AliasMetaData;
 import org.elasticsearch.index.VersionType;
 import org.elasticsearch.script.Script;
@@ -84,6 +86,7 @@ import org.springframework.data.elasticsearch.core.query.*;
 import org.springframework.data.util.CloseableIterator;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringRunner;
+import org.springframework.test.util.ReflectionTestUtils;
 
 /**
  * Base for testing rest/transport templates
@@ -1499,6 +1502,33 @@ public class ElasticsearchTemplateTests {
 		assertThat(indexedEntity.getMessage()).isEqualTo(messageAfterUpdate);
 	}
 
+	@Test // DATAES-227
+	public void shouldUseUpsertOnUpdate() throws IOException {
+
+		// given
+		Map<String, Object> doc = new HashMap<>();
+		doc.put("id", "1");
+		doc.put("message", "test");
+
+		UpdateRequest updateRequest = new UpdateRequest() //
+				.doc(doc) //
+				.upsert(doc);
+
+		UpdateQuery updateQuery = new UpdateQueryBuilder() //
+				.withClass(SampleEntity.class) //
+				.withId("1") //
+				.withUpdateRequest(updateRequest).build();
+
+		// when
+		UpdateRequest request = ((UpdateRequestBuilder) ReflectionTestUtils //
+				.invokeMethod(elasticsearchTemplate, "prepareUpdate", updateQuery)) //
+						.request();
+
+		// then
+		assertThat(request).isNotNull();
+		assertThat(request.upsertRequest()).isNotNull();
+	}
+
 	@Test
 	public void shouldDoUpsertIfDocumentDoesNotExist() {