DATAES-891 - Returning a Stream from a Query annotated repository method crashes.

Original PR: spring-projects#497

Author: sothawo

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

@@ -28,11 +28,12 @@
 
 public abstract class AbstractElasticsearchRepositoryQuery implements RepositoryQuery {
 
+	protected static final int DEFAULT_STREAM_BATCH_SIZE = 500;
 	protected ElasticsearchQueryMethod queryMethod;
 	protected ElasticsearchOperations elasticsearchOperations;
 
 	public AbstractElasticsearchRepositoryQuery(ElasticsearchQueryMethod queryMethod,
-												ElasticsearchOperations elasticsearchOperations) {
+			ElasticsearchOperations elasticsearchOperations) {
 		this.queryMethod = queryMethod;
 		this.elasticsearchOperations = elasticsearchOperations;
 	}

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

@@ -43,8 +43,6 @@
  */
 public class ElasticsearchPartQuery extends AbstractElasticsearchRepositoryQuery {
 
-	private static final int DEFAULT_STREAM_BATCH_SIZE = 500;
-
 	private final PartTree tree;
 	private final ElasticsearchConverter elasticsearchConverter;
 	private final MappingContext<?, ElasticsearchPersistentProperty> mappingContext;

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

@@ -19,13 +19,15 @@
 import java.util.regex.Pattern;
 
 import org.springframework.core.convert.support.GenericConversionService;
+import org.springframework.data.domain.PageRequest;
 import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
 import org.springframework.data.elasticsearch.core.SearchHitSupport;
 import org.springframework.data.elasticsearch.core.SearchHits;
 import org.springframework.data.elasticsearch.core.convert.DateTimeConverters;
 import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
 import org.springframework.data.elasticsearch.core.query.StringQuery;
 import org.springframework.data.repository.query.ParametersParameterAccessor;
+import org.springframework.data.util.StreamUtils;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.NumberUtils;
@@ -88,6 +90,13 @@ public Object execute(Object[] parameters) {
 			stringQuery.setPageable(accessor.getPageable());
 			SearchHits<?> searchHits = elasticsearchOperations.search(stringQuery, clazz, index);
 			result = SearchHitSupport.page(searchHits, stringQuery.getPageable());
+		} else if (queryMethod.isStreamQuery()) {
+			if (accessor.getPageable().isUnpaged()) {
+				stringQuery.setPageable(PageRequest.of(0, DEFAULT_STREAM_BATCH_SIZE));
+			} else {
+				stringQuery.setPageable(accessor.getPageable());
+			}
+			result = StreamUtils.createStreamFromIterator(elasticsearchOperations.searchForStream(stringQuery, clazz, index));
 		} else if (queryMethod.isCollectionQuery()) {
 			if (accessor.getPageable().isPaged()) {
 				stringQuery.setPageable(accessor.getPageable());

src/test/java/org/springframework/data/elasticsearch/repositories/custommethod/CustomMethodRepositoryBaseTests.java

@@ -1544,6 +1544,28 @@ private List<SampleEntity> createSampleEntities(String type, int numberOfEntitie
 		return entities;
 	}
 
+	@Test // DATAES-891
+	void shouldStreamEntitiesWithQueryAnnotatedMethod() {
+		List<SampleEntity> entities = createSampleEntities("abc", 20);
+		repository.saveAll(entities);
+
+		Stream<SampleEntity> stream = streamingRepository.streamEntitiesByType("abc");
+
+		long count = stream.peek(sampleEntity -> assertThat(sampleEntity).isInstanceOf(SampleEntity.class)).count();
+		assertThat(count).isEqualTo(20);
+	}
+
+	@Test // DATAES-891
+	void shouldStreamSearchHitsWithQueryAnnotatedMethod() {
+		List<SampleEntity> entities = createSampleEntities("abc", 20);
+		repository.saveAll(entities);
+
+		Stream<SearchHit<SampleEntity>> stream = streamingRepository.streamSearchHitsByType("abc");
+
+		long count = stream.peek(sampleEntity -> assertThat(sampleEntity).isInstanceOf(SearchHit.class)).count();
+		assertThat(count).isEqualTo(20);
+	}
+
 	@Data
 	@NoArgsConstructor
 	@AllArgsConstructor
@@ -1687,5 +1709,12 @@ public interface SampleStreamingCustomMethodRepository extends ElasticsearchRepo
 		Stream<SampleEntity> findByType(String type);
 
 		Stream<SampleEntity> findByType(String type, Pageable pageable);
+
+		@Query("{\"bool\": {\"must\": [{\"term\": {\"type\": \"?0\"}}]}}")
+		Stream<SampleEntity> streamEntitiesByType(String type);
+
+		@Query("{\"bool\": {\"must\": [{\"term\": {\"type\": \"?0\"}}]}}")
+		Stream<SearchHit<SampleEntity>> streamSearchHitsByType(String type);
+
 	}
 }