Cwikius-Spring/spring-data-elasticsearch
1
0
Fork 0
mirror of https://github.com/spring-projects/spring-data-elasticsearch.git synced 2025-05-31 09:12:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

DATAES-831 - SearchOperations.searchForStream does not use requested maxResults.

Browse Source 
Original PR: #459
This commit is contained in:
Peter-Josef Meisch 2020-05-17 10:49:50 +02:00 committed by GitHub
parent 391e240b49
commit 506f79a45a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 120 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/main/java/org/springframework/data/elasticsearch/core/AbstractElasticsearchTemplate.java
View File

@ -258,7 +258,11 @@ public abstract class AbstractElasticsearchTemplate implements ElasticsearchOper
long scrollTimeInMillis = TimeValue.timeValueMinutes(1).millis();
// noinspection ConstantConditions
int maxCount = query.isLimiting() ? query.getMaxResults() : 0;
return StreamQueries.streamResults( //
maxCount, //
searchScrollStart(scrollTimeInMillis, query, clazz, index), //
scrollId -> searchScrollContinue(scrollId, scrollTimeInMillis, clazz, index), //
this::searchScrollClear);

14
src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplate.java
View File

@ -40,6 +40,8 @@ import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.search.suggest.SuggestBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
import org.springframework.data.elasticsearch.core.document.DocumentAdapters;
import org.springframework.data.elasticsearch.core.document.SearchDocumentResponse;
@ -88,6 +90,8 @@ import org.springframework.util.Assert;
*/
public class ElasticsearchRestTemplate extends AbstractElasticsearchTemplate {
private static final Logger LOGGER = LoggerFactory.getLogger(ElasticsearchRestTemplate.class);
private RestHighLevelClient client;
private ElasticsearchExceptionTranslator exceptionTranslator;
@ -300,9 +304,13 @@ public class ElasticsearchRestTemplate extends AbstractElasticsearchTemplate {
@Override
public void searchScrollClear(List<String> scrollIds) {
ClearScrollRequest request = new ClearScrollRequest();
request.scrollIds(scrollIds);
execute(client -> client.clearScroll(request, RequestOptions.DEFAULT));
try {
ClearScrollRequest request = new ClearScrollRequest();
request.scrollIds(scrollIds);
execute(client -> client.clearScroll(request, RequestOptions.DEFAULT));
} catch (Exception e) {
LOGGER.warn("Could not clear scroll: {}", e.getMessage());
}
}
@Override

7
src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java
View File

@ -86,6 +86,7 @@ import org.springframework.util.Assert;
public class ElasticsearchTemplate extends AbstractElasticsearchTemplate {
private static final Logger QUERY_LOGGER = LoggerFactory
.getLogger("org.springframework.data.elasticsearch.core.QUERY");
private static final Logger LOGGER = LoggerFactory.getLogger(ElasticsearchTemplate.class);
private Client client;
@Nullable private String searchTimeout;
@ -322,7 +323,11 @@ public class ElasticsearchTemplate extends AbstractElasticsearchTemplate {
@Override
public void searchScrollClear(List<String> scrollIds) {
client.prepareClearScroll().setScrollIds(scrollIds).execute().actionGet();
try {
client.prepareClearScroll().setScrollIds(scrollIds).execute().actionGet();
} catch (Exception e) {
LOGGER.warn("Could not clear scroll: {}", e.getMessage());
}
}
@Override

34
src/main/java/org/springframework/data/elasticsearch/core/StreamQueries.java
View File

@ -18,6 +18,7 @@ package org.springframework.data.elasticsearch.core;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.function.Function;
@ -38,13 +39,15 @@ abstract class StreamQueries {
/**
* Stream query results using {@link SearchScrollHits}.
*
* @param maxCount the maximum number of entities to return, a value of 0 means that all available entities are
* returned
* @param searchHits the initial hits
* @param continueScrollFunction function to continue scrolling applies to the current scrollId.
* @param clearScrollConsumer consumer to clear the scroll context by accepting the scrollIds to clear.
* @param <T>
* @param <T> the entity type
* @return the {@link SearchHitsIterator}.
*/
static <T> SearchHitsIterator<T> streamResults(SearchScrollHits<T> searchHits,
static <T> SearchHitsIterator<T> streamResults(int maxCount, SearchScrollHits<T> searchHits,
Function<String, SearchScrollHits<T>> continueScrollFunction, Consumer<List<String>> clearScrollConsumer) {
Assert.notNull(searchHits, "searchHits must not be null.");
@ -59,20 +62,14 @@ abstract class StreamQueries {
return new SearchHitsIterator<T>() {
// As we couldn't retrieve single result with scroll, store current hits.
private volatile Iterator<SearchHit<T>> scrollHits = searchHits.iterator();
private volatile boolean continueScroll = scrollHits.hasNext();
private volatile AtomicInteger currentCount = new AtomicInteger();
private volatile Iterator<SearchHit<T>> currentScrollHits = searchHits.iterator();
private volatile boolean continueScroll = currentScrollHits.hasNext();
private volatile ScrollState scrollState = new ScrollState(searchHits.getScrollId());
@Override
public void close() {
try {
clearScrollConsumer.accept(scrollState.getScrollIds());
} finally {
scrollHits = null;
scrollState = null;
}
clearScrollConsumer.accept(scrollState.getScrollIds());
}
@Override
@ -99,24 +96,25 @@ abstract class StreamQueries {
@Override
public boolean hasNext() {
if (!continueScroll) {
if (!continueScroll || (maxCount > 0 && currentCount.get() >= maxCount)) {
return false;
}
if (!scrollHits.hasNext()) {
if (!currentScrollHits.hasNext()) {
SearchScrollHits<T> nextPage = continueScrollFunction.apply(scrollState.getScrollId());
scrollHits = nextPage.iterator();
currentScrollHits = nextPage.iterator();
scrollState.updateScrollId(nextPage.getScrollId());
continueScroll = scrollHits.hasNext();
continueScroll = currentScrollHits.hasNext();
}
return scrollHits.hasNext();
return currentScrollHits.hasNext();
}
@Override
public SearchHit<T> next() {
if (hasNext()) {
return scrollHits.next();
currentCount.incrementAndGet();
return currentScrollHits.next();
}
throw new NoSuchElementException();
}

50
src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java
View File

@ -77,7 +77,7 @@ import org.springframework.data.elasticsearch.annotations.ScriptedField;
import org.springframework.data.elasticsearch.core.geo.GeoPoint;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.*;
import org.springframework.data.util.CloseableIterator;
import org.springframework.data.util.StreamUtils;
import org.springframework.lang.Nullable;
/**
@ -1298,27 +1298,33 @@ public abstract class ElasticsearchTemplateTests {
assertThat(sampleEntities).hasSize(30);
}
@Test // DATAES-167
public void shouldReturnResultsWithStreamForGivenCriteriaQuery() {
@Test // DATAES-167, DATAES-831
public void shouldReturnAllResultsWithStreamForGivenCriteriaQuery() {
// given
List<IndexQuery> entities = createSampleEntitiesWithMessage("Test message", 30);
// when
operations.bulkIndex(entities, index);
operations.bulkIndex(createSampleEntitiesWithMessage("Test message", 30), index);
indexOperations.refresh();
// then
CriteriaQuery criteriaQuery = new CriteriaQuery(new Criteria());
criteriaQuery.setPageable(PageRequest.of(0, 10));
CloseableIterator<SearchHit<SampleEntity>> stream = operations.searchForStream(criteriaQuery, SampleEntity.class,
index);
List<SearchHit<SampleEntity>> sampleEntities = new ArrayList<>();
while (stream.hasNext()) {
sampleEntities.add(stream.next());
}
assertThat(sampleEntities).hasSize(30);
long count = StreamUtils
.createStreamFromIterator(operations.searchForStream(criteriaQuery, SampleEntity.class, index)).count();
assertThat(count).isEqualTo(30);
}
@Test // DATAES-831
void shouldLimitStreamResultToRequestedSize() {
operations.bulkIndex(createSampleEntitiesWithMessage("Test message", 30), index);
indexOperations.refresh();
CriteriaQuery criteriaQuery = new CriteriaQuery(new Criteria());
criteriaQuery.setMaxResults(10);
long count = StreamUtils
.createStreamFromIterator(operations.searchForStream(criteriaQuery, SampleEntity.class, index)).count();
assertThat(count).isEqualTo(10);
}
private static List<IndexQuery> createSampleEntitiesWithMessage(String message, int numberOfEntities) {
@ -3128,8 +3134,8 @@ public abstract class ElasticsearchTemplateTests {
operations.refresh(OptimisticEntity.class);
List<Query> queries = singletonList(queryForOne(saved.getId()));
List<SearchHits<OptimisticEntity>> retrievedHits = operations.multiSearch(queries,
OptimisticEntity.class, operations.getIndexCoordinatesFor(OptimisticEntity.class));
List<SearchHits<OptimisticEntity>> retrievedHits = operations.multiSearch(queries, OptimisticEntity.class,
operations.getIndexCoordinatesFor(OptimisticEntity.class));
OptimisticEntity retrieved = retrievedHits.get(0).getSearchHit(0).getContent();
assertThatSeqNoPrimaryTermIsFilled(retrieved);
@ -3162,8 +3168,7 @@ public abstract class ElasticsearchTemplateTests {
operations.save(forEdit1);
forEdit2.setMessage("It'll be great");
assertThatThrownBy(() -> operations.save(forEdit2))
.isInstanceOf(OptimisticLockingFailureException.class);
assertThatThrownBy(() -> operations.save(forEdit2)).isInstanceOf(OptimisticLockingFailureException.class);
}
@Test // DATAES-799
@ -3179,8 +3184,7 @@ public abstract class ElasticsearchTemplateTests {
operations.save(forEdit1);
forEdit2.setMessage("It'll be great");
assertThatThrownBy(() -> operations.save(forEdit2))
.isInstanceOf(OptimisticLockingFailureException.class);
assertThatThrownBy(() -> operations.save(forEdit2)).isInstanceOf(OptimisticLockingFailureException.class);
}
@Test // DATAES-799

57
src/test/java/org/springframework/data/elasticsearch/core/StreamQueriesTest.java
View File

@ -25,6 +25,7 @@ import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import org.junit.jupiter.api.Test;
import org.springframework.data.util.StreamUtils;
/**
* @author Sascha Woo
@ -45,6 +46,7 @@ public class StreamQueriesTest {
// when
SearchHitsIterator<String> iterator = StreamQueries.streamResults( //
0, //
searchHits, //
scrollId -> newSearchScrollHits(Collections.emptyList(), scrollId), //
scrollIds -> clearScrollCalled.set(true));
@ -70,6 +72,7 @@ public class StreamQueriesTest {
// when
SearchHitsIterator<String> iterator = StreamQueries.streamResults( //
0, //
searchHits, //
scrollId -> newSearchScrollHits(Collections.emptyList(), scrollId), //
scrollId -> {});
@ -90,10 +93,12 @@ public class StreamQueriesTest {
Collections.singletonList(new SearchHit<String>(null, 0, null, null, "one")), "s-2");
SearchScrollHits<String> searchHits4 = newSearchScrollHits(Collections.emptyList(), "s-3");
Iterator<SearchScrollHits<String>> searchScrollHitsIterator = Arrays.asList(searchHits1, searchHits2, searchHits3,searchHits4).iterator();
Iterator<SearchScrollHits<String>> searchScrollHitsIterator = Arrays
.asList(searchHits1, searchHits2, searchHits3, searchHits4).iterator();
List<String> clearedScrollIds = new ArrayList<>();
SearchHitsIterator<String> iterator = StreamQueries.streamResults( //
0, //
searchScrollHitsIterator.next(), //
scrollId -> searchScrollHitsIterator.next(), //
scrollIds -> clearedScrollIds.addAll(scrollIds));
@ -106,6 +111,56 @@ public class StreamQueriesTest {
assertThat(clearedScrollIds).isEqualTo(Arrays.asList("s-1", "s-2", "s-3"));
}
@Test // DATAES-831
void shouldReturnAllForRequestedSizeOf0() {
SearchScrollHits<String> searchHits1 = newSearchScrollHits(
Collections.singletonList(new SearchHit<String>(null, 0, null, null, "one")), "s-1");
SearchScrollHits<String> searchHits2 = newSearchScrollHits(
Collections.singletonList(new SearchHit<String>(null, 0, null, null, "one")), "s-2");
SearchScrollHits<String> searchHits3 = newSearchScrollHits(
Collections.singletonList(new SearchHit<String>(null, 0, null, null, "one")), "s-2");
SearchScrollHits<String> searchHits4 = newSearchScrollHits(Collections.emptyList(), "s-3");
Iterator<SearchScrollHits<String>> searchScrollHitsIterator = Arrays
.asList(searchHits1, searchHits2, searchHits3, searchHits4).iterator();
SearchHitsIterator<String> iterator = StreamQueries.streamResults( //
0, //
searchScrollHitsIterator.next(), //
scrollId -> searchScrollHitsIterator.next(), //
scrollIds -> {});
long count = StreamUtils.createStreamFromIterator(iterator).count();
assertThat(count).isEqualTo(3);
}
@Test // DATAES-831
void shouldOnlyReturnRequestedCount() {
SearchScrollHits<String> searchHits1 = newSearchScrollHits(
Collections.singletonList(new SearchHit<String>(null, 0, null, null, "one")), "s-1");
SearchScrollHits<String> searchHits2 = newSearchScrollHits(
Collections.singletonList(new SearchHit<String>(null, 0, null, null, "one")), "s-2");
SearchScrollHits<String> searchHits3 = newSearchScrollHits(
Collections.singletonList(new SearchHit<String>(null, 0, null, null, "one")), "s-2");
SearchScrollHits<String> searchHits4 = newSearchScrollHits(Collections.emptyList(), "s-3");
Iterator<SearchScrollHits<String>> searchScrollHitsIterator = Arrays
.asList(searchHits1, searchHits2, searchHits3, searchHits4).iterator();
SearchHitsIterator<String> iterator = StreamQueries.streamResults( //
2, //
searchScrollHitsIterator.next(), //
scrollId -> searchScrollHitsIterator.next(), //
scrollIds -> {});
long count = StreamUtils.createStreamFromIterator(iterator).count();
assertThat(count).isEqualTo(2);
}
private SearchScrollHits<String> newSearchScrollHits(List<SearchHit<String>> hits, String scrollId) {
return new SearchHitsImpl<String>(hits.size(), TotalHitsRelation.EQUAL_TO, 0, scrollId, hits, null);
}