Fix inaccurate total hit count in _search/template api

modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/RestSearchTemplateAction.java

@@ -72,7 +72,6 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
             searchTemplateRequest = SearchTemplateRequest.fromXContent(parser);
         }
         searchTemplateRequest.setRequest(searchRequest);
-        RestSearchAction.checkRestTotalHits(request, searchRequest);
 
         return channel -> client.execute(SearchTemplateAction.INSTANCE, searchTemplateRequest, new RestStatusToXContentListener<>(channel));
     }

modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/TransportSearchTemplateAction.java

@@ -32,15 +32,16 @@
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.rest.action.search.RestSearchAction;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.script.TemplateScript;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.transport.TransportService;
 
-import java.io.IOException;
 import java.util.Collections;
 
 public class TransportSearchTemplateAction extends HandledTransportAction<SearchTemplateRequest, SearchTemplateResponse> {
@@ -85,13 +86,13 @@ public void onFailure(Exception t) {
             } else {
                 listener.onResponse(response);
             }
-        } catch (IOException e) {
+        } catch (Exception e) {
             listener.onFailure(e);
         }
     }
 
     static SearchRequest convert(SearchTemplateRequest searchTemplateRequest, SearchTemplateResponse response, ScriptService scriptService,
-                                 NamedXContentRegistry xContentRegistry) throws IOException {
+                                 NamedXContentRegistry xContentRegistry) throws Exception {
         Script script = new Script(searchTemplateRequest.getScriptType(),
             searchTemplateRequest.getScriptType() == ScriptType.STORED ? null : TEMPLATE_LANG, searchTemplateRequest.getScript(),
                 searchTemplateRequest.getScriptParams() == null ? Collections.emptyMap() : searchTemplateRequest.getScriptParams());
@@ -110,6 +111,19 @@ static SearchRequest convert(SearchTemplateRequest searchTemplateRequest, Search
             builder.parseXContent(parser, false);
             builder.explain(searchTemplateRequest.isExplain());
             builder.profile(searchTemplateRequest.isProfile());
+            if (searchRequest.source() == null) {
+                searchRequest.source(new SearchSourceBuilder());
+            }
+            Integer trackTotalHitsUpTo = searchRequest.source().trackTotalHitsUpTo();
+            if (trackTotalHitsUpTo != null) {
+                if (builder.trackTotalHitsUpTo() == null) {
+                    builder.trackTotalHitsUpTo(trackTotalHitsUpTo);
+                } else if (builder.trackTotalHitsUpTo() != SearchContext.TRACK_TOTAL_HITS_ACCURATE
+                    && builder.trackTotalHitsUpTo() != SearchContext.TRACK_TOTAL_HITS_DISABLED) {
+                    throw new IllegalArgumentException("[" + RestSearchAction.TOTAL_HITS_AS_INT_PARAM + "] cannot be used " +
+                        "if the tracking of total hits is not accurate, got " + trackTotalHitsUpTo);
+                }
+            }
             searchRequest.source(builder);
         }
         return searchRequest;

modules/lang-mustache/src/test/java/org/elasticsearch/script/mustache/MultiSearchTemplateIT.java

@@ -19,12 +19,15 @@
 
 package org.elasticsearch.script.mustache;
 
+import org.elasticsearch.action.bulk.BulkRequestBuilder;
 import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.index.IndexNotFoundException;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.script.ScriptType;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.test.ESIntegTestCase;
 
 import java.util.Arrays;
@@ -36,6 +39,7 @@
 import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
 import static org.hamcrest.Matchers.arrayWithSize;
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.instanceOf;
@@ -173,4 +177,73 @@ public void testBasic() throws Exception {
         assertThat(searchTemplateResponse5.getSource().utf8ToString(),
                 equalTo("{\"query\":{\"terms\":{\"group\":[1,2,3,]}}}"));
     }
+
+    public void testMultiTemplateQueryHitCount() throws Exception {
+        BulkRequestBuilder bulkRequestBuilder = client().prepareBulk();
+        for (int i =0; i <10001; i ++) {
+            bulkRequestBuilder.add(client().prepareIndex("msearch").setId(String.valueOf(i)).setSource("{\"theField\":\"foo\"}",
+                XContentType.JSON));
+        }
+        bulkRequestBuilder.get();
+        client().admin().indices().prepareRefresh().get();
+
+        MultiSearchTemplateRequest multiRequest = new MultiSearchTemplateRequest();
+
+        // Search #1
+        SearchTemplateRequest search1 = new SearchTemplateRequest();
+        search1.setRequest(new SearchRequest().indices("msearch"));
+        search1.setScriptType(ScriptType.INLINE);
+        search1.setScript("{{! ignore me }}{\"query\":{\"match_all\":{}}}");
+        multiRequest.add(search1);
+
+        // Search #2
+        // When the request parameter `rest_total_hits_as_int` is set to true, `trackTotalHits` will also be set to true,
+        // we should test that we can get an accurate hits count.
+        SearchTemplateRequest search2 = new SearchTemplateRequest();
+        SearchRequest searchRequestWithRestTotalHitsAsInt = new SearchRequest();
+        searchRequestWithRestTotalHitsAsInt.indices("msearch");
+        SearchSourceBuilder builder = new SearchSourceBuilder();
+        searchRequestWithRestTotalHitsAsInt.source(builder.trackTotalHits(true));
+        search2.setRequest(searchRequestWithRestTotalHitsAsInt);
+        search2.setScriptType(ScriptType.INLINE);
+        search2.setScript("{{! ignore me }}{\"query\":{\"match_all\":{}}}");
+        multiRequest.add(search2);
+
+
+        // Search #3
+        // When `rest_total_hits_as_int` is set to true, `trackTotalHits` should not be set to a not accurate number
+        SearchTemplateRequest search3 = new SearchTemplateRequest();
+        search3.setRequest(searchRequestWithRestTotalHitsAsInt);
+        search3.setScriptType(ScriptType.INLINE);
+        search3.setScript("{{! ignore me }}{\"query\":{\"match_all\":{}}, \"track_total_hits\":100}");
+        multiRequest.add(search3);
+
+        // Search #4
+        SearchTemplateRequest search4 = new SearchTemplateRequest();
+        search4.setRequest(new SearchRequest().indices("msearch"));
+        search4.setScriptType(ScriptType.INLINE);
+        search4.setScript("{{! ignore me }}{\"query\":{\"match_all\":{}}, \"track_total_hits\":true}");
+        multiRequest.add(search4);
+
+        MultiSearchTemplateResponse response = client().execute(MultiSearchTemplateAction.INSTANCE, multiRequest).get();
+        assertThat(response.getResponses(), arrayWithSize(4));
+        assertThat(response.getTook().millis(), greaterThan(0L));
+
+        MultiSearchTemplateResponse.Item response1 = response.getResponses()[0];
+        assertThat(response1.isFailure(), is(false));
+        assertThat(response1.getResponse().getResponse().getHits().getTotalHits().value, equalTo(10000L));
+
+        MultiSearchTemplateResponse.Item response2 = response.getResponses()[1];
+        assertThat(response2.isFailure(), is(false));
+        assertThat(response2.getResponse().getResponse().getHits().getTotalHits().value, equalTo(10001L));
+
+        MultiSearchTemplateResponse.Item response3 = response.getResponses()[2];
+        assertThat(response3.isFailure(), is(true));
+        assertThat(response3.getFailureMessage(), containsString(
+            "[rest_total_hits_as_int] cannot be used if the tracking of total hits is not accurate"));
+
+        MultiSearchTemplateResponse.Item response4 = response.getResponses()[3];
+        assertThat(response4.isFailure(), is(false));
+        assertThat(response4.getResponse().getResponse().getHits().getTotalHits().value, equalTo(10001L));
+    }
 }

