Convert suggestion response parsing to use NamedXContentRegistry (#23355)

We recently added parsing code to parse suggesters responses into java api objects. This was done using a switch based on the type of the returned suggestion. We can now replace the switch with using NamedXContentRegistry, which will also be used for aggs parsing.

7 files changed, 68 insertions, 51 deletions

diff --git a/core/src/main/java/org/elasticsearch/search/suggest/Suggest.java b/core/src/main/java/org/elasticsearch/search/suggest/Suggest.java
index ba5ad712f4..f85ea18109 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/Suggest.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/Suggest.java
@@ -19,6 +19,7 @@
 package org.elasticsearch.search.suggest;
 
 import org.apache.lucene.util.CollectionUtil;
+import org.elasticsearch.common.CheckedFunction;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -48,7 +49,6 @@ import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.common.xcontent.ConstructingObjectParser.constructorArg;
@@ -384,13 +384,14 @@ public class Suggest implements Iterable<Suggest.Suggestion<? extends Entry<? ex
             return builder;
         }
 
+        @SuppressWarnings("unchecked")
         public static Suggestion<? extends Entry<? extends Option>> fromXContent(XContentParser parser) throws IOException {
             ensureExpectedToken(XContentParser.Token.FIELD_NAME, parser.currentToken(), parser::getTokenLocation);
             String typeAndName = parser.currentName();
             // we need to extract the type prefix from the name and throw error if it is not present
             int delimiterPos = typeAndName.indexOf(InternalAggregation.TYPED_KEYS_DELIMITER);
-            String type = null;
-            String name = null;
+            String type;
+            String name;
             if (delimiterPos > 0) {
                 type = typeAndName.substring(0, delimiterPos);
                 name = typeAndName.substring(delimiterPos + 1);
@@ -399,35 +400,17 @@ public class Suggest implements Iterable<Suggest.Suggestion<? extends Entry<? ex
                         "Cannot parse suggestion response without type information. Set [" + RestSearchAction.TYPED_KEYS_PARAM
                                 + "] parameter on the request to ensure the type information is added to the response output");
             }
