DATAES-1003 - add timeout to search query.

Original PR: #572

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

@@ -1173,6 +1173,11 @@ class RequestFactory {
 			request.routing(query.getRoute());
 		}
 		
+		TimeValue timeout = query.getTimeout();
+		if (timeout !=null) {
+			sourceBuilder.timeout(timeout);
+		}
+
 		request.source(sourceBuilder);
 		return request;
 	}
@@ -1248,6 +1253,11 @@ class RequestFactory {
 			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 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;
+	}
 }

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;
@@ -182,6 +184,11 @@ public class NativeSearchQueryBuilder {
 		return this;
 	}
 	
+	public NativeSearchQueryBuilder withTimeout(TimeValue timeout) {   
+		this.timeout = timeout;
+		return this;
+	}
+
 	public NativeSearchQuery build() {
 
 		NativeSearchQuery nativeSearchQuery = new NativeSearchQuery(queryBuilder, filterBuilder, sortBuilders,
@@ -244,6 +251,10 @@ public class NativeSearchQueryBuilder {
 
 		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 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();
 }

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;
@@ -482,6 +483,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();
 	}