Make ScriptSortBuilder implement NamedWritable

This adds methods and tests to ScriptSortBuilder that
makes it implement NamedWritable and adds the fromXContent
method needed to read itseld from xContent.

4 files changed, 293 insertions, 13 deletions

diff --git a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
index e554eb8846..a767faa3e0 100644
--- a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
+++ b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
@@ -19,24 +19,49 @@
 
 package org.elasticsearch.search.sort;
 
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.ParseFieldMatcher;
+import org.elasticsearch.common.io.stream.NamedWriteable;
+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.QueryBuilder;
+import org.elasticsearch.index.query.QueryParseContext;
 import org.elasticsearch.script.Script;
+import org.elasticsearch.script.Script.ScriptField;
+import org.elasticsearch.script.ScriptParameterParser;
+import org.elasticsearch.script.ScriptParameterParser.ScriptParameterValue;
 
 import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
 
 /**
  * Script sort builder allows to sort based on a custom script expression.
  */
-public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
+public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> implements NamedWriteable<ScriptSortBuilder>,
+    SortElementParserTemp<ScriptSortBuilder> {
 
-    private Script script;
+    private static final String NAME = "_script";
+    static final ScriptSortBuilder PROTOTYPE = new ScriptSortBuilder(new Script("_na_"), "_na_");
+    public static final ParseField TYPE_FIELD = new ParseField("type");
+    public static final ParseField SCRIPT_FIELD = new ParseField("script");
+    public static final ParseField SORTMODE_FIELD = new ParseField("mode");
+    public static final ParseField NESTED_PATH_FIELD = new ParseField("nested_path");
+    public static final ParseField NESTED_FILTER_FIELD = new ParseField("nested_filter");
+    public static final ParseField PARAMS_FIELD = new ParseField("params");
 
+    private final Script script;
+
+    // TODO make this an enum
     private final String type;
 
+    // TODO make this an enum
     private String sortMode;
 
-    private QueryBuilder nestedFilter;
+    private QueryBuilder<?> nestedFilter;
 
     private String nestedPath;
 
@@ -45,12 +70,40 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
      *
      * @param script
      *            The script to use.
+     * @param type
+     *            The type of the script, can be either {@link ScriptSortParser#STRING_SORT_TYPE} or
+     *            {@link ScriptSortParser#NUMBER_SORT_TYPE}
      */
     public ScriptSortBuilder(Script script, String type) {
+        Objects.requireNonNull(script, "script cannot be null");
+        Objects.requireNonNull(type, "type cannot be null");
         this.script = script;
         this.type = type;
     }
 
+    ScriptSortBuilder(ScriptSortBuilder original) {
+        this.script = original.script;
+        this.type = original.type;
+        this.order = original.order;
+        this.sortMode = original.sortMode;
+        this.nestedFilter = original.nestedFilter;
+        this.nestedPath = original.nestedPath;
+    }
+
+    /**
+     * Get the script used in this sort.
+     */
+    public Script script() {
+        return this.script;
+    }
+
+    /**
+     * Get the type used in this sort.
+     */
+    public String type() {
+        return this.type;
+    }
+
     /**
      * Defines which distance to use for sorting in the case a document contains multiple geo points.
      * Possible values: min and max
@@ -61,6 +114,13 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
     }
 
     /**
+     * Get the sort mode.
+     */
+    public String sortMode() {
+        return this.sortMode;
+    }
+
+    /**
      * Sets the nested filter that the nested objects should match with in order to be taken into account
      * for sorting.
      */
@@ -70,6 +130,13 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
     }
 
     /**
+     * Gets the nested filter.
+     */
+    public QueryBuilder<?> getNestedFilter() {
+        return this.nestedFilter;
+    }
+
+    /**
      * Sets the nested path if sorting occurs on a field that is inside a nested object. For sorting by script this
      * needs to be specified.
      */
@@ -78,22 +145,149 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
         return this;
     }
 
+    /**
+     * Gets the nested path.
+     */
+    public String getNestedPath() {
+        return this.nestedPath;
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params builderParams) throws IOException {
-        builder.startObject("_script");
-        builder.field("script", script);
-        builder.field("type", type);
+        builder.startObject(NAME);
+        builder.field(SCRIPT_FIELD.getPreferredName(), script);
+        builder.field(TYPE_FIELD.getPreferredName(), type);
         builder.field(ORDER_FIELD.getPreferredName(), order);
         if (sortMode != null) {
-            builder.field("mode", sortMode);
+            builder.field(SORTMODE_FIELD.getPreferredName(), sortMode);
         }
         if (nestedPath != null) {
-            builder.field("nested_path", nestedPath);
+            builder.field(NESTED_PATH_FIELD.getPreferredName(), nestedPath);
         }
         if (nestedFilter != null) {
-            builder.field("nested_filter", nestedFilter, builderParams);
+            builder.field(NESTED_FILTER_FIELD.getPreferredName(), nestedFilter, builderParams);
         }
         builder.endObject();
         return builder;
     }
+
+    @Override
+    public ScriptSortBuilder fromXContent(QueryParseContext context, String elementName) throws IOException {
+        ScriptParameterParser scriptParameterParser = new ScriptParameterParser();
+        XContentParser parser = context.parser();
+        ParseFieldMatcher parseField = context.parseFieldMatcher();
+        Script script = null;
+        String type = null;
+        String sortMode = null;
+        SortOrder order = null;
+        QueryBuilder<?> nestedFilter = null;
+        String nestedPath = null;
+        Map<String, Object> params = new HashMap<>();
+
+        XContentParser.Token token;
+        String currentName = parser.currentName();
+        while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+            if (token == XContentParser.Token.FIELD_NAME) {
+                currentName = parser.currentName();
+            } else if (token == XContentParser.Token.START_OBJECT) {
+                if (parseField.match(currentName, ScriptField.SCRIPT)) {
+                    script = Script.parse(parser, parseField);
+                } else if (parseField.match(currentName, PARAMS_FIELD)) {
+                    params = parser.map();
+                } else if (parseField.match(currentName, NESTED_FILTER_FIELD)) {
+                    nestedFilter = context.parseInnerQueryBuilder();
+                }
+            } else if (token.isValue()) {
+                if (parseField.match(currentName, ORDER_FIELD)) {
+                    order = SortOrder.fromString(parser.text());
+                } else if (scriptParameterParser.token(currentName, token, parser, parseField)) {
+                    // Do Nothing (handled by ScriptParameterParser
+                } else if (parseField.match(currentName, TYPE_FIELD)) {
+                    type = parser.text();
+                } else if (parseField.match(currentName, SORTMODE_FIELD)) {
+                    sortMode = parser.text();
+                } else if (parseField.match(currentName, NESTED_PATH_FIELD)) {
+                    nestedPath = parser.text();
+                }
+            }
+        }
+
+        if (script == null) { // Didn't find anything using the new API so try using the old one instead
+            ScriptParameterValue scriptValue = scriptParameterParser.getDefaultScriptParameterValue();
+            if (scriptValue != null) {
+                if (params == null) {
+                    params = new HashMap<>();
+                }
+                script = new Script(scriptValue.script(), scriptValue.scriptType(), scriptParameterParser.lang(), params);
+            }
+        }
+
+        ScriptSortBuilder result = new ScriptSortBuilder(script, type);
+        if (order != null) {
+            result.order(order);
+        }
+        if (sortMode != null) {
+            result.sortMode(sortMode);
+        }
+        if (nestedFilter != null) {
+            result.setNestedFilter(nestedFilter);
+        }
+        if (nestedPath != null) {
+            result.setNestedPath(nestedPath);
+        }
+        return result;
+    }
+
+    @Override
+    public boolean equals(Object object) {
+        if (this == object) {
+            return true;
+        }
+        if (object == null || getClass() != object.getClass()) {
+            return false;
+        }
+        ScriptSortBuilder other = (ScriptSortBuilder) object;
+        return Objects.equals(script, other.script) &&
+                Objects.equals(type, other.type) &&
+                Objects.equals(order, other.order) &&
+                Objects.equals(sortMode, other.sortMode) &&
+                Objects.equals(nestedFilter, other.nestedFilter) &&
+                Objects.equals(nestedPath, other.nestedPath);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(script, type, order, sortMode, nestedFilter, nestedPath);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        script.writeTo(out);
+        out.writeString(type);
+        order.writeTo(out);
+        out.writeOptionalString(sortMode);
+        out.writeOptionalString(nestedPath);
+        boolean hasNestedFilter = nestedFilter != null;
+        out.writeBoolean(hasNestedFilter);
+        if (hasNestedFilter) {
+            out.writeQuery(nestedFilter);
+        }
+    }
+
+    @Override
+    public ScriptSortBuilder readFrom(StreamInput in) throws IOException {
+        ScriptSortBuilder builder = new ScriptSortBuilder(Script.readScript(in), in.readString());
+        builder.order(SortOrder.readOrderFrom(in));
+        builder.sortMode = in.readOptionalString();
+        builder.nestedPath = in.readOptionalString();
+        if (in.readBoolean()) {
+            builder.nestedFilter = in.readQuery();
+        }
+        return builder;
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
 }
