DATAES-100: Allow configurable searchTimeout.

Author: gpanthe

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

@@ -31,6 +31,7 @@
 import java.util.*;
 
 import org.apache.commons.collections.CollectionUtils;
+import org.elasticsearch.action.ListenableActionFuture;
 import org.elasticsearch.action.admin.cluster.state.ClusterStateRequest;
 import org.elasticsearch.action.admin.indices.create.CreateIndexRequestBuilder;
 import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest;
@@ -51,9 +52,8 @@
 import org.elasticsearch.action.search.SearchRequestBuilder;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.search.SearchType;
-import org.elasticsearch.action.suggest.SuggestRequest;
-import org.elasticsearch.action.suggest.SuggestRequestBuilder;
-import org.elasticsearch.action.suggest.SuggestResponse;
+import org.elasticsearch.action.suggest.SuggestRequestBuilder;
+import org.elasticsearch.action.suggest.SuggestResponse;
 import org.elasticsearch.action.update.UpdateRequestBuilder;
 import org.elasticsearch.action.update.UpdateResponse;
 import org.elasticsearch.client.Client;
@@ -63,7 +63,6 @@
 import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.collect.MapBuilder;
 import org.elasticsearch.common.unit.TimeValue;
-import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.index.query.FilterBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
@@ -74,7 +73,7 @@
 import org.elasticsearch.search.highlight.HighlightBuilder;
 import org.elasticsearch.search.sort.SortBuilder;
 import org.elasticsearch.search.sort.SortOrder;
-import org.elasticsearch.search.suggest.SuggestBuilder;
+import org.elasticsearch.search.suggest.SuggestBuilder;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.BeansException;
@@ -111,6 +110,7 @@ public class ElasticsearchTemplate implements ElasticsearchOperations, Applicati
 	private Client client;
 	private ElasticsearchConverter elasticsearchConverter;
 	private ResultsMapper resultsMapper;
