1 files changed, 93 insertions, 39 deletions

diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java
index c493126577..8db9afe5ae 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java
@@ -21,12 +21,15 @@ package org.elasticsearch.search.suggest.completion.context;
 
 import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ObjectParser;
-import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.Objects;
 
 import static org.elasticsearch.search.suggest.completion.context.CategoryContextMapping.CONTEXT_BOOST;
 import static org.elasticsearch.search.suggest.completion.context.CategoryContextMapping.CONTEXT_PREFIX;
@@ -35,12 +38,15 @@ import static org.elasticsearch.search.suggest.completion.context.CategoryContex
 /**
  * Defines the query context for {@link CategoryContextMapping}
  */
-public final class CategoryQueryContext implements ToXContent {
-    private final CharSequence category;
+public final class CategoryQueryContext implements QueryContext {
+    public static final String NAME = "category";
+    public static final CategoryQueryContext PROTOTYPE = new CategoryQueryContext("", 1, false);
+
+    private final String category;
     private final boolean isPrefix;
     private final int boost;
 
-    private CategoryQueryContext(CharSequence category, int boost, boolean isPrefix) {
+    private CategoryQueryContext(String category, int boost, boolean isPrefix) {
         this.category = category;
         this.boost = boost;
         this.isPrefix = isPrefix;
@@ -49,7 +55,7 @@ public final class CategoryQueryContext implements ToXContent {
     /**
      * Returns the category of the context
      */
-    public CharSequence getCategory() {
+    public String getCategory() {
         return category;
     }
 
@@ -71,8 +77,81 @@ public final class CategoryQueryContext implements ToXContent {
         return new Builder();
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        CategoryQueryContext that = (CategoryQueryContext) o;
+
+        if (isPrefix != that.isPrefix) return false;
+        if (boost != that.boost) return false;
+        return category != null ? category.equals(that.category) : that.category == null;
+
+    }
+
+    @Override
+    public int hashCode() {
+        int result = category != null ? category.hashCode() : 0;
+        result = 31 * result + (isPrefix ? 1 : 0);
+        result = 31 * result + boost;
+        return result;
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    private static ObjectParser<Builder, Void> CATEGORY_PARSER = new ObjectParser<>(NAME, null);
+    static {
+        CATEGORY_PARSER.declareString(Builder::setCategory, new ParseField(CONTEXT_VALUE));
+        CATEGORY_PARSER.declareInt(Builder::setBoost, new ParseField(CONTEXT_BOOST));
+        CATEGORY_PARSER.declareBoolean(Builder::setPrefix, new ParseField(CONTEXT_PREFIX));
+    }
+
+    @Override
+    public CategoryQueryContext fromXContext(XContentParser parser) throws IOException {
+        XContentParser.Token token = parser.currentToken();
+        Builder builder = builder();
+        if (token == XContentParser.Token.START_OBJECT) {
+            CATEGORY_PARSER.parse(parser, builder);
+        } else if (token == XContentParser.Token.VALUE_STRING) {
+            builder.setCategory(parser.text());
+        } else {
+            throw new ElasticsearchParseException("category context must be an object or string");
+        }
+        return builder.build();
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        builder.field(CONTEXT_VALUE, category);
+        builder.field(CONTEXT_BOOST, boost);
+        builder.field(CONTEXT_PREFIX, isPrefix);
+        builder.endObject();
+        return builder;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeBoolean(isPrefix);
+        out.writeVInt(boost);
+        out.writeString(category);
+    }
+
+    @Override
+    public QueryContext readFrom(StreamInput in) throws IOException {
+        Builder builder = new Builder();
+        builder.isPrefix = in.readBoolean();
+        builder.boost = in.readVInt();
+        builder.category = in.readString();
+        return builder.build();
+    }
+
     public static class Builder {
-        private CharSequence category;
+        private String category;
         private boolean isPrefix = false;
         private int boost = 1;
 
@@ -80,11 +159,12 @@ public final class CategoryQueryContext implements ToXContent {
         }
 
         /**
-         * Sets the category of the context.
+         * Sets the category of the category.
          * This is a required field
          */
-        public Builder setCategory(CharSequence context) {
-            this.category = context;
+        public Builder setCategory(String category) {
+            Objects.requireNonNull(category, "category must not be null");
+            this.category = category;
             return this;
         }
 
@@ -102,42 +182,16 @@ public final class CategoryQueryContext implements ToXContent {
          * Defaults to 1.
          */
         public Builder setBoost(int boost) {
+            if (boost <= 0) {
+                throw new IllegalArgumentException("boost must be greater than 0");
+            }
             this.boost = boost;
             return this;
         }
 
         public CategoryQueryContext build() {
+            Objects.requireNonNull(category, "category must not be null");
             return new CategoryQueryContext(category, boost, isPrefix);
         }
     }
-
-    private static ObjectParser<Builder, Void> CATEGORY_PARSER = new ObjectParser<>("category", null);
-    static {
-        CATEGORY_PARSER.declareString(Builder::setCategory, new ParseField("context"));
-        CATEGORY_PARSER.declareInt(Builder::setBoost, new ParseField("boost"));
-        CATEGORY_PARSER.declareBoolean(Builder::setPrefix, new ParseField("prefix"));
-    }
-
-    public static CategoryQueryContext parse(XContentParser parser) throws IOException {
-        XContentParser.Token token = parser.currentToken();
-        Builder builder = builder();
-        if (token == XContentParser.Token.START_OBJECT) {
-            CATEGORY_PARSER.parse(parser, builder);
-        } else if (token == XContentParser.Token.VALUE_STRING) {
-            builder.setCategory(parser.text());
-        } else {
-            throw new ElasticsearchParseException("category context must be an object or string");
-        }
-        return builder.build();
-    }
-
-    @Override
-    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.startObject();
-        builder.field(CONTEXT_VALUE, category);
-        builder.field(CONTEXT_BOOST, boost);
-        builder.field(CONTEXT_PREFIX, isPrefix);
-        builder.endObject();
-        return builder;
-    }
 }