Query refactoring: refactored IdsQuery and added test

Author: Christoph Büscher

src/main/java/org/elasticsearch/index/query/IdsQueryBuilder.java

@@ -19,45 +19,73 @@
 
 package org.elasticsearch.index.query;
 
+import com.google.common.collect.Iterables;
+
+import org.apache.lucene.queries.TermsQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.util.BytesRef;
+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;
+import org.elasticsearch.index.mapper.Uid;
+import org.elasticsearch.index.mapper.internal.UidFieldMapper;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * A query that will return only documents matching specific ids (and a type).
  */
-public class IdsQueryBuilder extends BaseQueryBuilder implements BoostableQueryBuilder<IdsQueryBuilder> {
+public class IdsQueryBuilder extends BaseQueryBuilder implements Streamable, BoostableQueryBuilder<IdsQueryBuilder> {
 
-    private final List<String> types;
+    private List<String> types = new ArrayList<>();
 
-    private List<String> values = new ArrayList<>();
+    private List<String> ids = new ArrayList<>();
 
-    private float boost = -1;
+    private float boost = 1.0f;
 
     private String queryName;
 
     public IdsQueryBuilder(String... types) {
-        this.types = types == null ? null : Arrays.asList(types);
+        this.types = (types == null || types.length == 0) ? new ArrayList<String>() : Arrays.asList(types);
+    }
+
+    /**
+     * Get the types used in this query
+     * @return the types
+     */
+    public Collection<String> types() {
+        return this.types;
     }
 
     /**
-     * Adds ids to the filter.
+     * Adds ids to the query.
      */
     public IdsQueryBuilder addIds(String... ids) {
-        values.addAll(Arrays.asList(ids));
+        this.ids.addAll(Arrays.asList(ids));
         return this;
     }
 
     /**
-     * Adds ids to the filter.
+     * Adds ids to the query.
      */
     public IdsQueryBuilder ids(String... ids) {
         return addIds(ids);
     }
 
+    /**
+     * Gets the ids for the query.
+     */
+    public Collection<String> ids() {
+        return this.ids;
+    }
+
     /**
      * Sets the boost for this query.  Documents matching this query will (in addition to the normal
      * weightings) have their score multiplied by the boost provided.
@@ -69,29 +97,43 @@ public IdsQueryBuilder boost(float boost) {
     }
 
     /**
-     * Sets the query name for the filter that can be used when searching for matched_filters per hit.
+     * Gets the boost for this query.
+     */
+    public float boost() {
+        return this.boost;
+    }
+
+    /**
+     * Sets the query name for the query that can be used when searching for matched_filters per hit.
      */
     public IdsQueryBuilder queryName(String queryName) {
         this.queryName = queryName;
         return this;
     }
 
+    /**
+     * Gets the query name for the query.
+     */
+    public String queryName() {
+        return this.queryName;
+    }
+
     @Override
     protected void doXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject(IdsQueryParser.NAME);
         if (types != null) {
             if (types.size() == 1) {
-                builder.field("type", types.get(0));
+                builder.field("type", types.iterator().next());
             } else {
                 builder.startArray("types");
-                for (Object type : types) {
+                for (String type : types) {
                     builder.value(type);
                 }
                 builder.endArray();
             }
         }
         builder.startArray("values");
-        for (Object value : values) {
+        for (String value : ids) {
             builder.value(value);
         }
         builder.endArray();
@@ -108,4 +150,73 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
     protected String parserName() {
         return IdsQueryParser.NAME;
     }
+
+    public Query toQuery(QueryParseContext parseContext) throws IOException, QueryParsingException {
+        if (this.ids.isEmpty()) {
+            return Queries.newMatchNoDocsQuery();
+        }
+
+        Collection<String> typesForQuery = this.types;
+        if (typesForQuery == null || typesForQuery.isEmpty()) {
+            typesForQuery = parseContext.queryTypes();
+        } else if (typesForQuery.size() == 1 && Iterables.getFirst(typesForQuery, null).equals("_all")) {
+            typesForQuery = parseContext.mapperService().types();
+        }
+
+        List<BytesRef> ids = new ArrayList<>(this.ids.size());
+        for (String value : this.ids) {
+            BytesRef ref = new BytesRef(value);
+            ids.add(ref);
+        }
+
+        TermsQuery query = new TermsQuery(UidFieldMapper.NAME, Uid.createTypeUids(typesForQuery, ids));
+        query.setBoost(boost);
+        if (queryName != null) {
+            parseContext.addNamedQuery(queryName, query);
+        }
+        return query;
+    }
+
+    @Override
+    public QueryValidationException validate() {
+        // all fields can be empty or null
+        return null;
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public void readFrom(StreamInput in) throws IOException {
+        this.types = (List<String>) in.readGenericValue();
+        this.ids = (List<String>) in.readGenericValue();
+        queryName = in.readOptionalString();
+        boost = in.readFloat();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeGenericValue(this.types);
+        out.writeGenericValue(this.ids);
+        out.writeOptionalString(queryName);
+        out.writeFloat(boost);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(ids, types, boost, queryName);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) {
+            return true;
+        }
+        if (obj == null || getClass() != obj.getClass()) {
+            return false;
+        }
+        IdsQueryBuilder other = (IdsQueryBuilder) obj;
+        return Objects.equals(ids, other.ids) &&
+               Objects.equals(types, other.types) &&
+               Objects.equals(boost, other.boost) &&
+               Objects.equals(queryName, other.queryName);
+    }
 }

src/main/java/org/elasticsearch/index/query/IdsQueryParser.java

@@ -20,27 +20,18 @@
 package org.elasticsearch.index.query;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Iterables;
 
-
-import org.apache.lucene.queries.TermsQuery;
-import org.apache.lucene.search.Query;
-import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.inject.Inject;
-import org.elasticsearch.common.lucene.search.Queries;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.index.mapper.Uid;
-import org.elasticsearch.index.mapper.internal.UidFieldMapper;
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
 
 /**
- *
+ * Parser for the IdsQuery.
  */
-public class IdsQueryParser extends BaseQueryParserTemp {
+public class IdsQueryParser extends BaseQueryParser {
 
     public static final String NAME = "ids";
 
@@ -53,38 +44,38 @@ public String[] names() {
         return new String[]{NAME};
     }
 
+    /**
+     * @return a QueryBuilder representation of the query passed in as XContent in the parse context
+     */
     @Override
-    public Query parse(QueryParseContext parseContext) throws IOException, QueryParsingException {
+    public QueryBuilder fromXContent(QueryParseContext parseContext) throws IOException {
         XContentParser parser = parseContext.parser();
-
-        List<BytesRef> ids = new ArrayList<>();
-        Collection<String> types = null;
-        String currentFieldName = null;
+        List<String> ids = new ArrayList<>();
+        List<String> types = new ArrayList<>();
         float boost = 1.0f;
         String queryName = null;
+
+        String currentFieldName = null;
         XContentParser.Token token;
-        boolean idsProvided = false;
         while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
             if (token == XContentParser.Token.FIELD_NAME) {
                 currentFieldName = parser.currentName();
             } else if (token == XContentParser.Token.START_ARRAY) {
                 if ("values".equals(currentFieldName)) {
-                    idsProvided = true;
                     while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) {
                         if ((token == XContentParser.Token.VALUE_STRING) ||
                                 (token == XContentParser.Token.VALUE_NUMBER)) {
-                            BytesRef value = parser.utf8BytesOrNull();
-                            if (value == null) {
+                            String id = parser.textOrNull();
+                            if (id == null) {
                                 throw new QueryParsingException(parseContext, "No value specified for term filter");
                             }
-                            ids.add(value);
+                            ids.add(id);
                         } else {
                             throw new QueryParsingException(parseContext, "Illegal value for id, expecting a string or number, got: "
                                     + token);
                         }
                     }
                 } else if ("types".equals(currentFieldName) || "type".equals(currentFieldName)) {
-                    types = new ArrayList<>();
                     while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) {
                         String value = parser.textOrNull();
                         if (value == null) {
@@ -107,26 +98,13 @@ public Query parse(QueryParseContext parseContext) throws IOException, QueryPars
                 }
             }
         }
-
-        if (!idsProvided) {
-            throw new QueryParsingException(parseContext, "[ids] query, no ids values provided");
-        }
-
         if (ids.isEmpty()) {
-            return Queries.newMatchNoDocsQuery();
-        }
-
-        if (types == null || types.isEmpty()) {
-            types = parseContext.queryTypes();
-        } else if (types.size() == 1 && Iterables.getFirst(types, null).equals("_all")) {
-            types = parseContext.mapperService().types();
-        }
-
-        TermsQuery query = new TermsQuery(UidFieldMapper.NAME, Uid.createTypeUids(types, ids));
-        query.setBoost(boost);
-        if (queryName != null) {
-            parseContext.addNamedQuery(queryName, query);
+            throw new QueryParsingException(parseContext, "[ids] query, no ids values provided");
         }
+        IdsQueryBuilder query = new IdsQueryBuilder(types.toArray(new String[types.size()]));
+        query.addIds(ids.toArray(new String[ids.size()]));
+        query.boost(boost).queryName(queryName);
+        query.validate();
         return query;
     }
 }

src/test/java/org/elasticsearch/index/query/IdsQueryBuilderTest.java

@@ -0,0 +1,74 @@
+/*
+ * 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.index.query;
+
+
+import org.apache.lucene.queries.TermsQuery;
+import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.Query;
+import org.elasticsearch.index.mapper.internal.UidFieldMapper;
+
+import java.io.IOException;
+
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.is;
+
+public class IdsQueryBuilderTest extends BaseQueryTestCase<IdsQueryBuilder> {
+
+    @Override
+    protected IdsQueryBuilder createEmptyQueryBuilder() {
+        return new IdsQueryBuilder();
+    }
+
+    @Override
+    protected void assertLuceneQuery(IdsQueryBuilder queryBuilder, Query query, QueryParseContext context) throws IOException {
+        if (testQuery.ids().size() == 0) {
+            assertThat(query, is(instanceOf(BooleanQuery.class)));
+        } else {
+            assertThat(query, is(instanceOf(TermsQuery.class)));
+            TermsQuery termQuery = (TermsQuery) query;
+            assertThat(termQuery.getBoost(), is(testQuery.boost()));
+            // because internals of TermsFilter are well hidden, check string representation
+            String[] parts = termQuery.toString().split(" ");
+            assertThat(parts.length, is(queryBuilder.ids().size() * queryBuilder.types().size()));
+            assertThat(parts[0].substring(0, parts[0].indexOf(":")), is(UidFieldMapper.NAME));
+        }
+    }
+
+    public IdsQueryBuilder createTestQueryBuilder() {
+        IdsQueryBuilder query = new IdsQueryBuilder();
+        int numberOfTypes = randomIntBetween(1, 10);
+        String[] types = new String[numberOfTypes];
+        for (int i = 0; i < numberOfTypes; i++) {
+            types[i] = randomAsciiOfLength(8);
+        }
+        query = new IdsQueryBuilder(types);
+        if (randomBoolean()) {
+            int numberOfIds = randomIntBetween(1, 10);
+            for (int i = 0; i < numberOfIds; i++) {
+                query.addIds(randomAsciiOfLength(8));
+            }
+        }
+        if (randomBoolean()) {
+            query.boost(2.0f / randomIntBetween(1, 20));
+        }
+        return query;
+    }
+}