[HLRC][ML] Add ML update model snapshot API (#35537)

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

@@ -61,6 +61,7 @@
 import org.elasticsearch.client.ml.UpdateDatafeedRequest;
 import org.elasticsearch.client.ml.UpdateFilterRequest;
 import org.elasticsearch.client.ml.UpdateJobRequest;
+import org.elasticsearch.client.ml.UpdateModelSnapshotRequest;
 import org.elasticsearch.client.ml.job.util.PageParams;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.bytes.BytesReference;
@@ -391,6 +392,21 @@ static Request getModelSnapshots(GetModelSnapshotsRequest getModelSnapshotsReque
         return request;
     }
 
+    static Request updateModelSnapshot(UpdateModelSnapshotRequest updateModelSnapshotRequest) throws IOException {
+        String endpoint = new EndpointBuilder()
+            .addPathPartAsIs("_xpack")
+            .addPathPartAsIs("ml")
+            .addPathPartAsIs("anomaly_detectors")
+            .addPathPart(updateModelSnapshotRequest.getJobId())
+            .addPathPartAsIs("model_snapshots")
+            .addPathPart(updateModelSnapshotRequest.getSnapshotId())
+            .addPathPartAsIs("_update")
+            .build();
+        Request request = new Request(HttpPost.METHOD_NAME, endpoint);
+        request.setEntity(createEntity(updateModelSnapshotRequest, REQUEST_BODY_CONTENT_TYPE));
+        return request;
+    }
+
     static Request getOverallBuckets(GetOverallBucketsRequest getOverallBucketsRequest) throws IOException {
         String endpoint = new EndpointBuilder()
                 .addPathPartAsIs("_xpack")

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

@@ -79,6 +79,8 @@
 import org.elasticsearch.client.ml.UpdateDatafeedRequest;
 import org.elasticsearch.client.ml.UpdateFilterRequest;
 import org.elasticsearch.client.ml.UpdateJobRequest;
+import org.elasticsearch.client.ml.UpdateModelSnapshotRequest;
+import org.elasticsearch.client.ml.UpdateModelSnapshotResponse;
 import org.elasticsearch.client.ml.job.stats.JobStats;
 
 import java.io.IOException;
@@ -984,6 +986,47 @@ public void getModelSnapshotsAsync(GetModelSnapshotsRequest request, RequestOpti
             Collections.emptySet());
     }
 
+    /**
+     * Updates a snapshot for a Machine Learning Job.
+     * <p>
+     * For additional info
+     * see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/ml-update-snapshot.html">
+     * ML UPDATE model snapshots documentation</a>
+     *
+     * @param request The request
+     * @param options Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
+     * @throws IOException when there is a serialization issue sending the request or receiving the response
+     */
+    public UpdateModelSnapshotResponse updateModelSnapshot(UpdateModelSnapshotRequest request,
+                                                             RequestOptions options) throws IOException {
+        return restHighLevelClient.performRequestAndParseEntity(request,
+            MLRequestConverters::updateModelSnapshot,
+            options,
+            UpdateModelSnapshotResponse::fromXContent,
+            Collections.emptySet());
+    }
+
+    /**
+     * Updates a snapshot for a Machine Learning Job, notifies listener once the requested snapshots are retrieved.
+     * <p>
+     * For additional info
+     * see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/ml-update-snapshot.html">
+     * ML UPDATE model snapshots documentation</a>
+     *
+     * @param request  The request
+     * @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 updateModelSnapshotAsync(UpdateModelSnapshotRequest request, RequestOptions options,
+                                       ActionListener<UpdateModelSnapshotResponse> listener) {
+        restHighLevelClient.performRequestAsyncAndParseEntity(request,
+            MLRequestConverters::updateModelSnapshot,
+            options,
+            UpdateModelSnapshotResponse::fromXContent,
+            listener,
+            Collections.emptySet());
+    }
+
     /**
      * Gets overall buckets for a set of Machine Learning Jobs.
      * <p>

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

@@ -0,0 +1,135 @@
+/*
+ * 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.client.ml.job.config.Job;
+import org.elasticsearch.client.ml.job.process.ModelSnapshot;
+import org.elasticsearch.common.xcontent.ConstructingObjectParser;
+import org.elasticsearch.common.xcontent.ToXContentObject;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.Objects;
+
+/**
+ * A request to update information about an existing model snapshot for a given job
+ */
+public class UpdateModelSnapshotRequest extends ActionRequest implements ToXContentObject {
+
+
+    public static final ConstructingObjectParser<UpdateModelSnapshotRequest, Void> PARSER = new ConstructingObjectParser<>(
+        "update_model_snapshot_request", a -> new UpdateModelSnapshotRequest((String) a[0], (String) a[1]));
+
+
+    static {
+        PARSER.declareString(ConstructingObjectParser.constructorArg(), Job.ID);
+        PARSER.declareString(ConstructingObjectParser.constructorArg(), ModelSnapshot.SNAPSHOT_ID);
+        PARSER.declareStringOrNull(UpdateModelSnapshotRequest::setDescription, ModelSnapshot.DESCRIPTION);
+        PARSER.declareBoolean(UpdateModelSnapshotRequest::setRetain, ModelSnapshot.RETAIN);
+    }
+
+    private final String jobId;
+    private String snapshotId;
+    private String description;
+    private Boolean retain;
+
+    /**
+     * Constructs a request to update information for a snapshot of given job
+     * @param jobId id of the job from which to retrieve results
+     * @param snapshotId id of the snapshot from which to retrieve results
+     */
+    public UpdateModelSnapshotRequest(String jobId, String snapshotId) {
+        this.jobId = Objects.requireNonNull(jobId, "[" + Job.ID + "] must not be null");
+        this.snapshotId = Objects.requireNonNull(snapshotId, "[" + ModelSnapshot.SNAPSHOT_ID + "] must not be null");
+    }
+
+    public String getJobId() {
+        return jobId;
+    }
+
+    public String getSnapshotId() {
+        return snapshotId;
+    }
+
+    public String getDescription() {
+        return description;
+    }
+
+    /**
+     * The new description of the snapshot.
+     * @param description the updated snapshot description
+     */
+    public void setDescription(String description) {
+        this.description = description;
+    }
+
+    public Boolean getRetain() {
+        return retain;
+    }
+
+    /**
+     * The new value of the "retain" property of the snapshot
+     * @param retain the updated retain property
+     */
+    public void setRetain(boolean retain) {
+        this.retain = retain;
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        return null;
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        builder.field(Job.ID.getPreferredName(), jobId);
+        builder.field(ModelSnapshot.SNAPSHOT_ID.getPreferredName(), snapshotId);
+        if (description != null) {
+            builder.field(ModelSnapshot.DESCRIPTION.getPreferredName(), description);
+        }
+        if (retain != null) {
+            builder.field(ModelSnapshot.RETAIN.getPreferredName(), retain);
+        }
+        builder.endObject();
+        return builder;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == null) {
+            return false;
+        }
+        if (getClass() != obj.getClass()) {
+            return false;
+        }
+        UpdateModelSnapshotRequest request = (UpdateModelSnapshotRequest) obj;
+        return Objects.equals(jobId, request.jobId)
+            && Objects.equals(snapshotId, request.snapshotId)
+            && Objects.equals(description, request.description)
+            && Objects.equals(retain, request.retain);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(jobId, snapshotId, description, retain);
+    }
+}

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

