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 1c926c938..1838fb112 100644 --- a/src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java +++ b/src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java @@ -1172,6 +1172,11 @@ class RequestFactory { if (StringUtils.hasLength(query.getRoute())) { request.routing(query.getRoute()); } + + TimeValue timeout = query.getTimeout(); + if (timeout !=null) { + sourceBuilder.timeout(timeout); + } request.source(sourceBuilder); return request; @@ -1247,6 +1252,11 @@ class RequestFactory { if (StringUtils.hasLength(query.getRoute())) { searchRequestBuilder.setRouting(query.getRoute()); } + + TimeValue timeout = query.getTimeout(); + if (timeout !=null) { + searchRequestBuilder.setTimeout(timeout); + } return searchRequestBuilder; } diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/AbstractQuery.java b/src/main/java/org/springframework/data/elasticsearch/core/query/AbstractQuery.java index e0e284831..8f9a1d5a7 100755 --- a/src/main/java/org/springframework/data/elasticsearch/core/query/AbstractQuery.java +++ b/src/main/java/org/springframework/data/elasticsearch/core/query/AbstractQuery.java @@ -25,6 +25,7 @@ import java.util.Optional; import org.elasticsearch.action.search.SearchType; import org.elasticsearch.action.support.IndicesOptions; +import org.elasticsearch.common.unit.TimeValue; import org.springframework.data.domain.Pageable; import org.springframework.data.domain.Sort; import org.springframework.lang.Nullable; @@ -59,6 +60,7 @@ abstract class AbstractQuery implements Query { @Nullable private Boolean trackTotalHits; @Nullable private Integer trackTotalHitsUpTo; @Nullable private Duration scrollTime; + @Nullable private TimeValue timeout; @Override @Nullable @@ -252,4 +254,14 @@ abstract class AbstractQuery implements Query { public void setScrollTime(@Nullable Duration scrollTime) { this.scrollTime = scrollTime; } + + @Nullable + @Override + public TimeValue getTimeout() { + return timeout; + } + + public void setTimeout(TimeValue timeout) { + this.timeout = timeout; + } } diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java b/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java index 499c46444..2bf510e4e 100755 --- a/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java +++ b/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java @@ -23,6 +23,7 @@ import java.util.List; import org.elasticsearch.action.search.SearchType; import org.elasticsearch.action.support.IndicesOptions; +import org.elasticsearch.common.unit.TimeValue; import org.elasticsearch.index.query.QueryBuilder; import org.elasticsearch.search.aggregations.AbstractAggregationBuilder; import org.elasticsearch.search.collapse.CollapseBuilder; @@ -68,6 +69,7 @@ public class NativeSearchQueryBuilder { @Nullable private String preference; @Nullable private Integer maxResults; @Nullable private Boolean trackTotalHits; + @Nullable private TimeValue timeout; public NativeSearchQueryBuilder withQuery(QueryBuilder queryBuilder) { this.queryBuilder = queryBuilder; @@ -181,6 +183,11 @@ public class NativeSearchQueryBuilder { this.trackTotalHits = trackTotalHits; return this; } + + public NativeSearchQueryBuilder withTimeout(TimeValue timeout) { + this.timeout = timeout; + return this; + } public NativeSearchQuery build() { @@ -243,6 +250,10 @@ public class NativeSearchQueryBuilder { } nativeSearchQuery.setTrackTotalHits(trackTotalHits); + + if (timeout != null) { + nativeSearchQuery.setTimeout(timeout); + } return nativeSearchQuery; } diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/Query.java b/src/main/java/org/springframework/data/elasticsearch/core/query/Query.java index e5c4b907a..92b964477 100644 --- a/src/main/java/org/springframework/data/elasticsearch/core/query/Query.java +++ b/src/main/java/org/springframework/data/elasticsearch/core/query/Query.java @@ -22,6 +22,7 @@ import java.util.Optional; import org.elasticsearch.action.search.SearchType; import org.elasticsearch.action.support.IndicesOptions; +import org.elasticsearch.common.unit.TimeValue; import org.elasticsearch.index.query.QueryBuilders; import org.springframework.data.domain.PageRequest; import org.springframework.data.domain.Pageable; @@ -275,4 +276,12 @@ public interface Query { default boolean hasScrollTime() { return getScrollTime() != null; } + + /** + * Get timeout + * + * @return null if not set + */ + @Nullable + TimeValue getTimeout(); } diff --git a/src/test/java/org/springframework/data/elasticsearch/core/RequestFactoryTests.java b/src/test/java/org/springframework/data/elasticsearch/core/RequestFactoryTests.java index 6eb792b60..ac8548c41 100644 --- a/src/test/java/org/springframework/data/elasticsearch/core/RequestFactoryTests.java +++ b/src/test/java/org/springframework/data/elasticsearch/core/RequestFactoryTests.java @@ -39,6 +39,7 @@ import org.elasticsearch.action.search.SearchRequest; import org.elasticsearch.action.search.SearchRequestBuilder; import org.elasticsearch.client.Client; import org.elasticsearch.client.indices.PutIndexTemplateRequest; +import org.elasticsearch.common.unit.TimeValue; import org.elasticsearch.common.xcontent.ToXContent; import org.elasticsearch.common.xcontent.XContentHelper; import org.elasticsearch.common.xcontent.XContentType; @@ -481,6 +482,29 @@ class RequestFactoryTests { assertThat(indexRequest.opType()).isEqualTo(DocWriteRequest.OpType.INDEX); } + + @Test + @DisplayName("should set timeout to request") + void shouldSetTimeoutToRequest() { + Query query = new NativeSearchQueryBuilder().withQuery(matchAllQuery()).withTimeout(TimeValue.timeValueSeconds(1)).build(); + + SearchRequest searchRequest = requestFactory.searchRequest(query, Person.class, IndexCoordinates.of("persons")); + + assertThat(searchRequest.source().timeout()).isEqualTo(TimeValue.timeValueSeconds(1)); + } + + @Test + @DisplayName("should set timeout to requestbuilder") + void shouldSetTimeoutToRequestBuilder() { + when(client.prepareSearch(any())).thenReturn(new SearchRequestBuilder(client, SearchAction.INSTANCE)); + Query query = new NativeSearchQueryBuilder().withQuery(matchAllQuery()).withTimeout(TimeValue.timeValueSeconds(1)).build(); + + SearchRequestBuilder searchRequestBuilder = requestFactory.searchRequestBuilder(client, query, Person.class, + IndexCoordinates.of("persons")); + + assertThat(searchRequestBuilder.request().source().timeout()).isEqualTo(TimeValue.timeValueSeconds(1)); + } + private String requestToString(ToXContent request) throws IOException { return XContentHelper.toXContent(request, XContentType.JSON, true).utf8ToString();