Add to serialization of nested queries

elastic · May 13, 2015 · 0e56b75 · 0e56b75
1 parent 1715aac
commit 0e56b75
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 18 deletions.
diff --git a/src/main/java/org/elasticsearch/index/query/BaseQueryBuilder.java b/src/main/java/org/elasticsearch/index/query/BaseQueryBuilder.java
@@ -30,13 +30,22 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Base class with common code for all {@link QueryBuilder} implementations.
  */
 public abstract class BaseQueryBuilder implements QueryBuilder {
 
+    private final static Map<String, Class<?>> queryBuilderRegistry = new HashMap<>();
+    static {
+        queryBuilderRegistry.put(MatchAllQueryParser.NAME, MatchAllQueryBuilder.class);
+        queryBuilderRegistry.put(TermQueryParser.NAME, TermQueryBuilder.class);
+        queryBuilderRegistry.put(IdsQueryParser.NAME, IdsQueryBuilder.class);
+    }
+
     @Override
     public String toString() {
         try {
@@ -77,6 +86,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
      * Temporary default implementation for toQuery that parses the query using its query parser
      */
     //norelease to be removed once all query builders override toQuery providing their own specific implementation.
+    @Override
     public Query toQuery(QueryParseContext parseContext) throws QueryParsingException, IOException {
         return parseContext.indexQueryParserService().queryParser(parserName()).parse(parseContext);
     }
@@ -107,13 +117,18 @@ public void writeTo(StreamOutput out) throws IOException {
     }
 
     protected static QueryBuilder readQueryBuilderFrom(StreamInput in) throws IOException {
+        QueryBuilder newBuilder;
         String name = in.readString();
-        QueryBuilder newBuilder = null;
-        switch (name) {
-            case MatchAllQueryParser.NAME: newBuilder = new MatchAllQueryBuilder(); break;
-        }
-        if (newBuilder != null) {
+        Class<?> newBuilderClass = queryBuilderRegistry.get(name);
+        if (newBuilderClass != null) {
+            try {
+                newBuilder = (QueryBuilder) newBuilderClass.newInstance();
+            } catch (InstantiationException | IllegalAccessException e) {
+                throw new IOException(e);
+            }
             newBuilder.readFrom(in);
+        } else {
+            throw new IOException("query with name '" + name + "' not registered, could not be deserialized.");
         }
     return newBuilder;
 }

diff --git a/src/main/java/org/elasticsearch/index/query/BoolQueryBuilder.java b/src/main/java/org/elasticsearch/index/query/BoolQueryBuilder.java
@@ -26,7 +26,6 @@
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.common.io.stream.Streamable;
 import org.elasticsearch.common.lucene.search.Queries;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 
@@ -40,7 +39,7 @@
 /**
  * A Query that matches documents matching boolean combinations of other queries.
  */
-public class BoolQueryBuilder extends BaseQueryBuilder implements Streamable, BoostableQueryBuilder<BoolQueryBuilder> {
+public class BoolQueryBuilder extends BaseQueryBuilder implements BoostableQueryBuilder<BoolQueryBuilder> {
 
     private List<QueryBuilder> mustClauses = new ArrayList<>();
 
@@ -124,7 +123,7 @@ public BoolQueryBuilder should(QueryBuilder queryBuilder) {
      * @see #minimumNumberShouldMatch(int)
      */
     public BoolQueryBuilder should(List<QueryBuilder> queryBuilders) {
-        mustClauses.addAll(queryBuilders);
+        shouldClauses.addAll(queryBuilders);
         return this;
     }
 

diff --git a/src/test/java/org/elasticsearch/index/query/BaseQueryTestCase.java b/src/test/java/org/elasticsearch/index/query/BaseQueryTestCase.java
@@ -20,6 +20,7 @@
 package org.elasticsearch.index.query;
 
 import com.carrotsearch.randomizedtesting.annotations.Repeat;
+
 import org.apache.lucene.search.Query;
 import org.elasticsearch.Version;
 import org.elasticsearch.cluster.ClusterService;
@@ -30,7 +31,6 @@
 import org.elasticsearch.common.inject.util.Providers;
 import org.elasticsearch.common.io.stream.BytesStreamInput;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
-import org.elasticsearch.common.io.stream.Streamable;
 import org.elasticsearch.common.settings.ImmutableSettings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.settings.SettingsModule;
@@ -59,10 +59,10 @@
 
 import java.io.IOException;
 
-import static org.hamcrest.Matchers.*;
+import static org.hamcrest.Matchers.is;
 
 @Ignore
-public abstract class BaseQueryTestCase<QB extends BaseQueryBuilder & Streamable> extends ElasticsearchTestCase {
+public abstract class BaseQueryTestCase<QB extends BaseQueryBuilder> extends ElasticsearchTestCase {
 
     private static Injector injector;
     private static IndexQueryParserService queryParserService;

diff --git a/src/test/java/org/elasticsearch/index/query/BoolQueryBuilderTest.java b/src/test/java/org/elasticsearch/index/query/BoolQueryBuilderTest.java
@@ -19,8 +19,6 @@
 
 package org.elasticsearch.index.query;
 
-import com.carrotsearch.randomizedtesting.annotations.Repeat;
-
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
@@ -31,7 +29,6 @@
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
 
-@Repeat(iterations=20)
 public class BoolQueryBuilderTest extends BaseQueryTestCase<BoolQueryBuilder> {
 
     @Override
@@ -92,10 +89,10 @@ protected BoolQueryBuilder createTestQueryBuilder() {
 
     private static QueryBuilder createRandomQuery() {
         QueryBuilder query = null;
-        switch (randomIntBetween(0, 0)) {
-            case 0:  query = new MatchAllQueryBuilderTest().createTestQueryBuilder(); break;
-//            case 1: return TermQueryBuilderTest.createTestQuery();
-//            case 2: return IdsQueryBuilderTest.createTestQuery();
+        switch (randomIntBetween(0, 2)) {
+            case 0:  return new MatchAllQueryBuilderTest().createTestQueryBuilder();
+            case 1: return new TermQueryBuilderTest().createTestQueryBuilder();
+            case 2: return new IdsQueryBuilderTest().createTestQueryBuilder();
         }
         return query;
     }