Add stats params for Count API (#67528)

Author: kwnna15

rest-api-spec/src/main/resources/rest-api-spec/api/count.json

@@ -103,6 +103,10 @@
       "terminate_after":{
         "type":"number",
         "description":"The maximum count for each shard, upon reaching which the query execution will terminate early"
+      },
+      "stats":{
+        "type":"list",
+        "description":"Specific 'tag' of the request for logging and statistical purposes"
       }
     },
     "body":{

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/count/40_stats.yml

@@ -0,0 +1,36 @@
+setup:
+  - do:
+      indices.create:
+        index:  test
+  - do:
+      index:
+        index:  test
+        id:     "1"
+        body:   { foo: bar }
+  - do:
+      indices.refresh:
+        index: [test]
+
+---
+"count with single stat":
+  - do:
+      count:
+        index: test
+        body:
+          query:
+            match_all: {}
+          stats: ["stat_1"]
+
+  - match: {count : 1}
+
+---
+"count with multiple stats":
+  - do:
+      count:
+        index: test
+        body:
+          query:
+            match_all: {}
+          stats: ["stat_1","stat_2"]
+
+  - match: {count : 1}

server/src/main/java/org/elasticsearch/rest/action/RestActions.java

@@ -9,6 +9,8 @@
 
 package org.elasticsearch.rest.action;
 
+import java.util.Set;
+
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.FailedNodeException;
 import org.elasticsearch.action.ShardOperationFailedException;
@@ -264,6 +266,13 @@ public RestResponse buildResponse(NodesResponse response, XContentBuilder builde
      * Parses a top level query including the query element that wraps it
      */
     public static QueryBuilder getQueryContent(XContentParser parser) {
+        return getQueryContent(parser, Set.of());
+    }
+
+    /**
+     * Parses a top level query including the query element that wraps it. Ignoring explicitly supplied ignored tokens when validating.
+     */
+    public static QueryBuilder getQueryContent(XContentParser parser, Set<String> ignoredTokens) {
         try {
             QueryBuilder queryBuilder = null;
             XContentParser.Token first = parser.nextToken();
@@ -281,7 +290,7 @@ public static QueryBuilder getQueryContent(XContentParser parser) {
                     String currentName = parser.currentName();
                     if ("query".equals(currentName)) {
                         queryBuilder = parseTopLevelQuery(parser);
-                    } else {
+                    } else if (ignoredTokens.contains(currentName) == false) {
                         throw new ParsingException(parser.getTokenLocation(), "request does not support [" + parser.currentName() + "]");
                     }
                 }

server/src/main/java/org/elasticsearch/rest/action/search/RestCountAction.java

@@ -27,7 +27,9 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.List;
+import java.util.Set;
 
 import static org.elasticsearch.rest.RestRequest.Method.GET;
 import static org.elasticsearch.rest.RestRequest.Method.POST;
@@ -67,14 +69,18 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
                     searchSourceBuilder.query(queryBuilder);
                 }
             } else {
-                searchSourceBuilder.query(RestActions.getQueryContent(parser));
+                searchSourceBuilder.query(RestActions.getQueryContent(parser, Set.of("stats")));
             }
         });
         countRequest.routing(request.param("routing"));
         float minScore = request.paramAsFloat("min_score", -1f);
         if (minScore != -1f) {
             searchSourceBuilder.minScore(minScore);
         }
+        String sStats = request.param("stats");
+        if (sStats != null) {
+            searchSourceBuilder.stats(Arrays.asList(Strings.splitStringByCommaToArray(sStats)));
+        }
 
         countRequest.preference(request.param("preference"));