From 7f6a765a1efcfd5777f0cb1ad397b0fc3ec71c94 Mon Sep 17 00:00:00 2001
From: Nik Everett <nik9000@gmail.com>
Date: Thu, 14 Apr 2016 16:02:17 -0400
Subject: [PATCH] Cut the sampler and diversified_sampler aggregations to
 registerAggregation

and remove their PROTOTYPEs.

Relates to #17085
---
 .../elasticsearch/search/SearchModule.java    |  9 ++-
 .../sampler/DiversifiedAggregatorBuilder.java | 55 ++++++++------
 .../sampler/DiversifiedSamplerParser.java     | 13 ----
 .../sampler/SamplerAggregatorBuilder.java     | 67 ++++++++++++++---
 .../bucket/sampler/SamplerParser.java         | 75 -------------------
 5 files changed, 93 insertions(+), 126 deletions(-)
 delete mode 100644 core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerParser.java

diff --git a/core/src/main/java/org/elasticsearch/search/SearchModule.java b/core/src/main/java/org/elasticsearch/search/SearchModule.java
index c06fd014787..1318ac8ae1e 100644
--- a/core/src/main/java/org/elasticsearch/search/SearchModule.java
+++ b/core/src/main/java/org/elasticsearch/search/SearchModule.java
@@ -133,9 +133,10 @@ import org.elasticsearch.search.aggregations.bucket.range.geodistance.InternalGe
 import org.elasticsearch.search.aggregations.bucket.range.ipv4.IPv4RangeAggregatorBuilder;
 import org.elasticsearch.search.aggregations.bucket.range.ipv4.InternalIPv4Range;
 import org.elasticsearch.search.aggregations.bucket.range.ipv4.IpRangeParser;
+import org.elasticsearch.search.aggregations.bucket.sampler.DiversifiedAggregatorBuilder;
 import org.elasticsearch.search.aggregations.bucket.sampler.DiversifiedSamplerParser;
 import org.elasticsearch.search.aggregations.bucket.sampler.InternalSampler;
-import org.elasticsearch.search.aggregations.bucket.sampler.SamplerParser;
+import org.elasticsearch.search.aggregations.bucket.sampler.SamplerAggregatorBuilder;
 import org.elasticsearch.search.aggregations.bucket.sampler.UnmappedSampler;
 import org.elasticsearch.search.aggregations.bucket.significant.SignificantLongTerms;
 import org.elasticsearch.search.aggregations.bucket.significant.SignificantStringTerms;
@@ -459,8 +460,10 @@ public class SearchModule extends AbstractModule {
         registerAggregation(FilterAggregatorBuilder::new, FilterAggregatorBuilder::parse, FilterAggregatorBuilder.AGGREGATION_NAME_FIELD);
         registerAggregation(FiltersAggregatorBuilder::new, (n, c) -> FiltersAggregatorBuilder.parse(queryParserRegistry, n, c),
                 FiltersAggregatorBuilder.AGGREGATION_NAME_FIELD);
-        registerAggregatorParser(new SamplerParser());
-        registerAggregatorParser(new DiversifiedSamplerParser());
+        registerAggregation(SamplerAggregatorBuilder::new, SamplerAggregatorBuilder::parse,
+                SamplerAggregatorBuilder.AGGREGATION_NAME_FIELD);
+        registerAggregation(DiversifiedAggregatorBuilder::new, new DiversifiedSamplerParser(),
+                DiversifiedAggregatorBuilder.AGGREGATION_NAME_FIELD);
         registerAggregation(TermsAggregatorBuilder::new, new TermsParser(), TermsAggregatorBuilder.AGGREGATION_NAME_FIELD);
         registerAggregatorParser(new SignificantTermsParser(significanceHeuristicParserMapper, queryParserRegistry));
         registerAggregation(RangeAggregatorBuilder::new, new RangeParser(), RangeAggregatorBuilder.AGGREGATION_NAME_FIELD);
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedAggregatorBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedAggregatorBuilder.java
index d7e44ba8254..a8ff098ccc8 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedAggregatorBuilder.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedAggregatorBuilder.java
@@ -19,14 +19,14 @@
 
 package org.elasticsearch.search.aggregations.bucket.sampler;
 
+import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.XContentBuilder;
-import org.elasticsearch.search.aggregations.AggregatorFactory;
 import org.elasticsearch.search.aggregations.AggregatorFactories.Builder;
+import org.elasticsearch.search.aggregations.AggregatorFactory;
 import org.elasticsearch.search.aggregations.InternalAggregation.Type;
 import org.elasticsearch.search.aggregations.support.AggregationContext;
-import org.elasticsearch.search.aggregations.support.ValueType;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorBuilder;
 import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorFactory;
@@ -37,10 +37,9 @@ import java.io.IOException;
 import java.util.Objects;
 
 public class DiversifiedAggregatorBuilder extends ValuesSourceAggregatorBuilder<ValuesSource, DiversifiedAggregatorBuilder> {
-
-    public static final Type TYPE = new Type("diversified_sampler");
-
-    static final DiversifiedAggregatorBuilder PROTOTYPE = new DiversifiedAggregatorBuilder("");
+    public static final String NAME = "diversified_sampler";
+    public static final ParseField AGGREGATION_NAME_FIELD = new ParseField(NAME);
+    public static final Type TYPE = new Type(NAME);
 
     public static final int MAX_DOCS_PER_VALUE_DEFAULT = 1;
 
@@ -52,6 +51,28 @@ public class DiversifiedAggregatorBuilder extends ValuesSourceAggregatorBuilder<
         super(name, TYPE, ValuesSourceType.ANY, null);
     }
 
+    /**
+     * Read from a stream.
+     */
+    public DiversifiedAggregatorBuilder(StreamInput in) throws IOException {
+        super(in, TYPE, ValuesSourceType.ANY, null);
+        shardSize = in.readVInt();
+        maxDocsPerValue = in.readVInt();
+        executionHint = in.readOptionalString();
+    }
+
+    @Override
+    protected void innerWriteTo(StreamOutput out) throws IOException {
+        out.writeVInt(shardSize);
+        out.writeVInt(maxDocsPerValue);
+        out.writeOptionalString(executionHint);
+    }
+
+    @Override
+    protected boolean usesNewStyleSerialization() {
+        return true;
+    }
+
     /**
      * Set the max num docs to be returned from each shard.
      */
@@ -122,23 +143,6 @@ public class DiversifiedAggregatorBuilder extends ValuesSourceAggregatorBuilder<
         return builder;
     }
 
