From f0ae72141fca0792aaba6ef13a84b8644411b48a Mon Sep 17 00:00:00 2001
From: Dmitriy Yakovlev <hobbut@gmail.com>
Date: Thu, 18 Apr 2019 07:52:58 +0300
Subject: [PATCH] DATAES-565 - prepareScroll doesn't respect SourceFilter from
 the Query.

Original pull request: #272
---
 .../core/ElasticsearchRestTemplate.java       |  6 ++++
 .../core/ElasticsearchTemplate.java           |  6 ++++
 .../core/ElasticsearchTemplateTests.java      | 32 +++++++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplate.java b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplate.java
index 97701d682..0ea8705e1 100644
--- a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplate.java
+++ b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchRestTemplate.java
@@ -133,6 +133,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
  * @author Ivan Greene
  * @author Christoph Strobl
  * @author Lorenzo Spinelli
+ * @author Dmitriy Yakovlev
  */
 public class ElasticsearchRestTemplate
 		implements ElasticsearchOperations, EsClient<RestHighLevelClient>, ApplicationContextAware {
@@ -936,6 +937,11 @@ public class ElasticsearchRestTemplate
 			searchSourceBuilder.size(query.getPageable().getPageSize());
 		}
 
+		if (query.getSourceFilter() != null) {
+			SourceFilter sourceFilter = query.getSourceFilter();
+			searchSourceBuilder.fetchSource(sourceFilter.getIncludes(), sourceFilter.getExcludes());
+		}
+
 		if (!isEmpty(query.getFields())) {
 			searchSourceBuilder.fetchSource(toArray(query.getFields()), null);
 		}
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java
index 9dda4d255..27ae36ed7 100755
--- a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java
+++ b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java
@@ -125,6 +125,7 @@ import org.springframework.util.StringUtils;
  * @author Zetang Zeng
  * @author Ivan Greene
  * @author Christoph Strobl
+ * @author Dmitriy Yakovlev
  */
 public class ElasticsearchTemplate implements ElasticsearchOperations, EsClient<Client>, ApplicationContextAware {
 
@@ -816,6 +817,11 @@ public class ElasticsearchTemplate implements ElasticsearchOperations, EsClient<
 			requestBuilder.setSize(query.getPageable().getPageSize());
 		}
 
+		if (query.getSourceFilter() != null) {
+			SourceFilter sourceFilter = query.getSourceFilter();
+			requestBuilder.setFetchSource(sourceFilter.getIncludes(), sourceFilter.getExcludes());
+		}
+
 		if (!isEmpty(query.getFields())) {
 			requestBuilder.setFetchSource(toArray(query.getFields()), null);
 		}
diff --git a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java b/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java
index 44cf67e5e..4e4fdd0f1 100755
--- a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java
+++ b/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java
@@ -88,6 +88,7 @@ import org.springframework.data.util.CloseableIterator;
  * @author Zetang Zeng
  * @author Peter Nowak
  * @author Ivan Greene
+ * @author Dmitriy Yakovlev
  */
 
 @Ignore
@@ -2686,6 +2687,37 @@ public class ElasticsearchTemplateTests {
 				not(contains(notFindableMessage)));
 	}
 
+	@Test // DATAES-565
+	public void shouldRespectSourceFilterWithScanAndScrollForGivenSearchQuery() {
+		//given
+		List<IndexQuery> entities = createSampleEntitiesWithMessage("Test message", 3);
+		// when
+		elasticsearchTemplate.bulkIndex(entities);
+		elasticsearchTemplate.refresh(SampleEntity.class);
+		// then
+
+		SourceFilter sourceFilter = new FetchSourceFilter(new String[]{"id"}, new String[]{});
+
+		SearchQuery searchQuery = new NativeSearchQueryBuilder().withQuery(matchAllQuery()).withIndices(INDEX_NAME)
+				.withTypes(TYPE_NAME).withPageable(PageRequest.of(0, 10))
+				.withSourceFilter(sourceFilter)
+				.build();
+
+		ScrolledPage<SampleEntity> scroll = (ScrolledPage<SampleEntity>) elasticsearchTemplate.startScroll(1000, searchQuery, SampleEntity.class);
+		List<SampleEntity> sampleEntities = new ArrayList<>();
+		while (scroll.hasContent()) {
+			sampleEntities.addAll(scroll.getContent());
+			scroll = (ScrolledPage<SampleEntity>) elasticsearchTemplate.continueScroll(scroll.getScrollId() , 1000, SampleEntity.class);
+		}
+		elasticsearchTemplate.clearScroll(scroll.getScrollId());
+		assertThat(sampleEntities.size(), is(equalTo(3)));
+		assertThat(sampleEntities.stream().map(SampleEntity::getId).collect(Collectors.toList()),
+				everyItem(notNullValue()));
+		assertThat(sampleEntities.stream().map(SampleEntity::getMessage).collect(Collectors.toList()),
+				everyItem(nullValue()));
+	}
+
+
 	private IndexQuery getIndexQuery(SampleEntity sampleEntity) {
 		return new IndexQueryBuilder()
 				.withId(sampleEntity.getId())