ML: Add upgrade mode docs, hlrc, and fix bug

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

@@ -64,6 +64,7 @@
 import org.elasticsearch.client.ml.PutFilterRequest;
 import org.elasticsearch.client.ml.PutJobRequest;
 import org.elasticsearch.client.ml.RevertModelSnapshotRequest;
+import org.elasticsearch.client.ml.SetUpgradeModeRequest;
 import org.elasticsearch.client.ml.StartDatafeedRequest;
 import org.elasticsearch.client.ml.StopDatafeedRequest;
 import org.elasticsearch.client.ml.UpdateDatafeedRequest;
@@ -624,6 +625,17 @@ static Request deleteFilter(DeleteFilterRequest deleteFilterRequest) {
         return request;
     }
 
+    static Request setUpgradeMode(SetUpgradeModeRequest setUpgradeModeRequest) {
+        String endpoint = new EndpointBuilder().addPathPartAsIs("_ml", "set_upgrade_mode").build();
+        Request request = new Request(HttpPost.METHOD_NAME, endpoint);
+        RequestConverters.Params params = new RequestConverters.Params(request);
+        params.putParam(SetUpgradeModeRequest.ENABLED.getPreferredName(), Boolean.toString(setUpgradeModeRequest.isEnabled()));
+        if (setUpgradeModeRequest.getTimeout() != null) {
+            params.putParam(SetUpgradeModeRequest.TIMEOUT.getPreferredName(), setUpgradeModeRequest.getTimeout().toString());
+        }
+        return request;
+    }
+
     static Request mlInfo(MlInfoRequest infoRequest) {
         String endpoint = new EndpointBuilder()
             .addPathPartAsIs("_ml", "info")

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

@@ -86,6 +86,7 @@
 import org.elasticsearch.client.ml.PutJobResponse;
 import org.elasticsearch.client.ml.RevertModelSnapshotRequest;
 import org.elasticsearch.client.ml.RevertModelSnapshotResponse;
+import org.elasticsearch.client.ml.SetUpgradeModeRequest;
 import org.elasticsearch.client.ml.StartDatafeedRequest;
 import org.elasticsearch.client.ml.StartDatafeedResponse;
 import org.elasticsearch.client.ml.StopDatafeedRequest;
@@ -1838,4 +1839,42 @@ public void findFileStructureAsync(FindFileStructureRequest request, RequestOpti
             listener,
             Collections.emptySet());
     }
+
+    /**
+     * Sets the ML cluster setting upgrade_mode
+     * <p>
+     * For additional info
+     * see <a href="http://www.elastic.co/guide/en/elasticsearch/reference/current/ml-set-upgrade-mode.html">Set Upgrade Mode</a>
+     *
+     * @param request The request to set upgrade mode
+     * @param options  Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
+     * @return response
+     * @throws IOException when there is a serialization issue sending the request or receiving the response
+     */
+    public AcknowledgedResponse setUpgradeMode(SetUpgradeModeRequest request, RequestOptions options) throws IOException {
+        return restHighLevelClient.performRequestAndParseEntity(request,
+            MLRequestConverters::setUpgradeMode,
+            options,
+            AcknowledgedResponse::fromXContent,
+            Collections.emptySet());
+    }
+
+    /**
+     * Sets the ML cluster setting upgrade_mode asynchronously
+     * <p>
+     * For additional info
+     * see <a href="http://www.elastic.co/guide/en/elasticsearch/reference/current/ml-set-upgrade-mode.html">Set Upgrade Mode</a>
+     *
+     * @param request The request of Machine Learning info
+     * @param options  Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
+     * @param listener Listener to be notified upon request completion
+     */
+    public void setUpgradeModeAsync(SetUpgradeModeRequest request, RequestOptions options, ActionListener<AcknowledgedResponse> listener) {
+        restHighLevelClient.performRequestAsyncAndParseEntity(request,
+            MLRequestConverters::setUpgradeMode,
+            options,
+            AcknowledgedResponse::fromXContent,
+            listener,
+            Collections.emptySet());
+    }
 }

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

@@ -0,0 +1,93 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.client.ml;
+
+import org.elasticsearch.action.ActionRequest;
+import org.elasticsearch.action.ActionRequestValidationException;
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.unit.TimeValue;
+
+import java.util.Objects;
+
+/**
+ * Sets ML into upgrade_mode
+ */
+public class SetUpgradeModeRequest extends ActionRequest {
+
+
+    public static final ParseField ENABLED = new ParseField("enabled");
+    public static final ParseField TIMEOUT = new ParseField("timeout");
+
+    private boolean enabled;
+    private TimeValue timeout;
+
+    /**
+     * Create a new request
+     *
+     * @param enabled whether to enable `upgrade_mode` or not
+     */
+    public SetUpgradeModeRequest(boolean enabled) {
+        this.enabled = enabled;
+    }
+
+    public boolean isEnabled() {
+        return enabled;
+    }
+
+    public void setEnabled(boolean enabled) {
+        this.enabled = enabled;
+    }
+
+    public TimeValue getTimeout() {
+        return timeout;
+    }
+
+    /**
+     * How long to wait for the request to be completed
+     *
+     * @param timeout default value of 30 seconds
+     */
+    public void setTimeout(TimeValue timeout) {
+        this.timeout = timeout;
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        return null;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(enabled, timeout);
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        if (this == other) {
+            return true;
+        }
+
+        if (other == null || getClass() != other.getClass()) {
+            return false;
+        }
+
+        SetUpgradeModeRequest that = (SetUpgradeModeRequest) other;
+        return Objects.equals(enabled, that.enabled) && Objects.equals(timeout, that.timeout);
+    }
+}

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

