From f8630a09df74e2eb28e49f799cbbd90731519d91 Mon Sep 17 00:00:00 2001 From: Peter-Josef Meisch Date: Mon, 23 Mar 2020 20:59:30 +0100 Subject: [PATCH] Dataes 768 add missing query parameters for an update query. Original PR: #410 --- .../elasticsearch/core/RequestFactory.java | 66 ++++++++++ .../elasticsearch/core/query/UpdateQuery.java | 121 +++++++++++++++++- .../core/ElasticsearchRestTemplateTests.java | 50 ++++++++ .../core/ElasticsearchTemplateTests.java | 4 + .../ElasticsearchTransportTemplateTests.java | 48 +++++++ 5 files changed, 287 insertions(+), 2 deletions(-) diff --git a/src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java b/src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java index 32192d32f..73220fce2 100644 --- a/src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java +++ b/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 includes = query.getFetchSourceIncludes() != null ? query.getFetchSourceIncludes() + : Collections.emptyList(); + List 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 includes = query.getFetchSourceIncludes() != null ? query.getFetchSourceIncludes() + : Collections.emptyList(); + List 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; } diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/UpdateQuery.java b/src/main/java/org/springframework/data/elasticsearch/core/query/UpdateQuery.java index a346bf344..ef68ea7da 100644 --- a/src/main/java/org/springframework/data/elasticsearch/core/query/UpdateQuery.java +++ b/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 fetchSourceIncludes; + @Nullable private List 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 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 fetchSourceIncludes, @Nullable List 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 getFetchSourceIncludes() { + return fetchSourceIncludes; + } + + @Nullable + public List 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 fetchSourceIncludes; + @Nullable private List 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 fetchSourceIncludes) { + this.fetchSourceIncludes = fetchSourceIncludes; + return this; + } + + public Builder withFetchSourceExcludes(List 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 + } } diff --git a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplateTests.java b/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplateTests.java index d3654b964..367cf8c42 100644 --- a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplateTests.java +++ b/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 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"); + } + } diff --git a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java b/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java index 4b8a631d4..fa83206ba 100755 --- a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java +++ b/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; diff --git a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTransportTemplateTests.java b/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTransportTemplateTests.java index 5c2e2469f..6f3126627 100644 --- a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTransportTemplateTests.java +++ b/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 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 {