-    @Override
-    protected DiversifiedAggregatorBuilder innerReadFrom(String name, ValuesSourceType valuesSourceType,
-            ValueType targetValueType, StreamInput in) throws IOException {
-        DiversifiedAggregatorBuilder factory = new DiversifiedAggregatorBuilder(name);
-        factory.shardSize = in.readVInt();
-        factory.maxDocsPerValue = in.readVInt();
-        factory.executionHint = in.readOptionalString();
-        return factory;
-    }
-
-    @Override
-    protected void innerWriteTo(StreamOutput out) throws IOException {
-        out.writeVInt(shardSize);
-        out.writeVInt(maxDocsPerValue);
-        out.writeOptionalString(executionHint);
-    }
-
     @Override
     protected int innerHashCode() {
         return Objects.hash(shardSize, maxDocsPerValue, executionHint);
@@ -151,4 +155,9 @@ public class DiversifiedAggregatorBuilder extends ValuesSourceAggregatorBuilder<
                 && Objects.equals(maxDocsPerValue, other.maxDocsPerValue)
                 && Objects.equals(executionHint, other.executionHint);
     }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
 }
\ No newline at end of file
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedSamplerParser.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedSamplerParser.java
index 79ac7ba7ed5..cb87e53f2c0 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedSamplerParser.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedSamplerParser.java
@@ -22,7 +22,6 @@ package org.elasticsearch.search.aggregations.bucket.sampler;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.ParseFieldMatcher;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.search.aggregations.AggregatorBuilder;
 import org.elasticsearch.search.aggregations.support.AbstractValuesSourceParser.AnyValuesSourceParser;
 import org.elasticsearch.search.aggregations.support.ValueType;
 import org.elasticsearch.search.aggregations.support.ValuesSourceType;
@@ -34,16 +33,10 @@ import java.util.Map;
  *
  */
 public class DiversifiedSamplerParser extends AnyValuesSourceParser {
-
     public DiversifiedSamplerParser() {
         super(true, false);
     }
 
-    @Override
-    public String type() {
-        return DiversifiedAggregatorBuilder.TYPE.name();
-    }
-
     @Override
     protected DiversifiedAggregatorBuilder createFactory(String aggregationName, ValuesSourceType valuesSourceType,
             ValueType targetValueType, Map<ParseField, Object> otherOptions) {
@@ -85,10 +78,4 @@ public class DiversifiedSamplerParser extends AnyValuesSourceParser {
         }
         return false;
     }
-
-    @Override
-    public AggregatorBuilder<?> getFactoryPrototypes() {
-        return DiversifiedAggregatorBuilder.PROTOTYPE;
-    }
-
 }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerAggregatorBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerAggregatorBuilder.java
index d00ea80ae29..2f8a3fc29f5 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerAggregatorBuilder.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerAggregatorBuilder.java
@@ -19,20 +19,24 @@
 
 package org.elasticsearch.search.aggregations.bucket.sampler;
 
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.index.query.QueryParseContext;
 import org.elasticsearch.search.aggregations.AggregatorBuilder;
-import org.elasticsearch.search.aggregations.AggregatorFactory;
 import org.elasticsearch.search.aggregations.AggregatorFactories.Builder;