@@ -0,0 +1,109 @@
+/*
+ * 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.ActionResponse;
+import org.elasticsearch.client.ml.job.process.ModelSnapshot;
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.xcontent.ConstructingObjectParser;
+import org.elasticsearch.common.xcontent.ToXContent;
+import org.elasticsearch.common.xcontent.ToXContentObject;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.util.Objects;
+
+/**
+ * A response acknowledging the update of information for an existing model snapshot for a given job
+ */
+public class UpdateModelSnapshotResponse extends ActionResponse implements ToXContentObject {
+
+    private static final ParseField ACKNOWLEDGED = new ParseField("acknowledged");
+    private static final ParseField MODEL = new ParseField("model");
+
+    public UpdateModelSnapshotResponse(boolean acknowledged, ModelSnapshot.Builder modelSnapshot) {
+        this.acknowledged = acknowledged;
+        this.model = modelSnapshot.build();
+    }
+
+    public static final ConstructingObjectParser<UpdateModelSnapshotResponse, Void> PARSER =
+        new ConstructingObjectParser<>("update_model_snapshot_response", true,
+            a -> new UpdateModelSnapshotResponse((Boolean) a[0], ((ModelSnapshot.Builder) a[1])));
+
+    static {
+        PARSER.declareBoolean(ConstructingObjectParser.constructorArg(), ACKNOWLEDGED);
+        PARSER.declareObject(ConstructingObjectParser.constructorArg(), ModelSnapshot.PARSER, MODEL);
+    }
+
+    public static UpdateModelSnapshotResponse fromXContent(XContentParser parser) throws IOException {
+        return PARSER.parse(parser, null);
+    }
+
+    private final Boolean acknowledged;
+    private final ModelSnapshot model;
+
+    /**
+     * Get the action acknowledgement
+     * @return  a {@code boolean} that indicates whether the model snapshot was updated successfully.
+     */
+    public Boolean getAcknowledged() {
+        return acknowledged;
+    }
+
+    /**
+     * Get the updated snapshot of the model
+     * @return the updated model snapshot.
+     */
+    public  ModelSnapshot getModel() {
+        return model;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(acknowledged, model);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
+        builder.startObject();
+        if (acknowledged != null) {
+            builder.field(ACKNOWLEDGED.getPreferredName(), acknowledged);
+        }
+        if (model != null) {
+            builder.field(MODEL.getPreferredName(), model);
+        }
+        builder.endObject();
+        return builder;
+    }
+
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == null) {
+            return false;
+        }
+        if (getClass() != obj.getClass()) {
+            return false;
+        }
+        UpdateModelSnapshotResponse request = (UpdateModelSnapshotResponse) obj;
+        return Objects.equals(acknowledged, request.acknowledged)
+            && Objects.equals(model, request.model);
+    }
+}

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/job/process/ModelSnapshot.java

