Repository: searchbox-io/Jest
Pull request #677 (2020-05-09 23:02:45)

Change class returned by search scroll to be search result.
Allow you to use the getHits functionality rather than having to do your own object mapping.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/30677225/badge)](https://coveralls.io/builds/30677225)

Coverage decreased (-0.03%) to 58.008% when pulling **6fd0b6184d009651bbd55eed831d4d10b3b97970 on AndreasAbdi:scrollsearch** into **debfca8aafe698d3a4a236068c347d019b2d16b2 on searchbox-io:master**.


Patch:
diff --git a/jest-common/src/main/java/io/searchbox/core/SearchScroll.java b/jest-common/src/main/java/io/searchbox/core/SearchScroll.java
index bff5a91d..85ca9302 100644
--- a/jest-common/src/main/java/io/searchbox/core/SearchScroll.java
+++ b/jest-common/src/main/java/io/searchbox/core/SearchScroll.java
@@ -2,7 +2,9 @@ package io.searchbox.core;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
+import com.google.gson.Gson;
 import com.google.gson.JsonObject;
+import io.searchbox.action.AbstractAction;
 import io.searchbox.action.AbstractMultiIndexActionBuilder;
 import io.searchbox.action.GenericResultAbstractAction;
 import io.searchbox.client.config.ElasticsearchVersion;
@@ -11,13 +13,11 @@ import io.searchbox.params.Parameters;
 /**
  * @author ferhat
  */
-public class SearchScroll extends GenericResultAbstractAction {
+public class SearchScroll extends AbstractAction<SearchResult> {
     @VisibleForTesting
     static final int MAX_SCROLL_ID_LENGTH = 1900;
     private final String restMethodName;
     public static final String SCROLL_ID = "scroll_id";
-    public static final String COMMA = ",";
-
 
     protected SearchScroll(Builder builder) {
         super(builder);
@@ -33,6 +33,11 @@ public class SearchScroll extends GenericResultAbstractAction {
         }
     }
 
+    @Override
+    public SearchResult createNewElasticSearchResult(String responseBody, int statusCode, String reasonPhrase, Gson gson) {
+        return createNewElasticSearchResult(new SearchResult(gson), responseBody, statusCode, reasonPhrase, gson);
+    }
+
     @Override
     protected String buildURI(ElasticsearchVersion elasticsearchVersion) {
         return super.buildURI(elasticsearchVersion) + "/_search/scroll";
diff --git a/jest-common/src/main/java/io/searchbox/core/search/aggregation/AggregationField.java b/jest-common/src/main/java/io/searchbox/core/search/aggregation/AggregationField.java
index 8bf49006..a5df26be 100644
--- a/jest-common/src/main/java/io/searchbox/core/search/aggregation/AggregationField.java
+++ b/jest-common/src/main/java/io/searchbox/core/search/aggregation/AggregationField.java
@@ -31,7 +31,8 @@ public enum AggregationField {
     AVG("avg"),
     SUM("sum"),
     DOC_COUNT_ERROR_UPPER_BOUND("doc_count_error_upper_bound"),
-    SUM_OTHER_DOC_COUNT("sum_other_doc_count");
+    SUM_OTHER_DOC_COUNT("sum_other_doc_count"),
+    AFTER_KEY("after_key");
 
     private final String field;
 
diff --git a/jest-common/src/main/java/io/searchbox/core/search/aggregation/CompositeAggregation.java b/jest-common/src/main/java/io/searchbox/core/search/aggregation/CompositeAggregation.java
new file mode 100644
index 00000000..6c0d9c98
--- /dev/null
+++ b/jest-common/src/main/java/io/searchbox/core/search/aggregation/CompositeAggregation.java
@@ -0,0 +1,70 @@
+package io.searchbox.core.search.aggregation;
+
+import com.google.gson.JsonArray;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+import static io.searchbox.core.search.aggregation.AggregationField.*;
+
+public class CompositeAggregation extends BucketAggregation{
+
+    private List<Entry> buckets;
+    private JsonObject afterKey;
+
+    public CompositeAggregation(String name, JsonObject agg) {
+        super(name, agg);
+        if (agg.has(AFTER_KEY.toString()) && agg.get(AFTER_KEY.toString()).isJsonObject()) {
+            afterKey = agg.get(AFTER_KEY.toString()).getAsJsonObject();
+        }
+        if (agg.has(BUCKETS.toString()) && agg.get(BUCKETS.toString()).isJsonArray()) {
+            final JsonArray bucketsArray = agg.get(BUCKETS.toString()).getAsJsonArray();
+            buckets = StreamSupport.stream(bucketsArray.spliterator(), false)
+                    .map(JsonElement::getAsJsonObject)
+                    .map(bucket ->   new Entry(bucket,
+                                        bucket.get(KEY.toString()).getAsJsonObject(),
+                                        bucket.get(DOC_COUNT.toString()).getAsLong()))
+                    .collect(Collectors.toList());
+        }
+    }
+
+    @Override
+    public List<Entry> getBuckets() {
+        return buckets;
+    }
+
+    public JsonObject getAfterKey() {
+        return afterKey;
+    }
+
+    public class Entry extends Bucket {
+        private JsonObject key;
+
+        public Entry(JsonObject bucketRoot,  JsonObject key, Long count) {
+            super(bucketRoot, count);
+            this.key = key;
+        }
+
+        public JsonObject getKey() {
+            return key;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            if (!super.equals(o)) return false;
+            Entry entry = (Entry) o;
+            return Objects.equals(key, entry.key);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(super.hashCode(), key);
+        }
+    }
+}
diff --git a/jest-common/src/main/java/io/searchbox/core/search/aggregation/MetricAggregation.java b/jest-common/src/main/java/io/searchbox/core/search/aggregation/MetricAggregation.java
index e2bcb91c..eec9ca41 100644
--- a/jest-common/src/main/java/io/searchbox/core/search/aggregation/MetricAggregation.java
+++ b/jest-common/src/main/java/io/searchbox/core/search/aggregation/MetricAggregation.java
@@ -73,6 +73,14 @@ public abstract class MetricAggregation extends Aggregation {
         return jsonRoot.has(aggName) ? new CardinalityAggregation(aggName, jsonRoot.getAsJsonObject(aggName)) : null;
     }
 
+    /**
+     * @param aggName Name of the CompositeAggregation
+     * @return a new CompositeAggregation object if aggName is found within sub-aggregations of current aggregation level or null if not found
+     */
+    public CompositeAggregation getCompositeAggregation(String aggName) {
+        return jsonRoot.has(aggName) ? new CompositeAggregation(aggName, jsonRoot.getAsJsonObject(aggName)) : null;
+    }
+
     /**
      * @param aggName Name of the DateHistogramAggregation
      * @return a new DateHistogramAggregation object if aggName is found within sub-aggregations of current aggregation level or null if not found


Test patch:
diff --git a/jest-common/src/test/java/io/searchbox/core/search/aggregation/CompositeAggregationTest.java b/jest-common/src/test/java/io/searchbox/core/search/aggregation/CompositeAggregationTest.java
new file mode 100644
index 00000000..f5e55271
--- /dev/null
+++ b/jest-common/src/test/java/io/searchbox/core/search/aggregation/CompositeAggregationTest.java
@@ -0,0 +1,65 @@
+package io.searchbox.core.search.aggregation;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonObject;
+
+import java.util.Arrays;
+import java.util.List;
+
+import org.junit.Test;
+import static org.junit.Assert.*;
+
+public class CompositeAggregationTest {
+
+    final String afterKey = "'after_key':{'field':'valB'}";
+    final String bucketsArr = "'buckets':[" +
+            "{  'key':{'field':'val'}," +
+            "   'doc_count':1," +
+            "   'sub_agg':{}" +
+            "}," +
+            "{  'key':{'field':'valA'}," +
+            "   'doc_count':2," +
+            "   'sub_agg':{}" +
+            "}" + "]";
+    final Gson gson = new Gson();
+
+    @Test
+    public void testParseBuckets() {
+        for (String input : Arrays.asList("{" + afterKey + ", " + bucketsArr + "}", "{" + bucketsArr + "}")) {
+            JsonObject compositeAggregationJson = gson.fromJson(
+                    input.replace('\'', '\"'), JsonObject.class);
+            CompositeAggregation compositeAggregation = new CompositeAggregation("composite", compositeAggregationJson);
+            List<CompositeAggregation.Entry> buckets = compositeAggregation.getBuckets();
+            assertNotNull(buckets);
+            assertEquals(2, buckets.size());
+
+            assertEquals("val", buckets.get(0).getKey().get("field").getAsString());
+            assertEquals(Long.valueOf(1L), buckets.get(0).getCount());
+            assertTrue(buckets.get(0).getTermsAggregation("sub_agg").getBuckets().isEmpty());
+
+            assertEquals("valA", buckets.get(1).getKey().get("field").getAsString());
+            assertEquals(Long.valueOf(2L), buckets.get(1).getCount());
+            assertTrue(buckets.get(1).getTermsAggregation("sub_agg").getBuckets().isEmpty());
+        }
+    }
+
+    @Test
+    public void testParseAfterKey() {
+        for (String input : Arrays.asList("{" + afterKey + ", " + bucketsArr + "}", "{" + afterKey + "}")) {
+            JsonObject compositeAggregationJson = gson.fromJson(
+                    input.replace('\'', '\"'), JsonObject.class);
+            CompositeAggregation compositeAggregation = new CompositeAggregation("composite", compositeAggregationJson);
+            assertEquals("valB", compositeAggregation.getAfterKey().get("field").getAsString());
+        }
+    }
+
+    @Test
+    public void testNoAfterKey() {
+        for (String input : Arrays.asList("{"  + bucketsArr + "}", "{" + "}")) {
+            JsonObject compositeAggregationJson = gson.fromJson(
+                    input.replace('\'', '\"'), JsonObject.class);
+            CompositeAggregation compositeAggregation = new CompositeAggregation("composite", compositeAggregationJson);
+            assertNull(compositeAggregation.getAfterKey());
+        }
+    }
+}
\ No newline at end of file
diff --git a/jest/src/test/java/io/searchbox/core/SearchScrollIntegrationTest.java b/jest/src/test/java/io/searchbox/core/SearchScrollIntegrationTest.java
index 54828f76..7e2a53a9 100644
--- a/jest/src/test/java/io/searchbox/core/SearchScrollIntegrationTest.java
+++ b/jest/src/test/java/io/searchbox/core/SearchScrollIntegrationTest.java
@@ -46,7 +46,7 @@ public class SearchScrollIntegrationTest extends AbstractIntegrationTest {
                 .addSort(new Sort("code"))
                 .setParameter(Parameters.SCROLL, "5m")
                 .build();
-        JestResult result = client.execute(search);
+        SearchResult result = client.execute(search);
         assertTrue(result.getErrorMessage(), result.isSucceeded());
         JsonArray hits = result.getJsonObject().getAsJsonObject("hits").getAsJsonArray("hits");
         assertEquals(
@@ -81,8 +81,8 @@ public class SearchScrollIntegrationTest extends AbstractIntegrationTest {
 
         // clear a single scroll id
         ClearScroll clearScroll = new ClearScroll.Builder().addScrollId(scrollId).build();
-        result = client.execute(clearScroll);
-        assertTrue(result.getErrorMessage(), result.isSucceeded());
+        JestResult clearResult = client.execute(clearScroll);
+        assertTrue(clearResult.getErrorMessage(), clearResult.isSucceeded());
     }
 
     @Test
@@ -103,7 +103,7 @@ public class SearchScrollIntegrationTest extends AbstractIntegrationTest {
             .addSort(new Sort("code"))
             .setParameter(Parameters.SCROLL, "5m")
             .build();
-        JestResult result = client.execute(search);
+        SearchResult result = client.execute(search);
         assertTrue(result.getErrorMessage(), result.isSucceeded());
         JsonArray hits = result.getJsonObject().getAsJsonObject("hits").getAsJsonArray("hits");
         assertEquals(
@@ -138,8 +138,8 @@ public class SearchScrollIntegrationTest extends AbstractIntegrationTest {
 
         // clear all scroll ids
         ClearScroll clearScroll = new ClearScroll.Builder().build();
-        result = client.execute(clearScroll);
-        assertTrue(result.getErrorMessage(), result.isSucceeded());
+        JestResult clearResult = client.execute(clearScroll);
+        assertTrue(clearResult.getErrorMessage(), clearResult.isSucceeded());
     }
 
 }
diff --git a/jest/src/test/java/io/searchbox/core/search/aggregation/CompositeAggregationIntegrationTest.java b/jest/src/test/java/io/searchbox/core/search/aggregation/CompositeAggregationIntegrationTest.java
new file mode 100644
index 00000000..2eb964b9
--- /dev/null
+++ b/jest/src/test/java/io/searchbox/core/search/aggregation/CompositeAggregationIntegrationTest.java
@@ -0,0 +1,114 @@
+package io.searchbox.core.search.aggregation;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonObject;
+import io.searchbox.common.AbstractIntegrationTest;
+import io.searchbox.core.Search;
+import io.searchbox.core.SearchResult;
+import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
+import org.elasticsearch.action.admin.indices.mapping.put.PutMappingResponse;
+import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.test.ESIntegTestCase;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+@ESIntegTestCase.ClusterScope (scope = ESIntegTestCase.Scope.TEST, numDataNodes = 1)
+public class CompositeAggregationIntegrationTest extends AbstractIntegrationTest {
+
+    public static final Gson gson = new Gson();
+    private final String INDEX = "composite_aggregation";
+    private final String TYPE = "document";
+
+    @Test
+    public void testGetCompositeAggregation()
+            throws IOException {
+        createIndex(INDEX);
+        PutMappingResponse putMappingResponse = client().admin().indices().putMapping(new PutMappingRequest(INDEX)
+                .type(TYPE)
+                .source("{\"document\":{\"properties\":{\"gender\":{\"store\":true,\"type\":\"keyword\"}}}}", XContentType.JSON)
+        ).actionGet();
+
+        assertTrue(putMappingResponse.isAcknowledged());
+
+        index(INDEX, TYPE, null, "{\"gender\":\"male\"}");
+        index(INDEX, TYPE, null, "{\"gender\":\"male\"}");
+        index(INDEX, TYPE, null, "{\"gender\":\"female\"}");
+        refresh();
+        ensureSearchable(INDEX);
+        final JsonObject afterKey;
+        {
+            String query = "{\n" +
+                    "    \"query\" : {\n" +
+                    "        \"match_all\" : {}\n" +
+                    "    },\n" +
+                    "    \"aggs\" : {\n" +
+                    "        \"composite1\" : {\n" +
+                    "            \"composite\" : {\n" +
+                    "            \"size\" : 10,\n" +
+                    "                \"sources\":[{\"term1\": { \"terms\":{\"field\" : \"gender\"}}}]\n" +
+                    "            }\n" +
+                    "        }\n" +
+                    "    }\n" +
+                    "}";
+            Search search = new Search.Builder(query)
+                    .addIndex(INDEX)
+                    .addType(TYPE)
+                    .build();
+            SearchResult result = client.execute(search);
+            assertTrue(result.getErrorMessage(), result.isSucceeded());
+
+            CompositeAggregation composite = result.getAggregations().getCompositeAggregation("composite1");
+            assertEquals("composite1", composite.getName());
+            assertEquals(2, composite.getBuckets().size());
+            assertTrue(2L == composite.getBuckets().get(1).getCount());
+            assertEquals(gson.fromJson("{\"term1\":\"male\"}", JsonObject.class), composite.getBuckets().get(1).getKey());
+            assertTrue(1L == composite.getBuckets().get(0).getCount());
+            assertEquals(gson.fromJson("{\"term1\":\"female\"}", JsonObject.class), composite.getBuckets().get(0).getKey());
+
+            Aggregation aggregation = result.getAggregations().getAggregation("composite1", CompositeAggregation.class);
+            assertTrue(aggregation instanceof CompositeAggregation);
+            CompositeAggregation compositeByType = (CompositeAggregation) aggregation;
+            assertEquals(composite, compositeByType);
+
+            Map<String, Class> nameToTypeMap = new HashMap<String, Class>();
+            nameToTypeMap.put("composite1", CompositeAggregation.class);
+            List<Aggregation> aggregations = result.getAggregations().getAggregations(nameToTypeMap);
+            assertEquals(1, aggregations.size());
+            assertTrue(aggregations.get(0) instanceof CompositeAggregation);
+            CompositeAggregation termsWithMap = (CompositeAggregation) aggregations.get(0);
+            assertEquals(termsWithMap, compositeByType);
+
+            afterKey = composite.getAfterKey();
+            assertNotNull("ain't sure why", afterKey);
+        }
+
+        String query2 = "{\n" +
+                "    \"query\" : {\n" +
+                "        \"match_all\" : {}\n" +
+                "    },\n" +
+                "    \"aggs\" : {\n" +
+                "        \"composite1\" : {\n" +
+                "            \"composite\" : {\n" +
+                "            \"after\" : "+ gson.toJson(afterKey)+",\n" +
+                "                \"sources\":[{\"term1\": { \"terms\":{\"field\" : \"gender\"}}}]\n" +
+                "            }\n" +
+                "        }\n" +
+                "    }\n" +
+                "}";
+        Search search2 = new Search.Builder(query2)
+                .addIndex(INDEX)
+                .addType(TYPE)
+                .build();
+        SearchResult result2 = client.execute(search2);
+        assertTrue(result2.getErrorMessage(), result2.isSucceeded());
+
+        CompositeAggregation composite2 = result2.getAggregations().getCompositeAggregation("composite1");
+        assertEquals("composite1", composite2.getName());
+        assertEquals(0, composite2.getBuckets().size());
+        assertNull(composite2.getAfterKey());
+    }
+}