+	private String searchTimeout;
 
 	public ElasticsearchTemplate(Client client) {
 		this(client, null, null);
@@ -135,6 +135,10 @@ public ElasticsearchTemplate(Client client, ElasticsearchConverter elasticsearch
 		this.resultsMapper = (resultsMapper == null) ? new DefaultResultMapper(this.elasticsearchConverter.getMappingContext()) : resultsMapper;
 	}
 
+	public void setSearchTimeout(String searchTimeout) {
+		this.searchTimeout = searchTimeout;
+	}
+
 	@Override
 	public <T> boolean createIndex(Class<T> clazz) {
 		return createIndexIfNotCreated(clazz);
@@ -285,7 +289,7 @@ public <T> List<String> queryForIds(SearchQuery query) {
 		if (query.getFilter() != null) {
 			request.setPostFilter(query.getFilter());
 		}
-		SearchResponse response = request.execute().actionGet();
+		SearchResponse response = getSearchResponse(request.execute());
 		return extractIds(response);
 	}
 
@@ -308,8 +312,8 @@ public <T> Page<T> queryForPage(CriteriaQuery criteriaQuery, Class<T> clazz) {
 		if (elasticsearchFilter != null)
 			searchRequestBuilder.setPostFilter(elasticsearchFilter);
 
-		SearchResponse response = searchRequestBuilder
-				.execute().actionGet();
+		SearchResponse response = getSearchResponse(searchRequestBuilder
+				.execute());
 		return resultsMapper.mapResults(response, clazz, criteriaQuery.getPageable());
 	}
 
@@ -320,7 +324,7 @@ public <T> FacetedPage<T> queryForPage(StringQuery query, Class<T> clazz) {
 
 	@Override
 	public <T> FacetedPage<T> queryForPage(StringQuery query, Class<T> clazz, SearchResultMapper mapper) {
-		SearchResponse response = prepareSearch(query, clazz).setQuery(query.getSource()).execute().actionGet();
+		SearchResponse response = getSearchResponse(prepareSearch(query, clazz).setQuery(query.getSource()).execute());
 		return mapper.mapResults(response, clazz, query.getPageable());
 	}
 
@@ -569,20 +573,20 @@ public String scan(SearchQuery searchQuery, long scrollTimeInMillis, boolean noF
 		if (noFields) {
 			requestBuilder.setNoFields();
 		}
-		return requestBuilder.execute().actionGet().getScrollId();
+		return getSearchResponse(requestBuilder.execute()).getScrollId();
 	}
 
 	@Override
 	public <T> Page<T> scroll(String scrollId, long scrollTimeInMillis, Class<T> clazz) {
-		SearchResponse response = client.prepareSearchScroll(scrollId)
-				.setScroll(TimeValue.timeValueMillis(scrollTimeInMillis)).execute().actionGet();
+		SearchResponse response = getSearchResponse(client.prepareSearchScroll(scrollId)
+				.setScroll(TimeValue.timeValueMillis(scrollTimeInMillis)).execute());
 		return resultsMapper.mapResults(response, clazz, null);
 	}
 
 	@Override
 	public <T> Page<T> scroll(String scrollId, long scrollTimeInMillis, SearchResultMapper mapper) {
-		SearchResponse response = client.prepareSearchScroll(scrollId)
-				.setScroll(TimeValue.timeValueMillis(scrollTimeInMillis)).execute().actionGet();
+		SearchResponse response = getSearchResponse(client.prepareSearchScroll(scrollId)
+				.setScroll(TimeValue.timeValueMillis(scrollTimeInMillis)).execute());
 		return mapper.mapResults(response, null, null);
 	}
 
@@ -645,7 +649,7 @@ public <T> Page<T> moreLikeThis(MoreLikeThisQuery query, Class<T> clazz) {
 			requestBuilder.setBoostTerms(query.getBoostTerms());
 		}
 
-		SearchResponse response = requestBuilder.execute().actionGet();
+		SearchResponse response = getSearchResponse(requestBuilder.execute());
 		return resultsMapper.mapResults(response, clazz, query.getPageable());
 	}
 
@@ -681,8 +685,11 @@ private SearchResponse doSearch(SearchRequestBuilder searchRequest, SearchQuery
 				searchRequest.addAggregation(aggregationBuilder);
 			}
 		}
+		return getSearchResponse(searchRequest.setQuery(searchQuery.getQuery()).execute());
+	}
 
-		return searchRequest.setQuery(searchQuery.getQuery()).execute().actionGet();
+	private SearchResponse getSearchResponse(ListenableActionFuture<SearchResponse> response) {
+		return searchTimeout == null ? response.actionGet() : response.actionGet(searchTimeout);
 	}
 
 	private <T> boolean createIndexIfNotCreated(Class<T> clazz) {
@@ -981,14 +988,14 @@ public static String readFileFromClasspath(String url) {
 
 		return stringBuilder.toString();
 	}
-
-    public SuggestResponse suggest(SuggestBuilder.SuggestionBuilder<?> suggestion, String... indices) {
-	SuggestRequestBuilder suggestRequestBuilder = client.prepareSuggest(indices);
-	suggestRequestBuilder.addSuggestion(suggestion);
-	return suggestRequestBuilder.execute().actionGet();
-    }
-
-    public SuggestResponse suggest(SuggestBuilder.SuggestionBuilder<?> suggestion, Class clazz) {
-	return suggest(suggestion, retrieveIndexNameFromPersistentEntity(clazz));
-    }
+
+	public SuggestResponse suggest(SuggestBuilder.SuggestionBuilder<?> suggestion, String... indices) {
+		SuggestRequestBuilder suggestRequestBuilder = client.prepareSuggest(indices);
+		suggestRequestBuilder.addSuggestion(suggestion);
+		return suggestRequestBuilder.execute().actionGet();
+	}
+
+	public SuggestResponse suggest(SuggestBuilder.SuggestionBuilder<?> suggestion, Class clazz) {
+		return suggest(suggestion, retrieveIndexNameFromPersistentEntity(clazz));
+	}
 }