@@ -61,6 +61,7 @@
 import org.elasticsearch.client.ml.PutFilterRequest;
 import org.elasticsearch.client.ml.PutJobRequest;
 import org.elasticsearch.client.ml.RevertModelSnapshotRequest;
+import org.elasticsearch.client.ml.SetUpgradeModeRequest;
 import org.elasticsearch.client.ml.StartDatafeedRequest;
 import org.elasticsearch.client.ml.StartDatafeedRequestTests;
 import org.elasticsearch.client.ml.StopDatafeedRequest;
@@ -818,6 +819,22 @@ public void testFindFileStructure() throws Exception {
         assertEquals(sample, requestEntityToString(request));
     }
 
+    public void testSetUpgradeMode() {
+        SetUpgradeModeRequest setUpgradeModeRequest = new SetUpgradeModeRequest(true);
+
+        Request request = MLRequestConverters.setUpgradeMode(setUpgradeModeRequest);
+        assertThat(request.getEndpoint(), equalTo("/_ml/set_upgrade_mode"));
+        assertThat(request.getMethod(), equalTo(HttpPost.METHOD_NAME));
+        assertThat(request.getParameters().get(SetUpgradeModeRequest.ENABLED.getPreferredName()), equalTo(Boolean.toString(true)));
+        assertThat(request.getParameters().containsKey(SetUpgradeModeRequest.TIMEOUT.getPreferredName()), is(false));
+
+        setUpgradeModeRequest.setTimeout(TimeValue.timeValueHours(1));
+        setUpgradeModeRequest.setEnabled(false);
+        request = MLRequestConverters.setUpgradeMode(setUpgradeModeRequest);
+        assertThat(request.getParameters().get(SetUpgradeModeRequest.ENABLED.getPreferredName()), equalTo(Boolean.toString(false)));
+        assertThat(request.getParameters().get(SetUpgradeModeRequest.TIMEOUT.getPreferredName()), is("1h"));
+    }
+
     private static Job createValidJob(String jobId) {
         AnalysisConfig.Builder analysisConfig = AnalysisConfig.builder(Collections.singletonList(
                 Detector.builder().setFunction("count").build()));

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

@@ -84,6 +84,7 @@
 import org.elasticsearch.client.ml.PutJobResponse;
 import org.elasticsearch.client.ml.RevertModelSnapshotRequest;
 import org.elasticsearch.client.ml.RevertModelSnapshotResponse;
+import org.elasticsearch.client.ml.SetUpgradeModeRequest;
 import org.elasticsearch.client.ml.StartDatafeedRequest;
 import org.elasticsearch.client.ml.StartDatafeedResponse;
 import org.elasticsearch.client.ml.StopDatafeedRequest;
@@ -1614,4 +1615,30 @@ public void testFindFileStructure() throws IOException {
         assertEquals("timestamp", structure.getTimestampField());
         assertFalse(structure.needClientTimezone());
     }
+
+    public void testEnableUpgradeMode() throws Exception {
+        MachineLearningClient machineLearningClient = highLevelClient().machineLearning();
+
+        MlInfoResponse mlInfoResponse = machineLearningClient.getMlInfo(new MlInfoRequest(), RequestOptions.DEFAULT);
+        assertThat(mlInfoResponse.getInfo().get("upgrade_mode"), equalTo(false));
+
+        AcknowledgedResponse setUpgrademodeResponse = execute(new SetUpgradeModeRequest(true),
+            machineLearningClient::setUpgradeMode,
+            machineLearningClient::setUpgradeModeAsync);
+
+        assertThat(setUpgrademodeResponse.isAcknowledged(), is(true));
+
+
+        mlInfoResponse = machineLearningClient.getMlInfo(new MlInfoRequest(), RequestOptions.DEFAULT);
+        assertThat(mlInfoResponse.getInfo().get("upgrade_mode"), equalTo(true));
+
+        setUpgrademodeResponse = execute(new SetUpgradeModeRequest(false),
+            machineLearningClient::setUpgradeMode,
+            machineLearningClient::setUpgradeModeAsync);
+
+        assertThat(setUpgrademodeResponse.isAcknowledged(), is(true));
+
+        mlInfoResponse = machineLearningClient.getMlInfo(new MlInfoRequest(), RequestOptions.DEFAULT);
+        assertThat(mlInfoResponse.getInfo().get("upgrade_mode"), equalTo(false));
+    }
 }

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

@@ -99,6 +99,7 @@
 import org.elasticsearch.client.ml.PutJobResponse;
 import org.elasticsearch.client.ml.RevertModelSnapshotRequest;
 import org.elasticsearch.client.ml.RevertModelSnapshotResponse;
