DATAES-1003 - add timeout to search query.

Original PR: #572
This commit is contained in:
simpleusr 2020-12-18 18:57:40 +03:00 committed by GitHub
parent b49b053d06
commit e950752e79
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 66 additions and 0 deletions

View File

@ -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;
}

View File

@ -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;
}
}

View File

@ -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;
}

View File

@ -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();
}

View File

@ -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();