@@ -176,6 +176,10 @@ public Quantiles getQuantiles() {
         return quantiles;
     }
 
+    public boolean getRetain() {
+        return retain;
+    }
+
     public Date getLatestRecordTimeStamp() {
         return latestRecordTimeStamp;
     }

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

@@ -57,6 +57,7 @@
 import org.elasticsearch.client.ml.StopDatafeedRequest;
 import org.elasticsearch.client.ml.UpdateFilterRequest;
 import org.elasticsearch.client.ml.UpdateJobRequest;
+import org.elasticsearch.client.ml.UpdateModelSnapshotRequest;
 import org.elasticsearch.client.ml.calendars.Calendar;
 import org.elasticsearch.client.ml.calendars.CalendarTests;
 import org.elasticsearch.client.ml.datafeed.DatafeedConfig;
@@ -422,6 +423,22 @@ public void testGetModelSnapshots() throws IOException {
         }
     }
 
+    public void testUpdateModelSnapshot() throws IOException {
+        String jobId = randomAlphaOfLength(10);
+        String snapshotId = randomAlphaOfLength(10);
+        UpdateModelSnapshotRequest updateModelSnapshotRequest = new UpdateModelSnapshotRequest(jobId, snapshotId);
+        updateModelSnapshotRequest.setDescription("My First Snapshot");
+        updateModelSnapshotRequest.setRetain(true);
+
+        Request request = MLRequestConverters.updateModelSnapshot(updateModelSnapshotRequest);
+        assertEquals(HttpPost.METHOD_NAME, request.getMethod());
+        assertEquals("/_xpack/ml/anomaly_detectors/" + jobId + "/model_snapshots/" + snapshotId + "/_update", request.getEndpoint());
+        try (XContentParser parser = createParser(JsonXContent.jsonXContent, request.getEntity().getContent())) {
+            UpdateModelSnapshotRequest parsedRequest = UpdateModelSnapshotRequest.PARSER.apply(parser, null);
+            assertThat(parsedRequest, equalTo(updateModelSnapshotRequest));
+        }
+    }
+
     public void testGetOverallBuckets() throws IOException {
         String jobId = randomAlphaOfLength(10);
         GetOverallBucketsRequest getOverallBucketsRequest = new GetOverallBucketsRequest(jobId);