+import org.elasticsearch.client.ml.SetUpgradeModeRequest;
 import org.elasticsearch.client.ml.StartDatafeedRequest;
 import org.elasticsearch.client.ml.StartDatafeedResponse;
 import org.elasticsearch.client.ml.StopDatafeedRequest;
@@ -3078,6 +3079,57 @@ public void onFailure(Exception e) {
         }
     }
 
+    public void testSetUpgradeMode() throws Exception {
+        RestHighLevelClient client = highLevelClient();
+        {
+            // tag::set-upgrade-mode-request
+            SetUpgradeModeRequest request = new SetUpgradeModeRequest(true); // <1>
+            request.setTimeout(TimeValue.timeValueMinutes(10)); // <2>
+            // end::set-upgrade-mode-request
+
+            // Set to false so that the cluster setting does not have to be unset at the end of the test.
+            request.setEnabled(false);
+
+            // tag::set-upgrade-mode-execute
+            AcknowledgedResponse acknowledgedResponse = client.machineLearning().setUpgradeMode(request, RequestOptions.DEFAULT);
+            // end::set-upgrade-mode-execute
+
+            // tag::set-upgrade-mode-response
+            boolean acknowledged = acknowledgedResponse.isAcknowledged(); // <1>
+            // end::set-upgrade-mode-response
+            assertThat(acknowledged, is(true));
+        }
+        {
+            SetUpgradeModeRequest request = new SetUpgradeModeRequest(false);
+
+            // tag::set-upgrade-mode-execute-listener
+            ActionListener<AcknowledgedResponse> listener = new ActionListener<AcknowledgedResponse>() {
+                @Override
+                public void onResponse(AcknowledgedResponse acknowledgedResponse) {
+                    // <1>
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    // <2>
+                }
+            };
+            // end::set-upgrade-mode-execute-listener
+
+            // Replace the empty listener by a blocking listener in test
+            final CountDownLatch latch = new CountDownLatch(1);
+            listener = new LatchedActionListener<>(listener, latch);
+
+            // tag::set-upgrade-mode-execute-async
+            client.machineLearning()
+                .setUpgradeModeAsync(request, RequestOptions.DEFAULT, listener); // <1>
+            // end::set-upgrade-mode-execute-async
+
+            assertTrue(latch.await(30L, TimeUnit.SECONDS));
+
+        }
+    }
+
     private String createFilter(RestHighLevelClient client) throws IOException {
         MlFilter.Builder filterBuilder = MlFilter.builder("my_safe_domains")
             .setDescription("A list of safe domains")

docs/java-rest/high-level/ml/set-upgrade-mode.asciidoc

@@ -0,0 +1,40 @@
+--
+:api: set-upgrade-mode
+:request: SetUpgradeModeRequest
+:response: AcknowledgedResponse
+--
+[id="{upid}-{api}"]
+=== Set Upgrade Mode API
+
+The Set Upgrade Mode API temporarily halts all {ml} job and {dfeed} tasks when `enabled=true`. Their
+reported states remain unchanged. Consequently, when exiting upgrade mode the halted {ml} jobs and
+{dfeeds} will return to their previous state.
+
+It accepts a +{request}+ object and responds with a +{response}+ object.
+
+When `enabled=true`, no new jobs can be opened, and no job or {dfeed} tasks will
+be running. Be sure to set `enabled=false` once upgrade actions are completed.
+
+[id="{upid}-{api}-request"]
+==== Set Upgrade Mode Request
+
+A +{request}+ object gets created setting the desired `enabled` state.
+
+["source","java",subs="attributes,callouts,macros"]
+--------------------------------------------------
+include-tagged::{doc-tests-file}[{api}-request]
+--------------------------------------------------
+<1> Constructing a new request referencing enabling upgrade mode
+<2> Optionally setting the `timeout` value for how long the
+execution should wait.
+
+[id="{upid}-{api}-response"]
+==== Set Upgrade Mode Response
+
+["source","java",subs="attributes,callouts,macros"]
+--------------------------------------------------
+include-tagged::{doc-tests-file}[{api}-response]
+--------------------------------------------------
+<1> `isAcknowledged()` from the +{response}+ indicates if the setting was set successfully.
+
+include::../execution.asciidoc[]

docs/java-rest/high-level/supported-apis.asciidoc

@@ -295,6 +295,7 @@ The Java High Level REST Client supports the following Machine Learning APIs:
 * <<{upid}-update-model-snapshot>>
 * <<{upid}-get-ml-info>>
 * <<{upid}-delete-expired-data>>
+* <<{upid}-set-upgrade-mode>>
 
 include::ml/put-job.asciidoc[]
 include::ml/get-job.asciidoc[]
@@ -338,6 +339,7 @@ include::ml/revert-model-snapshot.asciidoc[]
 include::ml/update-model-snapshot.asciidoc[]
 include::ml/get-info.asciidoc[]
 include::ml/delete-expired-data.asciidoc[]
+include::ml/set-upgrade-mode.asciidoc[]
 
 == Migration APIs