modules/lang-mustache/src/test/java/org/elasticsearch/script/mustache/SearchTemplateIT.java

@@ -27,6 +27,7 @@
 import org.elasticsearch.common.xcontent.json.JsonXContent;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.script.ScriptType;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.test.ESSingleNodeTestCase;
 import org.junit.Before;
 
@@ -106,6 +107,67 @@ public void testTemplateQueryAsEscapedString() throws Exception {
         assertThat(searchResponse.getResponse().getHits().getHits().length, equalTo(1));
     }
 
+    /**
+     * Test the template query's hit count with an index which has more than 10k documents, relates to #52801.
+     */
+    public void testTemplateQueryHitCount() throws Exception {
+        BulkRequestBuilder bulkRequestBuilder = client().prepareBulk();
+        for (int i =0; i <10001; i ++) {
+            bulkRequestBuilder.add(client().prepareIndex("test").setId(String.valueOf(i)).setSource("{\"theField\":\"foo\"}",
+                XContentType.JSON));
+        }
+        bulkRequestBuilder.get();
+        client().admin().indices().prepareRefresh().get();
+
+        SearchRequest searchRequest = new SearchRequest();
+        searchRequest.indices("test");
+
+        String queryWithNoTrackTotalHits = "{\"source\" : \"{ \\\"query\\\":{\\\"match_all\\\":{}}}\"}";
+        SearchTemplateRequest request = SearchTemplateRequest.fromXContent(createParser(JsonXContent.jsonXContent,
+            queryWithNoTrackTotalHits));
+        request.setRequest(searchRequest);
+        SearchTemplateResponse searchResponse1 = client().execute(SearchTemplateAction.INSTANCE, request).get();
+        assertThat(searchResponse1.getResponse().getHits().getTotalHits().value, equalTo(10000L));
+
+        // When the request parameter `rest_total_hits_as_int` is set to true, `trackTotalHits` will also be set to true,
+        // we should test that we can get an accurate hits count.
+        SearchSourceBuilder builder = new SearchSourceBuilder();
+        searchRequest.source(builder.trackTotalHits(true));
+        request.setRequest(searchRequest);
+        SearchTemplateResponse searchResponse2 = client().execute(SearchTemplateAction.INSTANCE, request).get();
+        assertThat(searchResponse2.getResponse().getHits().getTotalHits().value, equalTo(10001L));
+
+        // When `rest_total_hits_as_int` is set to true, `trackTotalHits` should not be set to a not accurate number
+        String queryWithInvalidTrackTotalHits =
+            "{" + "  \"source\" : \"{ \\\"track_total_hits\\\": \\\"{{trackTotalHits}}\\\", \\\"query\\\":{\\\"match_all\\\":{}}}\","
+                + "  \"params\":{"
+                + "    \"trackTotalHits\": 100"
+                + "  }"
+                + "}";
+        SearchTemplateRequest requestWithInvalidTrackTotalHits = SearchTemplateRequest.fromXContent(createParser(JsonXContent.jsonXContent,
+            queryWithInvalidTrackTotalHits));
+        requestWithInvalidTrackTotalHits.setRequest(searchRequest);
+        Exception e = expectThrows(Exception.class,
+            () -> client().execute(SearchTemplateAction.INSTANCE, requestWithInvalidTrackTotalHits).get());
+        assertThat(e.getMessage(), containsString(
+            "[rest_total_hits_as_int] cannot be used if the tracking of total hits is not accurate"));
+
+        // When `rest_total_hits_as_int` is not set but `trackTotalHits` is set to true, we can also get an accurate hits count
+        String queryWithValidTrackTotalHits =
+            "{" + "  \"source\" : \"{ \\\"track_total_hits\\\": \\\"{{trackTotalHits}}\\\", \\\"query\\\":{\\\"match_all\\\":{}}}\","
+                + "  \"params\":{"
+                + "    \"trackTotalHits\": true"
+                + "  }"
+                + "}";
+        SearchTemplateRequest requestWithValidTrackTotalHits = SearchTemplateRequest.fromXContent(createParser(JsonXContent.jsonXContent,
+            queryWithValidTrackTotalHits));
+        SearchRequest newSearchRequest = new SearchRequest();
+        newSearchRequest.indices("test");
+        requestWithValidTrackTotalHits.setRequest(newSearchRequest);
+        SearchTemplateResponse searchResponse3 = client().execute(SearchTemplateAction.INSTANCE, requestWithValidTrackTotalHits).get();
+        assertThat(searchResponse3.getResponse().getHits().getTotalHits().value, equalTo(10001L));
+    }
+
     /**
      * Test that template can contain conditional clause. In this case it is at
      * the beginning of the string.