Add parsing method for binary range aggregation (#24706)

Related to #23331

2 files changed, 169 insertions, 0 deletions

diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/range/InternalBinaryRange.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/range/InternalBinaryRange.java
index 640b1cfb46..60fb5b99fb 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/range/InternalBinaryRange.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/range/InternalBinaryRange.java
@@ -44,6 +44,7 @@ import static java.util.Collections.unmodifiableList;
 public final class InternalBinaryRange
         extends InternalMultiBucketAggregation<InternalBinaryRange, InternalBinaryRange.Bucket>
         implements Range {
+
     public static class Bucket extends InternalMultiBucketAggregation.InternalBucket implements Range.Bucket {
 
         private final transient DocValueFormat format;
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/range/ParsedBinaryRange.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/range/ParsedBinaryRange.java
new file mode 100644
index 0000000000..760bd23c09
--- /dev/null
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/range/ParsedBinaryRange.java
@@ -0,0 +1,168 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.search.aggregations.bucket.range;
+
+import org.elasticsearch.common.xcontent.ObjectParser;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.common.xcontent.XContentParserUtils;
+import org.elasticsearch.search.aggregations.Aggregation;
+import org.elasticsearch.search.aggregations.Aggregations;
+import org.elasticsearch.search.aggregations.ParsedMultiBucketAggregation;
+import org.elasticsearch.search.aggregations.bucket.range.ip.IpRangeAggregationBuilder;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.elasticsearch.common.xcontent.XContentParserUtils.ensureExpectedToken;
+
+public class ParsedBinaryRange extends ParsedMultiBucketAggregation<ParsedBinaryRange.ParsedBucket> implements Range {
+
+    @Override
+    public String getType() {
+        return IpRangeAggregationBuilder.NAME;
+    }
+
+    @Override
+    public List<? extends Range.Bucket> getBuckets() {
+        return buckets;
+    }
+
+    private static ObjectParser<ParsedBinaryRange, Void> PARSER =
+            new ObjectParser<>(ParsedBinaryRange.class.getSimpleName(), true, ParsedBinaryRange::new);
+    static {
+        declareMultiBucketAggregationFields(PARSER,
+                parser -> ParsedBucket.fromXContent(parser, false),
+                parser -> ParsedBucket.fromXContent(parser, true));
+    }
+
+    public static ParsedBinaryRange fromXContent(XContentParser parser, String name) throws IOException {
+        ParsedBinaryRange aggregation = PARSER.parse(parser, null);
+        aggregation.setName(name);
+        return aggregation;
+    }
+
+    public static class ParsedBucket extends ParsedMultiBucketAggregation.ParsedBucket implements Range.Bucket {
+
+        private String key;
+        private String from;
+        private String to;
+
+        @Override
+        public Object getKey() {
+            return key;
+        }
+
+        @Override
+        public String getKeyAsString() {
+            return key;
+        }
+
+        @Override
+        public Object getFrom() {
+            return from;
+        }
+
+        @Override
+        public String getFromAsString() {
+            return from;
+        }
+
+        @Override
+        public Object getTo() {
+            return to;
+        }
+
+        @Override
+        public String getToAsString() {
+            return to;
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            if (isKeyed()) {
+                builder.startObject(key != null ? key : rangeKey(from, to));
+            } else {
+                builder.startObject();
+                if (key != null) {
+                    builder.field(CommonFields.KEY.getPreferredName(), key);
+                }
+            }
+            if (from != null) {
+                builder.field(CommonFields.FROM.getPreferredName(), getFrom());
+            }
+            if (to != null) {
+                builder.field(CommonFields.TO.getPreferredName(), getTo());
+            }
+            builder.field(CommonFields.DOC_COUNT.getPreferredName(), getDocCount());
+            getAggregations().toXContentInternal(builder, params);
+            builder.endObject();
+            return builder;
+        }
+
+        static ParsedBucket fromXContent(final XContentParser parser, final boolean keyed) throws IOException {
+            final ParsedBucket bucket = new ParsedBucket();
+            bucket.setKeyed(keyed);
+            XContentParser.Token token = parser.currentToken();
+            String currentFieldName = parser.currentName();
+
+            String rangeKey = null;
+            if (keyed) {
+                ensureExpectedToken(XContentParser.Token.FIELD_NAME, token, parser::getTokenLocation);
+                rangeKey = currentFieldName;
+                ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser::getTokenLocation);
+            }
+
+            List<Aggregation> aggregations = new ArrayList<>();
+            while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+                if (token == XContentParser.Token.FIELD_NAME) {
+                    currentFieldName = parser.currentName();
+                } else if (token.isValue()) {
+                    if (CommonFields.KEY.getPreferredName().equals(currentFieldName)) {
+                        bucket.key = parser.text();
+                    } else if (CommonFields.DOC_COUNT.getPreferredName().equals(currentFieldName)) {
+                        bucket.setDocCount(parser.longValue());
+                    } else if (CommonFields.FROM.getPreferredName().equals(currentFieldName)) {
+                        bucket.from = parser.text();
+                    } else if (CommonFields.TO.getPreferredName().equals(currentFieldName)) {
+                        bucket.to = parser.text();
+                    }
+                } else if (token == XContentParser.Token.START_OBJECT) {
+                    aggregations.add(XContentParserUtils.parseTypedKeysObject(parser, Aggregation.TYPED_KEYS_DELIMITER, Aggregation.class));
+                }
+            }
+            bucket.setAggregations(new Aggregations(aggregations));
+
+            if (keyed) {
+                if (rangeKey(bucket.from, bucket.to).equals(rangeKey)) {
+                    bucket.key = null;
+                } else {
+                    bucket.key = rangeKey;
+                }
+            }
+            return bucket;
+        }
+
+        private static String rangeKey(String from, String to) {
+            return (from == null ? "*" : from) + '-' + (to == null ? "*" : to);
+        }
+    }
+}