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

Author: kevinleturc

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

@@ -19,6 +19,7 @@
 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}
 	 *

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

@@ -78,6 +78,7 @@
 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 <T> FacetedPage<T> queryForPage(StringQuery query, Class<T> clazz, Search
 		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());

src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchPartQuery.java

@@ -24,6 +24,8 @@
 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 Object execute(Object[] parameters) {
 		} 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());

src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java

@@ -42,6 +42,7 @@
 import org.springframework.data.elasticsearch.entities.HetroEntity2;
 import org.springframework.data.elasticsearch.entities.SampleEntity;
 import org.springframework.data.elasticsearch.entities.SampleMappingEntity;
+import org.springframework.data.util.CloseableIterator;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 
@@ -896,6 +897,31 @@ public <T> FacetedPage<T> mapResults(SearchResponse response, Class<T> clazz, Pa
 		assertThat(sampleEntities.size(), is(equalTo(30)));
 	}
 
+	/*
+	DATAES-167
+	 */
+	@Test
+	public void shouldReturnResultsWithStreamForGivenCriteriaQuery() {
+		//given
+		List<IndexQuery> entities = createSampleEntitiesWithMessage("Test message", 30);
+		// when
+		elasticsearchTemplate.bulkIndex(entities);
+		elasticsearchTemplate.refresh(SampleEntity.class, true);
+		// then
+
+		CriteriaQuery criteriaQuery = new CriteriaQuery(new Criteria());
+		criteriaQuery.addIndices(INDEX_NAME);
+		criteriaQuery.addTypes(TYPE_NAME);
+		criteriaQuery.setPageable(new PageRequest(0, 10));
+
+		CloseableIterator<SampleEntity> stream = elasticsearchTemplate.stream(criteriaQuery, SampleEntity.class);
+		List<SampleEntity> sampleEntities = new ArrayList<SampleEntity>();
+		while (stream.hasNext()) {
+			sampleEntities.add(stream.next());
+		}
+		assertThat(sampleEntities.size(), is(equalTo(30)));
+	}
+
 	private static List<IndexQuery> createSampleEntitiesWithMessage(String message, int numberOfEntities) {
 		List<IndexQuery> indexQueries = new ArrayList<IndexQuery>();
 		for (int i = 0; i < numberOfEntities; i++) {

src/test/java/org/springframework/data/elasticsearch/repositories/CustomMethodRepositoryTests.java

@@ -19,8 +19,10 @@
 import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.*;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.stream.Stream;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -641,6 +643,22 @@ public void shouldExecuteCustomMethodWithNearPointAndDistance() {
 		assertThat(page.getTotalElements(), is(equalTo(1L)));
 	}
 
+	/*
+	DATAES-165
+	 */
+	@Test
+	public void shouldAllowReturningJava8StreamInCustomQuery() {
+		// given
+		List<SampleEntity> entities = createSampleEntities("abc", 30);
+		repository.save(entities);
+
+		// when
+		Stream<SampleEntity> stream = repository.findByType("abc");
+		// then
+		assertThat(stream, is(notNullValue()));
+		assertThat(stream.count(), is(equalTo(30L)));
+	}
+
 	/*
 	DATAES-106
 	 */
@@ -1174,5 +1192,18 @@ public void shouldCountCustomMethodWithNearPointAndDistance() {
 		// then
 		assertThat(count, is(equalTo(1L)));
 	}
+
+	private List<SampleEntity> createSampleEntities(String type, int numberOfEntities) {
+		List<SampleEntity> entities = new ArrayList<SampleEntity>();
+		for (int i = 0; i < numberOfEntities; i++) {
+			SampleEntity entity = new SampleEntity();
+			entity.setId(randomNumeric(numberOfEntities));
+			entity.setAvailable(true);
+			entity.setMessage("Message");
+			entity.setType(type);
+			entities.add(entity);
+		}
+		return entities;
+	}
 }
 

src/test/java/org/springframework/data/elasticsearch/repositories/custom/SampleCustomMethodRepository.java

@@ -16,6 +16,7 @@
 package org.springframework.data.elasticsearch.repositories.custom;
 
 import java.util.List;
+import java.util.stream.Stream;
 
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.Pageable;
@@ -87,6 +88,8 @@ public interface SampleCustomMethodRepository extends ElasticsearchRepository<Sa
 
 	Page<SampleEntity> findByLocationNear(GeoPoint point, String distance, Pageable pageable);
 
+	Stream<SampleEntity> findByType(String type);
+
 	long countByType(String type);
 
 	long countByTypeNot(String type);