DATAES-1003 - add timeout to search query.

Original PR: #572

Author: simpleusr

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

@@ -1172,6 +1172,11 @@ private SearchRequest prepareSearchRequest(Query query, @Nullable Class<?> clazz
 		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 @@ private SearchRequestBuilder prepareSearchRequestBuilder(Query query, Client cli
 		if (StringUtils.hasLength(query.getRoute())) {
 			searchRequestBuilder.setRouting(query.getRoute());
 		}
+		
+		TimeValue timeout = query.getTimeout();
+		if (timeout !=null) {
+			searchRequestBuilder.setTimeout(timeout);
+		}
 
 		return searchRequestBuilder;
 	}

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

@@ -25,6 +25,7 @@
 
 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 @@ public Duration getScrollTime() {
 	public void setScrollTime(@Nullable Duration scrollTime) {
 		this.scrollTime = scrollTime;
 	}
+	
+	@Nullable
+	@Override
+	public TimeValue getTimeout() {
+		return timeout;
+	}
+
+	public void setTimeout(TimeValue timeout) {
+		this.timeout = timeout;
+	}
 }

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

@@ -23,6 +23,7 @@
 
 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 NativeSearchQueryBuilder withTrackTotalHits(Boolean trackTotalHits) {
 		this.trackTotalHits = trackTotalHits;
 		return this;
 	}
+	
+	public NativeSearchQueryBuilder withTimeout(TimeValue timeout) {   
+		this.timeout = timeout;
+		return this;
+	}
 
 	public NativeSearchQuery build() {
 
@@ -243,6 +250,10 @@ public NativeSearchQuery build() {
 		}
 
 		nativeSearchQuery.setTrackTotalHits(trackTotalHits);
+		
+		if (timeout != null) {
+			nativeSearchQuery.setTimeout(timeout);
+		}
 
 		return nativeSearchQuery;
 	}

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

@@ -22,6 +22,7 @@
 
 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 @@ default Optional<HighlightQuery> getHighlightQuery() {
 	default boolean hasScrollTime() {
 		return getScrollTime() != null;
 	}
+	
+	/**
+	 * Get timeout
+	 *
+	 * @return null if not set
+	 */
+	@Nullable
+	TimeValue getTimeout();
 }

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

@@ -39,6 +39,7 @@
 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 @@ void shouldSetOpTypeIndexIfSpecified() {
 
 		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();