-            Suggestion suggestion = null;
-            Function<XContentParser, Entry> entryParser = null;
-            // the "size" parameter and the SortBy for TermSuggestion cannot be parsed from the response, use default values
-            // TODO investigate if we can use NamedXContentRegistry instead of this switch
-            switch (type) {
-                case Suggestion.NAME:
-                    suggestion = new Suggestion(name, -1);
-                    entryParser = Suggestion.Entry::fromXContent;
-                    break;
-                case PhraseSuggestion.NAME:
-                    suggestion = new PhraseSuggestion(name, -1);
-                    entryParser = PhraseSuggestion.Entry::fromXContent;
-                    break;
-                case TermSuggestion.NAME:
-                    suggestion = new TermSuggestion(name, -1, SortBy.SCORE);
-                    entryParser = TermSuggestion.Entry::fromXContent;
-                    break;
-                case CompletionSuggestion.NAME:
-                    suggestion = new CompletionSuggestion(name, -1);
-                    entryParser = CompletionSuggestion.Entry::fromXContent;
-                    break;
-                default:
-                    throw new ParsingException(parser.getTokenLocation(), "Unknown suggestion type [{}]", type);
-            }
+
+            return parser.namedObject(Suggestion.class, type, name);
+        }
+
+        protected static <E extends Suggestion.Entry<?>> void parseEntries(XContentParser parser, Suggestion<E> suggestion,
+                                                                           CheckedFunction<XContentParser, E, IOException> entryParser)
+                throws IOException {
             ensureExpectedToken(XContentParser.Token.START_ARRAY, parser.nextToken(), parser::getTokenLocation);
             while ((parser.nextToken()) != XContentParser.Token.END_ARRAY) {
                 suggestion.addTerm(entryParser.apply(parser));
             }
-            return suggestion;
         }
 
         /**
@@ -584,6 +567,7 @@ public class Suggest implements Iterable<Suggest.Suggestion<? extends Entry<? ex
                 }
             }
 
+            @SuppressWarnings("unchecked")
             protected O newOption(){
                 return (O) new Option();
             }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestion.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestion.java
index ed45104d99..229b77aad2 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestion.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestion.java
@@ -94,6 +94,12 @@ public final class CompletionSuggestion extends Suggest.Suggestion<CompletionSug
         return getOptions().size() > 0;
     }
 
+    public static CompletionSuggestion fromXContent(XContentParser parser, String name) throws IOException {
+        CompletionSuggestion suggestion = new CompletionSuggestion(name, -1);
+        parseEntries(parser, suggestion, CompletionSuggestion.Entry::fromXContent);
+        return suggestion;
+    }
+
     private static final class OptionPriorityQueue extends org.apache.lucene.util.PriorityQueue<Entry.Option> {
 
         private final Comparator<Suggest.Suggestion.Entry.Option> comparator;
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestion.java b/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestion.java
index fb646b03ae..bd6c828bd4 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestion.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestion.java
@@ -60,10 +60,13 @@ public class PhraseSuggestion extends Suggest.Suggestion<PhraseSuggestion.Entry>
         return new Entry();
     }
 
+    public static PhraseSuggestion fromXContent(XContentParser parser, String name) throws IOException {
+        PhraseSuggestion suggestion = new PhraseSuggestion(name, -1);
+        parseEntries(parser, suggestion, PhraseSuggestion.Entry::fromXContent);
+        return suggestion;
+    }
+
     public static class Entry extends Suggestion.Entry<Suggestion.Entry.Option> {
-        static class Fields {
-            static final String CUTOFF_SCORE = "cutoff_score";
-        }
 
         protected double cutoffScore = Double.MIN_VALUE;
 
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/term/TermSuggestion.java b/core/src/main/java/org/elasticsearch/search/suggest/term/TermSuggestion.java
index 31621fb63b..185e4c76b3 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/term/TermSuggestion.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/term/TermSuggestion.java
@@ -132,6 +132,13 @@ public class TermSuggestion extends Suggestion<TermSuggestion.Entry> {
         sort.writeTo(out);
     }
 
+    public static TermSuggestion fromXContent(XContentParser parser, String name) throws IOException {
+        // the "size" parameter and the SortBy for TermSuggestion cannot be parsed from the response, use default values
+        TermSuggestion suggestion = new TermSuggestion(name, -1, SortBy.SCORE);
+        parseEntries(parser, suggestion, TermSuggestion.Entry::fromXContent);
+        return suggestion;
+    }
+
     @Override
     protected Entry newEntry() {
         return new Entry();
@@ -140,8 +147,7 @@ public class TermSuggestion extends Suggestion<TermSuggestion.Entry> {
     /**
      * Represents a part from the suggest text with suggested options.
      */