diff --git a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortParser.java b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortParser.java
index c30ea503d8..9bf4dde711 100644
--- a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortParser.java
+++ b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortParser.java
@@ -48,8 +48,6 @@ import org.elasticsearch.script.ScriptParameterParser;
 import org.elasticsearch.script.ScriptParameterParser.ScriptParameterValue;
 import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.MultiValueMode;
-import org.elasticsearch.search.SearchParseException;
-import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -61,8 +59,8 @@ import java.util.Map;
  */
 public class ScriptSortParser implements SortParser {
 
-    private static final String STRING_SORT_TYPE = "string";
-    private static final String NUMBER_SORT_TYPE = "number";
+    public static final String STRING_SORT_TYPE = "string";
+    public static final String NUMBER_SORT_TYPE = "number";
 
     @Override
     public String[] names() {
diff --git a/core/src/test/java/org/elasticsearch/search/sort/AbstractSortTestCase.java b/core/src/test/java/org/elasticsearch/search/sort/AbstractSortTestCase.java
index dc61f0ef34..71866c34d2 100644
--- a/core/src/test/java/org/elasticsearch/search/sort/AbstractSortTestCase.java
+++ b/core/src/test/java/org/elasticsearch/search/sort/AbstractSortTestCase.java
@@ -55,6 +55,7 @@ public abstract class AbstractSortTestCase<T extends SortBuilder & NamedWriteabl
         namedWriteableRegistry = new NamedWriteableRegistry();
         namedWriteableRegistry.registerPrototype(SortBuilder.class, GeoDistanceSortBuilder.PROTOTYPE);
         namedWriteableRegistry.registerPrototype(SortBuilder.class, ScoreSortBuilder.PROTOTYPE);
+        namedWriteableRegistry.registerPrototype(SortBuilder.class, ScriptSortBuilder.PROTOTYPE);
         indicesQueriesRegistry = new SearchModule(Settings.EMPTY, namedWriteableRegistry).buildQueryParserRegistry();
     }
 
diff --git a/core/src/test/java/org/elasticsearch/search/sort/ScriptSortBuilderTests.java b/core/src/test/java/org/elasticsearch/search/sort/ScriptSortBuilderTests.java
new file mode 100644
index 0000000000..2b46858b4c
--- /dev/null
+++ b/core/src/test/java/org/elasticsearch/search/sort/ScriptSortBuilderTests.java
@@ -0,0 +1,87 @@
+/*
+ * 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.sort;
+
+
+import org.elasticsearch.script.Script;
+
+import java.io.IOException;
+
+public class ScriptSortBuilderTests extends AbstractSortTestCase<ScriptSortBuilder> {
+
+    @Override
+    protected ScriptSortBuilder createTestItem() {
+        ScriptSortBuilder builder = new ScriptSortBuilder(new Script(randomAsciiOfLengthBetween(5, 10)),
+                randomBoolean() ? ScriptSortParser.NUMBER_SORT_TYPE : ScriptSortParser.STRING_SORT_TYPE);
+        if (randomBoolean()) {
+            builder.order(RandomSortDataGenerator.order(builder.order()));
+        }
+        if (randomBoolean()) {
+            builder.sortMode(RandomSortDataGenerator.mode(builder.sortMode()));
+        }
+        if (randomBoolean()) {
+            builder.setNestedFilter(RandomSortDataGenerator.nestedFilter(builder.getNestedFilter()));
+        }
+        if (randomBoolean()) {
+            builder.setNestedPath(RandomSortDataGenerator.randomAscii(builder.getNestedPath()));
+        }
+        return builder;
+    }
+
+    @Override
+    protected ScriptSortBuilder mutate(ScriptSortBuilder original) throws IOException {
+        ScriptSortBuilder result;
+        if (randomBoolean()) {
+            // change one of the constructor args, copy the rest over
+            Script script = original.script();
+            String type = original.type();
+            if (randomBoolean()) {
+                result = new ScriptSortBuilder(new Script(script.getScript() + "_suffix"), type);
+            } else {
+                result = new ScriptSortBuilder(script, type + "_suffix");
+            }
+            result.order(original.order());
+            result.sortMode(original.sortMode());
+            result.setNestedFilter(original.getNestedFilter());
+            result.setNestedPath(original.getNestedPath());
+            return result;
+        }
+        result = new ScriptSortBuilder(original);
+        switch (randomIntBetween(0, 3)) {
+            case 0:
+                if (original.order() == SortOrder.ASC) {
+                    result.order(SortOrder.DESC);
+                } else {
+                    result.order(SortOrder.ASC);
+                }
+                break;
+            case 1:
+                result.sortMode(RandomSortDataGenerator.mode(original.sortMode()));
+                break;
+            case 2:
+                result.setNestedFilter(RandomSortDataGenerator.nestedFilter(original.getNestedFilter()));
+                break;
+            case 3:
+                result.setNestedPath(original.getNestedPath() + "_some_suffix");
+                break;
+        }
+        return result;
+    }
+}