Change HLRC count request to accept a QueryBuilder

client/rest-high-level/src/main/java/org/elasticsearch/client/RequestConverters.java

@@ -505,7 +505,7 @@ static Request count(CountRequest countRequest) throws IOException {
             params.putParam("min_score", String.valueOf(countRequest.minScore()));
         }
         request.addParameters(params.asMap());
-        request.setEntity(createEntity(countRequest.source(), REQUEST_BODY_CONTENT_TYPE));
+        request.setEntity(createEntity(countRequest, REQUEST_BODY_CONTENT_TYPE));
         return request;
     }
 

client/rest-high-level/src/main/java/org/elasticsearch/client/core/CountRequest.java

@@ -24,9 +24,13 @@
 import org.elasticsearch.action.IndicesRequest;
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.xcontent.ToXContentObject;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.internal.SearchContext;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.Objects;
 
@@ -35,34 +39,43 @@
 /**
  * Encapsulates a request to _count API against one, several or all indices.
  */
-public final class CountRequest extends ActionRequest implements IndicesRequest.Replaceable {
+public final class CountRequest extends ActionRequest implements IndicesRequest.Replaceable, ToXContentObject {
 
     private String[] indices = Strings.EMPTY_ARRAY;
     private String[] types = Strings.EMPTY_ARRAY;
     private String routing;
     private String preference;
-    private SearchSourceBuilder searchSourceBuilder;
+    private QueryBuilder query;
     private IndicesOptions indicesOptions = DEFAULT_INDICES_OPTIONS;
     private int terminateAfter = SearchContext.DEFAULT_TERMINATE_AFTER;
     private Float minScore;
 
-    public CountRequest() {
-        this.searchSourceBuilder = new SearchSourceBuilder();
-    }
+    public CountRequest() {}
 
     /**
      * Constructs a new count request against the indices. No indices provided here means that count will execute on all indices.
      */
     public CountRequest(String... indices) {
-        this(indices, new SearchSourceBuilder());
+        indices(indices);
     }
 
     /**
      * Constructs a new search request against the provided indices with the given search source.
+     *
+     * @deprecated The count api only supports a query. Use {@link #CountRequest(String[], QueryBuilder)} instead.
      */
+    @Deprecated
     public CountRequest(String[] indices, SearchSourceBuilder searchSourceBuilder) {
         indices(indices);
-        this.searchSourceBuilder = searchSourceBuilder;
+        this.query = Objects.requireNonNull(searchSourceBuilder, "source must not be null").query();
+    }
+
+    /**
+     * Constructs a new search request against the provided indices with the given query.
+     */
+    public CountRequest(String[] indices, QueryBuilder query) {
+        indices(indices);
+        this.query = Objects.requireNonNull(query, "query must not be null");;
     }
 
     @Override
@@ -84,9 +97,20 @@ public CountRequest indices(String... indices) {
 
     /**
      * The source of the count request.
+     *
+     * @deprecated The count api only supports a query. Use {@link #query(QueryBuilder)} instead.
      */
+    @Deprecated
     public CountRequest source(SearchSourceBuilder searchSourceBuilder) {
-        this.searchSourceBuilder = Objects.requireNonNull(searchSourceBuilder, "source must not be null");
+        this.query = Objects.requireNonNull(searchSourceBuilder, "source must not be null").query();
+        return this;
+    }
+
+    /**
+     * Sets the query to execute for this count request.
+     */
+    public CountRequest query(QueryBuilder query) {
+        this.query = Objects.requireNonNull(query, "query must not be null");
         return this;
     }
 
@@ -188,8 +212,31 @@ public String[] types() {
         return Arrays.copyOf(this.types, this.types.length);
     }
 
+    /**
+     * @return the source builder
+     * @deprecated The count api only supports a query. Use {@link #query()} instead.
+     */
+    @Deprecated
     public SearchSourceBuilder source() {
-        return this.searchSourceBuilder;
+        return new SearchSourceBuilder().query(query);
+    }
+
+    /**
+     * @return The provided query to execute with the count request or
+     * <code>null</code> if no query was provided.
+     */
+    public QueryBuilder query() {
+        return query;
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        if (query != null) {
+            builder.field("query", query);
+        }
+        builder.endObject();
+        return builder;
     }
 
     @Override
@@ -205,12 +252,15 @@ public boolean equals(Object o) {
             Arrays.equals(indices, that.indices) &&
             Arrays.equals(types, that.types) &&
             Objects.equals(routing, that.routing) &&
-            Objects.equals(preference, that.preference);
+            Objects.equals(preference, that.preference) &&
+            Objects.equals(terminateAfter, that.terminateAfter) &&
+            Objects.equals(minScore, that.minScore) &&
+            Objects.equals(query, that.query);
     }
 
     @Override
     public int hashCode() {
-        int result = Objects.hash(indicesOptions, routing, preference);
+        int result = Objects.hash(indicesOptions, routing, preference, terminateAfter, minScore, query);
         result = 31 * result + Arrays.hashCode(indices);
         result = 31 * result + Arrays.hashCode(types);
         return result;

client/rest-high-level/src/test/java/org/elasticsearch/client/AbstractRequestTestCase.java

@@ -49,13 +49,20 @@ public final void testFromXContent() throws IOException {
 
         final XContent xContent = XContentFactory.xContent(xContentType);
         final XContentParser parser = xContent.createParser(
-            NamedXContentRegistry.EMPTY,
+            xContentRegistry(),
             LoggingDeprecationHandler.INSTANCE,
             bytes.streamInput());
         final S serverInstance = doParseToServerInstance(parser);
         assertInstances(serverInstance, clientTestInstance);
     }
 
+    /**
+     * The {@link NamedXContentRegistry} to use for this test. Subclasses may override this to have a more realistic registry.
+     */
+    protected NamedXContentRegistry xContentRegistry() {
+        return NamedXContentRegistry.EMPTY;
+    }
+
     /**
      * @return The client test instance to be serialized to xcontent as bytes
      */

client/rest-high-level/src/test/java/org/elasticsearch/client/RequestConvertersTests.java

@@ -72,6 +72,7 @@
 import org.elasticsearch.common.xcontent.json.JsonXContent;
 import org.elasticsearch.index.VersionType;
 import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.index.query.TermQueryBuilder;
 import org.elasticsearch.index.rankeval.PrecisionAtK;
@@ -1151,13 +1152,12 @@ public void testCount() throws Exception {
         setRandomCountParams(countRequest, expectedParams);
         setRandomIndicesOptions(countRequest::indicesOptions, countRequest::indicesOptions, expectedParams);
 
-        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
-        if (frequently()) {
-            if (randomBoolean()) {
-                searchSourceBuilder.minScore(randomFloat());
-            }
+        if (randomBoolean()) {
+            SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
+            countRequest.source(searchSourceBuilder);
+        } else {
+            countRequest.query(new MatchAllQueryBuilder());
         }
-        countRequest.source(searchSourceBuilder);
         Request request = RequestConverters.count(countRequest);
         StringJoiner endpoint = new StringJoiner("/", "/", "");
         String index = String.join(",", indices);
@@ -1172,7 +1172,7 @@ public void testCount() throws Exception {
         assertEquals(HttpPost.METHOD_NAME, request.getMethod());
         assertEquals(endpoint.toString(), request.getEndpoint());
         assertEquals(expectedParams, request.getParameters());
-        assertToXContentBody(searchSourceBuilder, request.getEntity());
+        assertToXContentBody(countRequest, request.getEntity());
     }
 
     public void testCountNullIndicesAndTypes() {

client/rest-high-level/src/test/java/org/elasticsearch/client/SearchIT.java

@@ -45,6 +45,7 @@
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.MatchQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.index.query.ScriptQueryBuilder;
 import org.elasticsearch.index.query.TermsQueryBuilder;
@@ -1352,19 +1353,27 @@ public void testCountOneIndexMatchQuery() throws IOException {
     }
 
     public void testCountMultipleIndicesMatchQueryUsingConstructor() throws IOException {
-
-        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().query(new MatchQueryBuilder("field", "value1"));
-        CountRequest countRequest = new CountRequest(new String[]{"index1", "index2", "index3"}, sourceBuilder);
+        CountRequest countRequest;
+        if (randomBoolean()) {
+            SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().query(new MatchQueryBuilder("field", "value1"));
+            countRequest = new CountRequest(new String[]{"index1", "index2", "index3"}, sourceBuilder);
+        } else {
+            QueryBuilder query = new MatchQueryBuilder("field", "value1");
+            countRequest = new CountRequest(new String[]{"index1", "index2", "index3"}, query);
+        }
         CountResponse countResponse = execute(countRequest, highLevelClient()::count, highLevelClient()::countAsync);
         assertCountHeader(countResponse);
         assertEquals(3, countResponse.getCount());
 
     }
 
     public void testCountMultipleIndicesMatchQuery() throws IOException {
-
         CountRequest countRequest = new CountRequest("index1", "index2", "index3");
-        countRequest.source(new SearchSourceBuilder().query(new MatchQueryBuilder("field", "value1")));
+        if (randomBoolean()) {
+            countRequest.source(new SearchSourceBuilder().query(new MatchQueryBuilder("field", "value1")));
+        } else {
+            countRequest.query(new MatchQueryBuilder("field", "value1"));
+        }
         CountResponse countResponse = execute(countRequest, highLevelClient()::count, highLevelClient()::countAsync);
         assertCountHeader(countResponse);
         assertEquals(3, countResponse.getCount());
@@ -1378,7 +1387,7 @@ public void testCountAllIndicesMatchQuery() throws IOException {
         assertCountHeader(countResponse);
         assertEquals(3, countResponse.getCount());
     }
-    
+
     public void testSearchWithBasicLicensedQuery() throws IOException {
         SearchRequest searchRequest = new SearchRequest("index");
         SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
@@ -1390,7 +1399,7 @@ public void testSearchWithBasicLicensedQuery() throws IOException {
         assertFirstHit(searchResponse, hasId("2"));
         assertSecondHit(searchResponse, hasId("1"));
     }
-    
+
     private static void assertCountHeader(CountResponse countResponse) {
         assertEquals(0, countResponse.getSkippedShards());
         assertEquals(0, countResponse.getFailedShards());

client/rest-high-level/src/test/java/org/elasticsearch/client/core/CountRequestTests.java

@@ -20,18 +20,55 @@
 package org.elasticsearch.client.core;
 
 import org.elasticsearch.action.support.IndicesOptions;
+import org.elasticsearch.client.AbstractRequestTestCase;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.ArrayUtils;
+import org.elasticsearch.common.xcontent.NamedXContentRegistry;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.MatchQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.rest.action.RestActions;
+import org.elasticsearch.search.SearchModule;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
-import org.elasticsearch.test.ESTestCase;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 
 import static org.elasticsearch.test.EqualsHashCodeTestUtils.checkEqualsAndHashCode;
+import static org.hamcrest.Matchers.equalTo;
 
-//similar to SearchRequestTests as CountRequest inline several members (and functionality) from SearchRequest
-public class CountRequestTests extends ESTestCase {
+// similar to SearchRequestTests as CountRequest inline several members (and functionality) from SearchRequest
+// In RestCountAction the request body is parsed as QueryBuilder (top level query field),
+// so that is why this is chosen as server side instance.
+public class CountRequestTests extends AbstractRequestTestCase<CountRequest, QueryBuilder> {
+
+    @Override
+    protected CountRequest createClientTestInstance() {
+        CountRequest countRequest = new CountRequest();
+        // query is the only property that is serialized as xcontent:
+        if (randomBoolean()) {
+            countRequest.query(new MatchAllQueryBuilder());
+        }
+        return countRequest;
+    }
+
+    @Override
+    protected QueryBuilder doParseToServerInstance(XContentParser parser) throws IOException {
+        return RestActions.getQueryContent(parser);
+    }
+
+    @Override
+    protected void assertInstances(QueryBuilder serverInstance, CountRequest clientTestInstance) {
+        // query is the only property that is serialized as xcontent:
+        assertThat(serverInstance, equalTo(clientTestInstance.query()));
+    }
+
+    @Override
+    protected NamedXContentRegistry xContentRegistry() {
+        return new NamedXContentRegistry(new SearchModule(Settings.EMPTY, List.of()).getNamedXContents());
+    }
 
     public void testIllegalArguments() {
         CountRequest countRequest = new CountRequest();
@@ -55,6 +92,8 @@ public void testIllegalArguments() {
         e = expectThrows(NullPointerException.class, () -> countRequest.source(null));
         assertEquals("source must not be null", e.getMessage());
 
+        e = expectThrows(NullPointerException.class, () -> countRequest.query(null));
+        assertEquals("query must not be null", e.getMessage());
     }
 
     public void testEqualsAndHashcode() {
@@ -63,7 +102,11 @@ public void testEqualsAndHashcode() {
 
     private CountRequest createCountRequest() {
         CountRequest countRequest = new CountRequest("index");
-        countRequest.source(new SearchSourceBuilder().query(new MatchQueryBuilder("num", 10)));
+        if (randomBoolean()) {
+            countRequest.source(new SearchSourceBuilder().query(new MatchQueryBuilder("num", 10)));
+        } else {
+            countRequest.query(new MatchQueryBuilder("num", 10));
+        }
         return countRequest;
     }
 
@@ -76,6 +119,10 @@ private CountRequest mutate(CountRequest countRequest) {
         mutators.add(() -> mutation.types(ArrayUtils.concat(countRequest.types(), new String[]{randomAlphaOfLength(10)})));
         mutators.add(() -> mutation.preference(randomValueOtherThan(countRequest.preference(), () -> randomAlphaOfLengthBetween(3, 10))));
         mutators.add(() -> mutation.routing(randomValueOtherThan(countRequest.routing(), () -> randomAlphaOfLengthBetween(3, 10))));
+        mutators.add(() -> mutation.terminateAfter(randomValueOtherThan(countRequest.terminateAfter(), () -> randomIntBetween(0, 10))));
+        mutators.add(() -> mutation.minScore(randomValueOtherThan(countRequest.minScore(), () -> (float) randomIntBetween(0, 10))));
+        mutators.add(() -> mutation.query(randomValueOtherThan(countRequest.query(),
+            () -> new MatchQueryBuilder(randomAlphaOfLength(4), randomAlphaOfLength(4)))));
         randomFrom(mutators).run();
         return mutation;
     }
@@ -87,9 +134,11 @@ private static CountRequest copyRequest(CountRequest countRequest) {
         result.types(countRequest.types());
         result.routing(countRequest.routing());
         result.preference(countRequest.preference());
-        if (countRequest.source() != null) {
-            result.source(countRequest.source());
+        if (countRequest.query() != null) {
+            result.query(countRequest.query());
         }
+        result.terminateAfter(countRequest.terminateAfter());
+        result.minScore(countRequest.minScore());
         return result;
     }
 }