From a833c89a699cf837429caf5ebc7e9f45953c4c35 Mon Sep 17 00:00:00 2001
From: "mohsin.husen" <smhdiu@gmail.com>
Date: Sat, 23 Nov 2013 13:45:42 +0000
Subject: [PATCH]  DATAES-39 : SearchResultMapper support for scroll and
 queryForPage

---
 .../core/ElasticsearchOperations.java         | 11 +++++
 .../core/ElasticsearchTemplate.java           |  7 +++
 .../core/ElasticsearchTemplateTests.java      | 44 +++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java
index 7996ec864..f8290353f 100644
--- a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java
+++ b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java
@@ -310,6 +310,17 @@ public interface ElasticsearchOperations {
 	 */
 	<T> Page<T> scroll(String scrollId, long scrollTimeInMillis, Class<T> clazz);
 
+    /**
+     * Scrolls the results for give scroll id using custom result mapper
+     *
+     * @param scrollId
+     * @param scrollTimeInMillis
+     * @param mapper
+     * @param <T>
+     * @return
+     */
+    <T> Page<T> scroll(String scrollId, long scrollTimeInMillis, SearchResultMapper mapper);
+
 	/**
 	 * more like this query to search for documents that are "like" a specific document.
 	 * 
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 e53e4ea23..18f3137b1 100644
--- a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java
+++ b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java
@@ -367,6 +367,13 @@ public class ElasticsearchTemplate implements ElasticsearchOperations {
         return resultsMapper.mapResults(response, clazz, null);
     }
 
+    @Override
+    public <T> Page<T> scroll(String scrollId, long scrollTimeInMillis, SearchResultMapper mapper) {
+        SearchResponse response = client.prepareSearchScroll(scrollId)
+                .setScroll(TimeValue.timeValueMillis(scrollTimeInMillis)).execute().actionGet();
+        return mapper.mapResults(response, null, null);
+    }
+
     @Override
     public <T> Page<T> moreLikeThis(MoreLikeThisQuery query, Class<T> clazz) {
         int startRecord = 0;
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 5247816c9..a036b691a 100644
--- a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java
+++ b/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java
@@ -556,6 +556,50 @@ public class ElasticsearchTemplateTests {
         assertThat(sampleEntities.size(), is(equalTo(30)));
     }
 
+    @Test
+    public void shouldReturnResultsForScanAndScrollWithCustomResultMapper() {
+        //given
+        List<IndexQuery> entities = createSampleEntitiesWithMessage("Test message", 30);
+        // when
+        elasticsearchTemplate.bulkIndex(entities);
+        elasticsearchTemplate.refresh(SampleEntity.class, true);
+        // then
+
+        SearchQuery searchQuery = new NativeSearchQueryBuilder().withQuery(matchAllQuery()).withIndices("test-index")
+                .withTypes("test-type").withPageable(new PageRequest(0, 10)).build();
+
+        String scrollId = elasticsearchTemplate.scan(searchQuery, 1000, false);
+        List<SampleEntity> sampleEntities = new ArrayList<SampleEntity>();
+        boolean hasRecords = true;
+        while (hasRecords) {
+            Page<SampleEntity> page = elasticsearchTemplate.scroll(scrollId, 5000L, new SearchResultMapper() {
+                @Override
+                public <T> FacetedPage<T> mapResults(SearchResponse response, Class<T> clazz, Pageable pageable) {
+                    List<SampleEntity> chunk = new ArrayList<SampleEntity>();
+                    for (SearchHit searchHit : response.getHits()) {
+                        if (response.getHits().getHits().length <= 0) {
+                            return null;
+                        }
+                        SampleEntity user = new SampleEntity();
+                        user.setId(searchHit.getId());
+                        user.setMessage((String) searchHit.getSource().get("message"));
+                        chunk.add(user);
+                    }
+                    if (chunk.size() > 0) {
+                        return new FacetedPageImpl<T>((List<T>) chunk);
+                    }
+                    return null;
+                }
+            });
+            if (page != null) {
+                sampleEntities.addAll(page.getContent());
+            } else {
+                hasRecords = false;
+            }
+        }
+        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++) {