DATAES-165 - Add Java 8 Steram support in custom repositories.

This commit is contained in:
Kevin Leturc
2015-05-31 18:10:26 +02:00
committed by Artur Konczak
parent 576495932e
commit 416daef2f9
6 changed files with 186 additions and 0 deletions

View File

@@ -19,6 +19,7 @@ import org.elasticsearch.action.update.UpdateResponse;
import org.springframework.data.domain.Page;
import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
import org.springframework.data.elasticsearch.core.query.*;
import org.springframework.data.util.CloseableIterator;
import java.util.LinkedList;
import java.util.List;
@@ -209,6 +210,46 @@ public interface ElasticsearchOperations {
*/
<T> FacetedPage<T> queryForPage(StringQuery query, Class<T> clazz, SearchResultMapper mapper);
/**
* Executes the given {@link CriteriaQuery} against elasticsearch and return result as {@link CloseableIterator}.
* <p>
* Returns a {@link CloseableIterator} that wraps an Elasticsearch scroll context that needs to be closed in case of error.
*
* @param <T> element return type
* @param query
* @param clazz
* @return
* @since 1.3
*/
<T> CloseableIterator<T> stream(CriteriaQuery query, Class<T> clazz);
/**
* Executes the given {@link SearchQuery} against elasticsearch and return result as {@link CloseableIterator}.
* <p>
* Returns a {@link CloseableIterator} that wraps an Elasticsearch scroll context that needs to be closed in case of error.
*
* @param <T> element return type
* @param query
* @param clazz
* @return
* @since 1.3
*/
<T> CloseableIterator<T> stream(SearchQuery query, Class<T> clazz);
/**
* Executes the given {@link SearchQuery} against elasticsearch and return result as {@link CloseableIterator} using custom mapper.
* <p>
* Returns a {@link CloseableIterator} that wraps an Elasticsearch scroll context that needs to be closed in case of error.
*
* @param <T> element return type
* @param query
* @param clazz
* @param mapper
* @return
* @since 1.3
*/
<T> CloseableIterator<T> stream(SearchQuery query, Class<T> clazz, SearchResultMapper mapper);
/**
* Execute the criteria query against elasticsearch and return result as {@link List}
*

View File

@@ -78,6 +78,7 @@ import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersiste
import org.springframework.data.elasticsearch.core.mapping.SimpleElasticsearchMappingContext;
import org.springframework.data.elasticsearch.core.query.*;
import org.springframework.data.mapping.PersistentProperty;
import org.springframework.data.util.CloseableIterator;
import org.springframework.util.Assert;
import java.io.BufferedReader;
@@ -330,6 +331,80 @@ public class ElasticsearchTemplate implements ElasticsearchOperations, Applicati
return mapper.mapResults(response, clazz, query.getPageable());
}
@Override
public <T> CloseableIterator<T> stream(CriteriaQuery query, Class<T> clazz) {
final long scrollTimeInMillis = TimeValue.timeValueMinutes(1).millis();
final String initScrollId = scan(query, scrollTimeInMillis, false);
return doStream(initScrollId, scrollTimeInMillis, clazz, resultsMapper);
}
@Override
public <T> CloseableIterator<T> stream(SearchQuery query, Class<T> clazz) {
return stream(query, clazz, resultsMapper);
}
@Override
public <T> CloseableIterator<T> stream(SearchQuery query, final Class<T> clazz, final SearchResultMapper mapper) {
final long scrollTimeInMillis = TimeValue.timeValueMinutes(1).millis();
final String initScrollId = scan(query, scrollTimeInMillis, false);
return doStream(initScrollId, scrollTimeInMillis, clazz, mapper);
}
private <T> CloseableIterator<T> doStream(final String initScrollId, final long scrollTimeInMillis, final Class<T> clazz, final SearchResultMapper mapper) {
return new CloseableIterator<T>() {
/** As we couldn't retrieve single result with scroll, store current hits. */
private volatile Iterator<T> currentHits;
/** The scroll id. */
private volatile String scrollId = initScrollId;
/** If stream is finished (ie: cluster returns no results. */
private volatile boolean finished;
@Override
public void close() {
try {
// Clear scroll on cluster only in case of error (cause elasticsearch auto clear scroll when it's done)
if (!finished && scrollId != null && currentHits != null && currentHits.hasNext()) {
client.prepareClearScroll().addScrollId(scrollId).execute().actionGet();
}
} finally {
currentHits = null;
scrollId = null;
}
}
@Override
public boolean hasNext() {
// Test if stream is finished
if (finished) {
return false;
}
// Test if it remains hits
if (currentHits == null || !currentHits.hasNext()) {
// Do a new request
SearchResponse response = getSearchResponse(client.prepareSearchScroll(scrollId)
.setScroll(TimeValue.timeValueMillis(scrollTimeInMillis)).execute());
// Save hits and scroll id
currentHits = mapper.mapResults(response, clazz, null).iterator();
finished = !currentHits.hasNext();
scrollId = response.getScrollId();
}
return currentHits.hasNext();
}
@Override
public T next() {
if (hasNext()) {
return currentHits.next();
}
throw new NoSuchElementException();
}
};
}
@Override
public <T> long count(CriteriaQuery criteriaQuery, Class<T> clazz) {
QueryBuilder elasticsearchQuery = new CriteriaQueryProcessor().createQueryFromCriteria(criteriaQuery.getCriteria());

View File

@@ -24,6 +24,8 @@ import org.springframework.data.mapping.context.MappingContext;
import org.springframework.data.repository.query.ParametersParameterAccessor;
import org.springframework.data.repository.query.parser.PartTree;
import org.springframework.util.ClassUtils;
import org.springframework.data.util.CloseableIterator;
import org.springframework.data.util.StreamUtils;
/**
* ElasticsearchPartQuery
@@ -54,6 +56,14 @@ public class ElasticsearchPartQuery extends AbstractElasticsearchRepositoryQuery
} else if (queryMethod.isPageQuery()) {
query.setPageable(accessor.getPageable());
return elasticsearchOperations.queryForPage(query, queryMethod.getEntityInformation().getJavaType());
} else if (queryMethod.isStreamQuery()) {
Class<?> entityType = queryMethod.getEntityInformation().getJavaType();
if (query.getPageable() == null) {
query.setPageable(new PageRequest(0, 20));
}
return StreamUtils.createStreamFromIterator((CloseableIterator<Object>) elasticsearchOperations.stream(query, entityType));
} else if (queryMethod.isCollectionQuery()) {
if (accessor.getPageable() == null) {
int itemCount = (int) elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType());