[7.x] [ML][Transforms] add wait_for_checkpoint flag to stop (#47935)

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

@@ -41,6 +41,7 @@
 import static org.elasticsearch.client.transform.DeleteTransformRequest.FORCE;
 import static org.elasticsearch.client.transform.GetTransformRequest.ALLOW_NO_MATCH;
 import static org.elasticsearch.client.transform.PutTransformRequest.DEFER_VALIDATION;
+import static org.elasticsearch.client.transform.StopTransformRequest.WAIT_FOR_CHECKPOINT;
 
 final class TransformRequestConverters {
 
@@ -135,6 +136,9 @@ static Request stopTransform(StopTransformRequest stopRequest) {
         if (stopRequest.getAllowNoMatch() != null) {
             request.addParameter(ALLOW_NO_MATCH, stopRequest.getAllowNoMatch().toString());
         }
+        if (stopRequest.getWaitForCheckpoint() != null) {
+            request.addParameter(WAIT_FOR_CHECKPOINT, stopRequest.getWaitForCheckpoint().toString());
+        }
         request.addParameters(params.asMap());
         return request;
     }

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

@@ -28,21 +28,23 @@
 
 public class StopTransformRequest implements Validatable {
 
+    public static final String WAIT_FOR_CHECKPOINT = "wait_for_checkpoint";
+
     private final String id;
     private Boolean waitForCompletion;
+    private Boolean waitForCheckpoint;
     private TimeValue timeout;
     private Boolean allowNoMatch;
 
     public StopTransformRequest(String id) {
-        this.id = id;
-        waitForCompletion = null;
-        timeout = null;
+        this(id, null, null, null);
     }
 
-    public StopTransformRequest(String id, Boolean waitForCompletion, TimeValue timeout) {
+    public StopTransformRequest(String id, Boolean waitForCompletion, TimeValue timeout, Boolean waitForCheckpoint) {
         this.id = id;
         this.waitForCompletion = waitForCompletion;
         this.timeout = timeout;
+        this.waitForCheckpoint = waitForCheckpoint;
     }
 
     public String getId() {
@@ -73,6 +75,14 @@ public void setAllowNoMatch(Boolean allowNoMatch) {
         this.allowNoMatch = allowNoMatch;
     }
 
+    public Boolean getWaitForCheckpoint() {
+        return waitForCheckpoint;
+    }
+
+    public void setWaitForCheckpoint(Boolean waitForCheckpoint) {
+        this.waitForCheckpoint = waitForCheckpoint;
+    }
+
     @Override
     public Optional<ValidationException> validate() {
         if (id == null) {
@@ -86,7 +96,7 @@ public Optional<ValidationException> validate() {
 
     @Override
     public int hashCode() {
-        return Objects.hash(id, waitForCompletion, timeout, allowNoMatch);
+        return Objects.hash(id, waitForCompletion, timeout, allowNoMatch, waitForCheckpoint);
     }
 
     @Override
@@ -102,6 +112,7 @@ public boolean equals(Object obj) {
         return Objects.equals(this.id, other.id)
                 && Objects.equals(this.waitForCompletion, other.waitForCompletion)
                 && Objects.equals(this.timeout, other.timeout)
+                && Objects.equals(this.waitForCheckpoint, other.waitForCheckpoint)
                 && Objects.equals(this.allowNoMatch, other.allowNoMatch);
     }
 

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

@@ -148,7 +148,7 @@ private void indexData(String indexName) throws IOException {
     public void cleanUpTransforms() throws Exception {
         for (String transformId : transformsToClean) {
             highLevelClient().transform().stopTransform(
-                    new StopTransformRequest(transformId, Boolean.TRUE, null), RequestOptions.DEFAULT);
+                    new StopTransformRequest(transformId, Boolean.TRUE, null, false), RequestOptions.DEFAULT);
         }
 
         for (String transformId : transformsToClean) {
@@ -311,7 +311,7 @@ public void testStartStop() throws IOException {
         assertThat(taskState, oneOf(TransformStats.State.STARTED, TransformStats.State.INDEXING,
             TransformStats.State.STOPPING, TransformStats.State.STOPPED));
 
-        StopTransformRequest stopRequest = new StopTransformRequest(id, Boolean.TRUE, null);
+        StopTransformRequest stopRequest = new StopTransformRequest(id, Boolean.TRUE, null, false);
         StopTransformResponse stopResponse =
                 execute(stopRequest, client::stopTransform, client::stopTransformAsync);
         assertTrue(stopResponse.isAcknowledged());

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

@@ -148,7 +148,12 @@ public void testStopDataFrameTransform() {
         if (randomBoolean()) {
             timeValue = TimeValue.parseTimeValue(randomTimeValue(), "timeout");
         }
-        StopTransformRequest stopRequest = new StopTransformRequest(id, waitForCompletion, timeValue);
+        Boolean waitForCheckpoint = null;
+        if (randomBoolean()) {
+            waitForCheckpoint = randomBoolean();
+        }
+
+        StopTransformRequest stopRequest = new StopTransformRequest(id, waitForCompletion, timeValue, waitForCheckpoint);
 
         Request request = TransformRequestConverters.stopTransform(stopRequest);
         assertEquals(HttpPost.METHOD_NAME, request.getMethod());
@@ -167,6 +172,12 @@ public void testStopDataFrameTransform() {
         } else {
             assertFalse(request.getParameters().containsKey("timeout"));
         }
+        if (waitForCheckpoint != null) {
+            assertTrue(request.getParameters().containsKey("wait_for_checkpoint"));
+            assertEquals(stopRequest.getWaitForCheckpoint(), Boolean.parseBoolean(request.getParameters().get("wait_for_checkpoint")));
+        } else {
+            assertFalse(request.getParameters().containsKey("wait_for_checkpoint"));
+        }
 
         assertFalse(request.getParameters().containsKey(ALLOW_NO_MATCH));
         stopRequest.setAllowNoMatch(randomBoolean());

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

@@ -81,7 +81,7 @@ public class TransformDocumentationIT extends ESRestHighLevelClientTestCase {
     public void cleanUpTransforms() throws Exception {
         for (String transformId : transformsToClean) {
             highLevelClient().transform().stopTransform(
-                    new StopTransformRequest(transformId, Boolean.TRUE, TimeValue.timeValueSeconds(20)), RequestOptions.DEFAULT);
+                    new StopTransformRequest(transformId, true, TimeValue.timeValueSeconds(20), false), RequestOptions.DEFAULT);
         }
 
         for (String transformId : transformsToClean) {

docs/reference/transform/apis/stop-transform.asciidoc

@@ -90,6 +90,10 @@ are no matches or only partial matches.
   state completely stops. If set to `false`, the API returns immediately and the
   indexer will be stopped asynchronously in the background. Defaults to `false`.
 
+`wait_for_checkpoint`::
+  (Optional, boolean) If set to `true`, the transform will not completely stop
+  until the current checkpoint is completed. If set to `false`, the transform
+  stops as soon as possible. Defaults to `false`.
 
 [[stop-transform-response-codes]]
 ==== {api-response-codes-title}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/TransformField.java

@@ -22,6 +22,7 @@ public final class TransformField {
     public static final ParseField GROUP_BY = new ParseField("group_by");
     public static final ParseField TIMEOUT = new ParseField("timeout");
     public static final ParseField WAIT_FOR_COMPLETION = new ParseField("wait_for_completion");
+    public static final ParseField WAIT_FOR_CHECKPOINT = new ParseField("wait_for_checkpoint");
     public static final ParseField STATS_FIELD = new ParseField("stats");
     public static final ParseField INDEX_DOC_TYPE = new ParseField("doc_type");
     public static final ParseField SOURCE = new ParseField("source");

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/action/StopTransformAction.java

@@ -5,6 +5,7 @@
  */
 package org.elasticsearch.xpack.core.transform.action;
 
+import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionRequestValidationException;
@@ -32,6 +33,8 @@
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
+import static org.elasticsearch.action.ValidateActions.addValidationError;
+
 public class StopTransformAction extends ActionType<StopTransformAction.Response> {
 
     public static final StopTransformAction INSTANCE = new StopTransformAction();
@@ -48,16 +51,23 @@ public static class Request extends BaseTasksRequest<Request> {
         private final boolean waitForCompletion;
         private final boolean force;
         private final boolean allowNoMatch;
+        private final boolean waitForCheckpoint;
         private Set<String> expandedIds;
 
-        public Request(String id, boolean waitForCompletion, boolean force, @Nullable TimeValue timeout, boolean allowNoMatch) {
+        public Request(String id,
+                       boolean waitForCompletion,
+                       boolean force,
+                       @Nullable TimeValue timeout,
+                       boolean allowNoMatch,
+                       boolean waitForCheckpoint) {
             this.id = ExceptionsHelper.requireNonNull(id, TransformField.ID.getPreferredName());
             this.waitForCompletion = waitForCompletion;
             this.force = force;
 
             // use the timeout value already present in BaseTasksRequest
             this.setTimeout(timeout == null ? DEFAULT_TIMEOUT : timeout);
             this.allowNoMatch = allowNoMatch;
+            this.waitForCheckpoint = waitForCheckpoint;
         }
 
         public Request(StreamInput in) throws IOException {
@@ -73,6 +83,11 @@ public Request(StreamInput in) throws IOException {
             } else {
                 this.allowNoMatch = true;
             }
+            if (in.getVersion().onOrAfter(Version.V_7_6_0)) {
+                this.waitForCheckpoint = in.readBoolean();
+            } else {
+                this.waitForCheckpoint = false;
+            }
         }
 
         public String getId() {
@@ -99,6 +114,10 @@ public boolean isAllowNoMatch() {
             return allowNoMatch;
         }
 
+        public boolean isWaitForCheckpoint() {
+            return waitForCheckpoint;
+        }
+
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             super.writeTo(out);
@@ -113,17 +132,27 @@ public void writeTo(StreamOutput out) throws IOException {
             if (out.getVersion().onOrAfter(Version.V_7_3_0)) {
                 out.writeBoolean(allowNoMatch);
             }
+            if (out.getVersion().onOrAfter(Version.V_7_6_0)) {
+                out.writeBoolean(waitForCheckpoint);
+            }
         }
 
         @Override
         public ActionRequestValidationException validate() {
+            if (force && waitForCheckpoint) {
+                return addValidationError(new ParameterizedMessage(
+                    "cannot set both [{}] and [{}] to true",
+                        TransformField.FORCE,
+                        TransformField.WAIT_FOR_CHECKPOINT).getFormattedMessage(),
+                    null);
+            }
             return null;
         }
 
         @Override
         public int hashCode() {
             // the base class does not implement hashCode, therefore we need to hash timeout ourselves
-            return Objects.hash(id, waitForCompletion, force, expandedIds, this.getTimeout(), allowNoMatch);
+            return Objects.hash(id, waitForCompletion, force, expandedIds, this.getTimeout(), allowNoMatch, waitForCheckpoint);
         }
 
         @Override
@@ -146,6 +175,7 @@ public boolean equals(Object obj) {
                     Objects.equals(waitForCompletion, other.waitForCompletion) &&
                     Objects.equals(force, other.force) &&
                     Objects.equals(expandedIds, other.expandedIds) &&
+                    Objects.equals(waitForCheckpoint, other.waitForCheckpoint) &&
                     allowNoMatch == other.allowNoMatch;
         }