Add support for rest_total_hits_as_int in watcher (#36035)

This change adds the support for rest_total_hits_as_int
in the watcher search inputs. Setting this parameter in the request
will transform the search response to contain the total hits as
a number (instead of an object).
Note that this parameter is currently a noop since #35849 is not
merged.

Closes #36008

Author: jimczi

x-pack/plugin/watcher/src/main/java/org/elasticsearch/xpack/watcher/input/search/ExecutableSearchInput.java

@@ -29,6 +29,7 @@
 import org.elasticsearch.xpack.watcher.support.search.WatcherSearchTemplateRequest;
 import org.elasticsearch.xpack.watcher.support.search.WatcherSearchTemplateService;
 
+import java.util.Collections;
 import java.util.Map;
 
 import static org.elasticsearch.xpack.watcher.input.search.SearchInput.TYPE;
@@ -37,11 +38,12 @@
  * An input that executes search and returns the search response as the initial payload
  */
 public class ExecutableSearchInput extends ExecutableInput<SearchInput, SearchInput.Result> {
-
     public static final SearchType DEFAULT_SEARCH_TYPE = SearchType.QUERY_THEN_FETCH;
 
     private static final Logger logger = LogManager.getLogger(ExecutableSearchInput.class);
 
+    private static final Params EMPTY_PARAMS = new MapParams(Collections.emptyMap());
+
     private final Client client;
     private final WatcherSearchTemplateService searchTemplateService;
     private final TimeValue timeout;
@@ -86,7 +88,13 @@ SearchInput.Result doExecute(WatchExecutionContext ctx, WatcherSearchTemplateReq
 
         final Payload payload;
         if (input.getExtractKeys() != null) {
-            BytesReference bytes = XContentHelper.toXContent(response, XContentType.JSON, false);
+            Params params;
+            if (request.isRestTotalHitsAsint()) {
+                params = new MapParams(Collections.singletonMap("rest_total_hits_a_int", "true"));
+            } else {
+                params = EMPTY_PARAMS;
+            }
+            BytesReference bytes = XContentHelper.toXContent(response, XContentType.JSON, params, false);
             // EMPTY is safe here because we never use namedObject
             try (XContentParser parser = XContentHelper
                     .createParser(NamedXContentRegistry.EMPTY, LoggingDeprecationHandler.INSTANCE, bytes)) {

x-pack/plugin/watcher/src/main/java/org/elasticsearch/xpack/watcher/support/search/WatcherSearchTemplateRequest.java

@@ -40,8 +40,8 @@ public class WatcherSearchTemplateRequest implements ToXContentObject {
     private final SearchType searchType;
     private final IndicesOptions indicesOptions;
     private final Script template;
-
     private final BytesReference searchSource;
+    private boolean restTotalHitsAsInt;
 
     public WatcherSearchTemplateRequest(String[] indices, String[] types, SearchType searchType, IndicesOptions indicesOptions,
                                         BytesReference searchSource) {
@@ -72,6 +72,7 @@ public WatcherSearchTemplateRequest(WatcherSearchTemplateRequest original, Bytes
         this.indicesOptions = original.indicesOptions;
         this.searchSource = source;
         this.template = original.template;
+        this.restTotalHitsAsInt = original.restTotalHitsAsInt;
     }
 
     private WatcherSearchTemplateRequest(String[] indices, String[] types, SearchType searchType, IndicesOptions indicesOptions,
@@ -105,6 +106,19 @@ public IndicesOptions getIndicesOptions() {
         return indicesOptions;
     }
 
+    public boolean isRestTotalHitsAsint() {
+        return restTotalHitsAsInt;
+    }
+
+    /**
+     * Indicates whether the total hits in the response should be
+     * serialized as number (<code>true</code>) or as an object (<code>false</code>).
+     * Defaults to false.
+     */
+    public void setRestTotalHitsAsInt(boolean value) {
+        this.restTotalHitsAsInt = restTotalHitsAsInt;
+    }
+
     public BytesReference getSearchSource() {
         return searchSource;
     }
@@ -129,6 +143,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (types != null) {
             builder.array(TYPES_FIELD.getPreferredName(), types);
         }
+        if (restTotalHitsAsInt) {
+            builder.field(REST_TOTAL_HITS_AS_INT_FIELD.getPreferredName(), restTotalHitsAsInt);
+        }
         if (searchSource != null && searchSource.length() > 0) {
             try (InputStream stream = searchSource.streamInput()) {
                 builder.rawField(BODY_FIELD.getPreferredName(), stream);
@@ -167,6 +184,7 @@ public static WatcherSearchTemplateRequest fromXContent(XContentParser parser, S
         IndicesOptions indicesOptions = DEFAULT_INDICES_OPTIONS;
         BytesReference searchSource = null;
         Script template = null;
+        boolean totalHitsAsInt = false;
 
         XContentParser.Token token;
         String currentFieldName = null;
@@ -263,10 +281,19 @@ public static WatcherSearchTemplateRequest fromXContent(XContentParser parser, S
                     types.addAll(Arrays.asList(Strings.delimitedListToStringArray(typesStr, ",", " \t")));
                 } else if (SEARCH_TYPE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
                     searchType = SearchType.fromString(parser.text().toLowerCase(Locale.ROOT));
+                } else if (REST_TOTAL_HITS_AS_INT_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
+                    totalHitsAsInt = parser.booleanValue();
                 } else {
                     throw new ElasticsearchParseException("could not read search request. unexpected string field [" +
                             currentFieldName + "]");
                 }
+            } else if (token == XContentParser.Token.VALUE_BOOLEAN) {
+                if (REST_TOTAL_HITS_AS_INT_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
+                    totalHitsAsInt = parser.booleanValue();
+                } else {
+                    throw new ElasticsearchParseException("could not read search request. unexpected boolean field [" +
+                        currentFieldName + "]");
+                }
             } else {
                 throw new ElasticsearchParseException("could not read search request. unexpected token [" + token + "]");
             }
@@ -276,8 +303,10 @@ public static WatcherSearchTemplateRequest fromXContent(XContentParser parser, S
             searchSource = BytesArray.EMPTY;
         }
 
-        return new WatcherSearchTemplateRequest(indices.toArray(new String[0]), types.toArray(new String[0]), searchType,
-                indicesOptions, searchSource, template);
+        WatcherSearchTemplateRequest request = new WatcherSearchTemplateRequest(indices.toArray(new String[0]),
+            types.toArray(new String[0]), searchType, indicesOptions, searchSource, template);
+        request.setRestTotalHitsAsInt(totalHitsAsInt);
+        return request;
     }
 
     @Override
@@ -291,13 +320,14 @@ public boolean equals(Object o) {
                 Objects.equals(searchType, other.searchType) &&
                 Objects.equals(indicesOptions, other.indicesOptions) &&
                 Objects.equals(searchSource, other.searchSource) &&
-                Objects.equals(template, other.template);
+                Objects.equals(template, other.template) &&
+                Objects.equals(restTotalHitsAsInt, other.restTotalHitsAsInt);
 
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(indices, types, searchType, indicesOptions, searchSource, template);
+        return Objects.hash(indices, types, searchType, indicesOptions, searchSource, template, restTotalHitsAsInt);
     }
 
     private static final ParseField INDICES_FIELD = new ParseField("indices");
@@ -309,6 +339,7 @@ public int hashCode() {
     private static final ParseField IGNORE_UNAVAILABLE_FIELD = new ParseField("ignore_unavailable");
     private static final ParseField ALLOW_NO_INDICES_FIELD = new ParseField("allow_no_indices");
     private static final ParseField TEMPLATE_FIELD = new ParseField("template");
+    private static final ParseField REST_TOTAL_HITS_AS_INT_FIELD = new ParseField("rest_total_hits_as_int");
 
     public static final IndicesOptions DEFAULT_INDICES_OPTIONS = IndicesOptions.lenientExpandOpen();
 }

x-pack/qa/smoke-test-watcher/src/test/resources/rest-api-spec/test/mustache/30_search_input.yml

@@ -163,3 +163,129 @@ setup:
   - match: { "watch_record.result.input.search.request.body.query.bool.must.0.term.value": "val_2" }
   - match: { "watch_record.result.input.search.request.template.id": "search-template" }
   - match: { "watch_record.result.input.search.request.template.params.num": 2 }
+
+---
+"Test search input mustache integration (using request body and rest_total_hits_as_int)":
+  - skip:
+      version: " - 6.99.99"
+      reason: "rest_total_hits_as_int support was added in 7.0"
+  - do:
+      xpack.watcher.execute_watch:
+        body: >
+          {
+            "trigger_data" : {
+              "scheduled_time" : "2015-01-04T00:00:00"
+            },
+            "watch" : {
+              "trigger" : { "schedule" : { "interval" : "10s" } },
+              "actions" : {
+                "dummy" : {
+                  "logging" : {
+                    "text" : "executed!"
+                  }
+                }
+              },
+              "input" : {
+                "search" : {
+                  "request" : {
+                    "indices" : "idx",
+                    "rest_total_hits_as_int": true,
+                    "body" : {
+                      "query" : {
+                        "bool" : {
+                          "filter" : [
+                            {
+                              "range" : {
+                                "date" : {
+                                  "lte" : "{{ctx.trigger.scheduled_time}}",
+                                  "gte" : "{{ctx.trigger.scheduled_time}}||-3d"
+                                }
+                              }
+                            }
+                          ]
+                        }
+                      }
+                    }
+                  }
+                }
+              }
+            }
+          }
+  - match: { "watch_record.result.input.type": "search" }
+  - match: { "watch_record.result.input.status": "success" }
+  - match: { "watch_record.result.input.payload.hits.total": 4 }
+  # makes sure that the mustache template snippets have been resolved  correctly:
+  - match: { "watch_record.result.input.search.request.body.query.bool.filter.0.range.date.gte": "2015-01-04T00:00:00.000Z||-3d" }
+  - match: { "watch_record.result.input.search.request.body.query.bool.filter.0.range.date.lte": "2015-01-04T00:00:00.000Z" }
+
+---
+"Test search input mustache integration (using request template and rest_total_hits_as_int)":
+  - skip:
+      version: " - 6.99.99"
+      reason: "rest_total_hits_as_int support was added in 7.0"
+
+  - do:
+      put_script:
+        id: "search-template"
+        body: {
+          "script": {
+            "lang": "mustache",
+            "source": {
+              "query" : {
+                "bool" : {
+                  "must" : [
+                  {
+                    "term" : {
+                      "value" : "val_{{num}}"
+                    }
+                  }
+                  ]
+                }
+              }
+            }
+          }
+        }
+  - match: { acknowledged: true }
+
+  - do:
+      xpack.watcher.execute_watch:
+        body: >
+          {
+            "trigger_data" : {
+              "scheduled_time" : "2015-01-04T00:00:00"
+            },
+            "watch" : {
+              "trigger" : { "schedule" : { "interval" : "10s" } },
+              "actions" : {
+                "dummy" : {
+                  "logging" : {
+                    "text" : "executed!"
+                  }
+                }
+              },
+              "input" : {
+                "search" : {
+                  "request" : {
+                    "rest_total_hits_as_int": true,
+                    "indices" : "idx",
+                    "template" : {
+                      "id": "search-template",
+                      "params": {
+                        "num": 2
+                      }
+                    }
+                  }
+                }
+              }
+            }
+          }
+  - match: { "watch_record.result.input.type": "search" }
+  - match: { "watch_record.result.input.status": "success" }
+  - match: { "watch_record.result.input.payload.hits.total": 1 }
+  - match: { "watch_record.result.input.payload.hits.hits.0._id": "2" }
+  # makes sure that the mustache template snippets have been resolved  correctly:
+  - match: { "watch_record.result.input.search.request.body.query.bool.must.0.term.value": "val_2" }
+  - match: { "watch_record.result.input.search.request.template.id": "search-template" }
+  - match: { "watch_record.result.input.search.request.template.params.num": 2 }
+
+