-    public static class Entry extends
-            org.elasticsearch.search.suggest.Suggest.Suggestion.Entry<TermSuggestion.Entry.Option> {
+    public static class Entry extends org.elasticsearch.search.suggest.Suggest.Suggestion.Entry<TermSuggestion.Entry.Option> {
 
         public Entry(Text text, int offset, int length) {
             super(text, offset, length);
@@ -235,7 +241,7 @@ public class TermSuggestion extends Suggestion<TermSuggestion.Entry> {
                 PARSER.declareFloat(constructorArg(), Suggestion.Entry.Option.SCORE);
             }
 
-            public static final Option fromXContent(XContentParser parser) {
+            public static Option fromXContent(XContentParser parser) {
                 return PARSER.apply(parser, null);
             }
         }
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/SuggestTests.java b/core/src/test/java/org/elasticsearch/search/suggest/SuggestTests.java
index 0948b9e2ff..24e98899e8 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/SuggestTests.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/SuggestTests.java
@@ -19,8 +19,10 @@
 
 package org.elasticsearch.search.suggest;
 
+import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.text.Text;
+import org.elasticsearch.common.xcontent.NamedXContentRegistry;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.XContentType;
@@ -47,6 +49,22 @@ import static org.hamcrest.Matchers.equalTo;
 
 public class SuggestTests extends ESTestCase {
 
+    static NamedXContentRegistry getSuggestersRegistry() {
+        List<NamedXContentRegistry.Entry> namedXContents = new ArrayList<>();
+        namedXContents.add(new NamedXContentRegistry.Entry(Suggest.Suggestion.class, new ParseField("term"),
+                (parser, context) -> TermSuggestion.fromXContent(parser, (String)context)));
+        namedXContents.add(new NamedXContentRegistry.Entry(Suggest.Suggestion.class, new ParseField("phrase"),
+                (parser, context) -> PhraseSuggestion.fromXContent(parser, (String)context)));
+        namedXContents.add(new NamedXContentRegistry.Entry(Suggest.Suggestion.class, new ParseField("completion"),
+                (parser, context) -> CompletionSuggestion.fromXContent(parser, (String)context)));
+        return new NamedXContentRegistry(namedXContents);
+    }
+
+    @Override
+    protected NamedXContentRegistry xContentRegistry() {
+        return getSuggestersRegistry();
+    }
+
     public static Suggest createTestItem() {
         int numEntries = randomIntBetween(0, 5);
         List<Suggestion<? extends Entry<? extends Option>>> suggestions = new ArrayList<>();
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/SuggestionEntryTests.java b/core/src/test/java/org/elasticsearch/search/suggest/SuggestionEntryTests.java
index f1bf602ae6..07f310dfd0 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/SuggestionEntryTests.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/SuggestionEntryTests.java
@@ -23,7 +23,6 @@ import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.text.Text;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.XContentType;
-import org.elasticsearch.search.suggest.Suggest.Suggestion;
 import org.elasticsearch.search.suggest.Suggest.Suggestion.Entry;
 import org.elasticsearch.search.suggest.Suggest.Suggestion.Entry.Option;
 import org.elasticsearch.search.suggest.completion.CompletionSuggestion;
@@ -44,10 +43,8 @@ import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertToXC
 
 public class SuggestionEntryTests extends ESTestCase {
 
-    @SuppressWarnings("rawtypes")
     private static final Map<Class<? extends Entry>, Function<XContentParser, ? extends Entry>> ENTRY_PARSERS = new HashMap<>();
     static {
-        ENTRY_PARSERS.put(Suggestion.Entry.class, Suggestion.Entry::fromXContent);
         ENTRY_PARSERS.put(TermSuggestion.Entry.class, TermSuggestion.Entry::fromXContent);
         ENTRY_PARSERS.put(PhraseSuggestion.Entry.class, PhraseSuggestion.Entry::fromXContent);
         ENTRY_PARSERS.put(CompletionSuggestion.Entry.class, CompletionSuggestion.Entry::fromXContent);
@@ -61,13 +58,9 @@ public class SuggestionEntryTests extends ESTestCase {
         Text entryText = new Text(randomAsciiOfLengthBetween(5, 15));
         int offset = randomInt();
         int length = randomInt();
-        @SuppressWarnings("rawtypes")
-        Entry entry = null;
-        Supplier<Option> supplier = null;
-        if (entryType == Suggestion.Entry.class) {
-            entry = new Suggestion.Entry<>(entryText, offset, length);
-            supplier = SuggestionOptionTests::createTestItem;
-        } else if (entryType == TermSuggestion.Entry.class) {
+        Entry entry;
+        Supplier<Option> supplier;
+        if (entryType == TermSuggestion.Entry.class) {
             entry = new TermSuggestion.Entry(entryText, offset, length);
             supplier = TermSuggestionOptionTests::createTestItem;
         } else if (entryType == PhraseSuggestion.Entry.class) {
@@ -76,6 +69,8 @@ public class SuggestionEntryTests extends ESTestCase {
         } else if (entryType == CompletionSuggestion.Entry.class) {
             entry = new CompletionSuggestion.Entry(entryText, offset, length);
             supplier = CompletionSuggestionOptionTests::createTestItem;
+        } else {
+            throw new UnsupportedOperationException("entryType not supported [" + entryType + "]");
         }
         int numOptions = randomIntBetween(0, 5);
         for (int i = 0; i < numOptions; i++) {
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/SuggestionTests.java b/core/src/test/java/org/elasticsearch/search/suggest/SuggestionTests.java
index e17ead287a..259ff50e0b 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/SuggestionTests.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/SuggestionTests.java
@@ -22,6 +22,7 @@ package org.elasticsearch.search.suggest;
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.text.Text;
+import org.elasticsearch.common.xcontent.NamedXContentRegistry;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContent;
 import org.elasticsearch.common.xcontent.XContentParser;
@@ -48,15 +49,20 @@ import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertToXC
 
 public class SuggestionTests extends ESTestCase {
 
-    @SuppressWarnings({ "unchecked" })
+    @SuppressWarnings("unchecked")
     private static final Class<Suggestion<? extends Entry<? extends Option>>>[] SUGGESTION_TYPES = new Class[] {
-        Suggestion.class, TermSuggestion.class, PhraseSuggestion.class, CompletionSuggestion.class
+        TermSuggestion.class, PhraseSuggestion.class, CompletionSuggestion.class
     };
 
     public static Suggestion<? extends Entry<? extends Option>> createTestItem() {
         return createTestItem(randomFrom(SUGGESTION_TYPES));
     }
 
+    @Override
+    protected NamedXContentRegistry xContentRegistry() {
+        return SuggestTests.getSuggestersRegistry();
+    }
+
     @SuppressWarnings({ "unchecked", "rawtypes" })
     public static Suggestion<? extends Entry<? extends Option>> createTestItem(Class<? extends Suggestion> type) {
         String name = randomAsciiOfLengthBetween(5, 10);
@@ -64,10 +70,7 @@ public class SuggestionTests extends ESTestCase {
         int size = randomInt();
         Supplier<Entry> entrySupplier = null;
         Suggestion suggestion = null;
-        if (type == Suggestion.class) {
-            suggestion = new Suggestion(name, size);
-            entrySupplier = () -> SuggestionEntryTests.createTestItem(Entry.class);
-        } else if (type == TermSuggestion.class) {
+        if (type == TermSuggestion.class) {
             suggestion = new TermSuggestion(name, size, randomFrom(SortBy.values()));
             entrySupplier = () -> SuggestionEntryTests.createTestItem(TermSuggestion.Entry.class);
         } else if (type == PhraseSuggestion.class) {
@@ -76,6 +79,8 @@ public class SuggestionTests extends ESTestCase {
         } else if (type == CompletionSuggestion.class) {
             suggestion = new CompletionSuggestion(name, size);
             entrySupplier = () -> SuggestionEntryTests.createTestItem(CompletionSuggestion.Entry.class);
+        } else {
+            throw new UnsupportedOperationException("type not supported [" + type + "]");
         }
         int numEntries;
         if (frequently()) {
@@ -151,7 +156,7 @@ public class SuggestionTests extends ESTestCase {
             ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser::getTokenLocation);
             ensureExpectedToken(XContentParser.Token.FIELD_NAME, parser.nextToken(), parser::getTokenLocation);
             ParsingException e = expectThrows(ParsingException.class, () -> Suggestion.fromXContent(parser));
-            assertEquals("Unknown suggestion type [unknownType]", e.getMessage());
+            assertEquals("Unknown Suggestion [unknownType]", e.getMessage());
         }
     }