+import org.elasticsearch.search.aggregations.AggregatorFactory;
 import org.elasticsearch.search.aggregations.support.AggregationContext;
 
 import java.io.IOException;
 import java.util.Objects;
 
 public class SamplerAggregatorBuilder extends AggregatorBuilder<SamplerAggregatorBuilder> {
-
-    static final SamplerAggregatorBuilder PROTOTYPE = new SamplerAggregatorBuilder("");
+    public static final String NAME = InternalSampler.TYPE.name();
+    public static final ParseField AGGREGATION_NAME_FIELD = new ParseField(NAME);
 
     public static final int DEFAULT_SHARD_SAMPLE_SIZE = 100;
 
@@ -42,6 +46,24 @@ public class SamplerAggregatorBuilder extends AggregatorBuilder<SamplerAggregato
         super(name, InternalSampler.TYPE);
     }
 
+    /**
+     * Read from a stream.
+     */
+    public SamplerAggregatorBuilder(StreamInput in) throws IOException {
+        super(in, InternalSampler.TYPE);
+        shardSize = in.readVInt();
+    }
+
+    @Override
+    protected void doWriteTo(StreamOutput out) throws IOException {
+        out.writeVInt(shardSize);
+    }
+
+    @Override
+    protected boolean usesNewStyleSerialization() {
+        return true;
+    }
+
     /**
      * Set the max num docs to be returned from each shard.
      */
@@ -71,16 +93,33 @@ public class SamplerAggregatorBuilder extends AggregatorBuilder<SamplerAggregato
         return builder;
     }
 
-    @Override
-    protected SamplerAggregatorBuilder doReadFrom(String name, StreamInput in) throws IOException {
-        SamplerAggregatorBuilder factory = new SamplerAggregatorBuilder(name);
-        factory.shardSize = in.readVInt();
-        return factory;
-    }
+    public static SamplerAggregatorBuilder parse(String aggregationName, QueryParseContext context) throws IOException {
+        XContentParser.Token token;
+        String currentFieldName = null;
+        Integer shardSize = null;
 
-    @Override
-    protected void doWriteTo(StreamOutput out) throws IOException {
-        out.writeVInt(shardSize);
+        XContentParser parser = context.parser();
+        while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+            if (token == XContentParser.Token.FIELD_NAME) {
+                currentFieldName = parser.currentName();
+            } else if (token == XContentParser.Token.VALUE_NUMBER) {
+                if (context.getParseFieldMatcher().match(currentFieldName, SamplerAggregator.SHARD_SIZE_FIELD)) {
+                    shardSize = parser.intValue();
+                } else {
+                    throw new ParsingException(parser.getTokenLocation(),
+                            "Unsupported property \"" + currentFieldName + "\" for aggregation \"" + aggregationName);
+                }
+            } else {
+                throw new ParsingException(parser.getTokenLocation(),
+                        "Unsupported property \"" + currentFieldName + "\" for aggregation \"" + aggregationName);
+            }
+        }
+
+        SamplerAggregatorBuilder factory = new SamplerAggregatorBuilder(aggregationName);
+        if (shardSize != null) {
+            factory.shardSize(shardSize);
+        }
+        return factory;
     }
 
     @Override
@@ -94,4 +133,8 @@ public class SamplerAggregatorBuilder extends AggregatorBuilder<SamplerAggregato
         return Objects.equals(shardSize, other.shardSize);
     }
 
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
 }
\ No newline at end of file
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerParser.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerParser.java
deleted file mode 100644
index 22efaf4f28c..00000000000
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerParser.java
+++ /dev/null
@@ -1,75 +0,0 @@
-/*
- * 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.sampler;
-
-
-import org.elasticsearch.common.ParsingException;
-import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.index.query.QueryParseContext;
-import org.elasticsearch.search.aggregations.Aggregator;
-
-import java.io.IOException;
-
-/**
- *
- */
-public class SamplerParser implements Aggregator.Parser {
-
-    @Override
-    public String type() {
-        return InternalSampler.TYPE.name();
-    }
-
-    @Override
-    public SamplerAggregatorBuilder parse(String aggregationName, QueryParseContext context) throws IOException {
-
-        XContentParser.Token token;
-        String currentFieldName = null;
-        Integer shardSize = null;
-
-        XContentParser parser = context.parser();
-        while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
-            if (token == XContentParser.Token.FIELD_NAME) {
-                currentFieldName = parser.currentName();
-            } else if (token == XContentParser.Token.VALUE_NUMBER) {
-                if (context.getParseFieldMatcher().match(currentFieldName, SamplerAggregator.SHARD_SIZE_FIELD)) {
-                    shardSize = parser.intValue();
-                } else {
-                    throw new ParsingException(parser.getTokenLocation(),
-                            "Unsupported property \"" + currentFieldName + "\" for aggregation \"" + aggregationName);
-                }
-            } else {
-                throw new ParsingException(parser.getTokenLocation(),
-                        "Unsupported property \"" + currentFieldName + "\" for aggregation \"" + aggregationName);
-            }
-        }
-
-        SamplerAggregatorBuilder factory = new SamplerAggregatorBuilder(aggregationName);
-        if (shardSize != null) {
-            factory.shardSize(shardSize);
-        }
-        return factory;
-    }
-
-    @Override
-    public SamplerAggregatorBuilder getFactoryPrototypes() {
-        return SamplerAggregatorBuilder.PROTOTYPE;
-    }
-
-}