From 323a3bf339af2d786fb7a1827a521c95db227dae Mon Sep 17 00:00:00 2001
From: turackangal <kangal.turac@gmail.com>
Date: Sun, 8 Sep 2019 23:08:02 +0300
Subject: [PATCH] [HLRC] Send terminate_after as query string parameter to the
 count API

Prior to this commit terminate_after was sent as request body parameter
(via SearchSourceBuilder), which is not possible in the count api.

Closes #46446
---
 .../java/org/elasticsearch/client/RequestConverters.java | 7 +++++++
 .../java/org/elasticsearch/client/core/CountRequest.java | 9 +++++++--
 .../org/elasticsearch/client/RequestConvertersTests.java | 4 ++++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/client/rest-high-level/src/main/java/org/elasticsearch/client/RequestConverters.java b/client/rest-high-level/src/main/java/org/elasticsearch/client/RequestConverters.java
index 5f21ea6a2f1f2..2bdd8124de28b 100644
--- a/client/rest-high-level/src/main/java/org/elasticsearch/client/RequestConverters.java
+++ b/client/rest-high-level/src/main/java/org/elasticsearch/client/RequestConverters.java
@@ -498,6 +498,9 @@ static Request count(CountRequest countRequest) throws IOException {
         params.withRouting(countRequest.routing());
         params.withPreference(countRequest.preference());
         params.withIndicesOptions(countRequest.indicesOptions());
+        if (countRequest.terminateAfter() != 0){
+            params.withTerminateAfter(countRequest.terminateAfter());
+        }
         request.addParameters(params.asMap());
         request.setEntity(createEntity(countRequest.source(), REQUEST_BODY_CONTENT_TYPE));
         return request;
@@ -910,6 +913,10 @@ Params withStoredFields(String[] storedFields) {
             return this;
         }
 
+        Params withTerminateAfter(int terminateAfter){
+            return putParam("terminate_after", String.valueOf(terminateAfter));
+        }
+
         Params withTimeout(TimeValue timeout) {
             return putParam("timeout", timeout);
         }
diff --git a/client/rest-high-level/src/main/java/org/elasticsearch/client/core/CountRequest.java b/client/rest-high-level/src/main/java/org/elasticsearch/client/core/CountRequest.java
index 4b40fde53e954..75516288255fc 100644
--- a/client/rest-high-level/src/main/java/org/elasticsearch/client/core/CountRequest.java
+++ b/client/rest-high-level/src/main/java/org/elasticsearch/client/core/CountRequest.java
@@ -25,6 +25,7 @@
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.internal.SearchContext;
 
 import java.util.Arrays;
 import java.util.Objects;
@@ -42,6 +43,7 @@ public final class CountRequest extends ActionRequest implements IndicesRequest.
     private String preference;
     private SearchSourceBuilder searchSourceBuilder;
     private IndicesOptions indicesOptions = DEFAULT_INDICES_OPTIONS;
+    private int terminateAfter = SearchContext.DEFAULT_TERMINATE_AFTER;
 
     public CountRequest() {
         this.searchSourceBuilder = new SearchSourceBuilder();
@@ -165,11 +167,14 @@ public CountRequest minScore(Float minScore) {
     }
 
     public int terminateAfter() {
-        return this.searchSourceBuilder.terminateAfter();
+        return this.terminateAfter;
     }
 
     public CountRequest terminateAfter(int terminateAfter) {
-        this.searchSourceBuilder.terminateAfter(terminateAfter);
+        if (terminateAfter < 0) {
+            throw new IllegalArgumentException("terminateAfter must be > 0");
+        }
+        this.terminateAfter = terminateAfter;
         return this;
     }
 
diff --git a/client/rest-high-level/src/test/java/org/elasticsearch/client/RequestConvertersTests.java b/client/rest-high-level/src/test/java/org/elasticsearch/client/RequestConvertersTests.java
index 2291c24eeef7d..0682ce8cbc046 100644
--- a/client/rest-high-level/src/test/java/org/elasticsearch/client/RequestConvertersTests.java
+++ b/client/rest-high-level/src/test/java/org/elasticsearch/client/RequestConvertersTests.java
@@ -1191,6 +1191,10 @@ private static void setRandomCountParams(CountRequest countRequest,
             countRequest.preference(randomAlphaOfLengthBetween(3, 10));
             expectedParams.put("preference", countRequest.preference());
         }
+        if (randomBoolean()){
+            countRequest.terminateAfter(randomIntBetween(0, Integer.MAX_VALUE));
+            expectedParams.put("terminate_after", String.valueOf(countRequest.terminateAfter()));
+        }
     }
 
     public void testMultiSearch() throws IOException {