[ML] Allow marking a model snapshot with retain (elastic/x-pack-elasticsearch#880)

This change adds a retain field to model snapshots.
A user can set retain to true/false via the update model snapshot API.
Model snapshots with retain set to true will not be deleted by
the daily maintenance service regardless of whether they expired.

This allows users to keep always keep certain snapshots around for
potentially reverting to in the future.

relates elastic/x-pack-elasticsearch#758

Original commit: elastic/x-pack-elasticsearch@2283989a33

plugin/src/main/java/org/elasticsearch/xpack/ml/action/UpdateModelSnapshotAction.java

@@ -42,9 +42,8 @@ import java.util.List;
 import java.util.Objects;
 import java.util.function.Consumer;
 
-public class UpdateModelSnapshotAction extends
+public class UpdateModelSnapshotAction extends Action<UpdateModelSnapshotAction.Request,
-Action<UpdateModelSnapshotAction.Request, UpdateModelSnapshotAction.Response,
+        UpdateModelSnapshotAction.Response, UpdateModelSnapshotAction.RequestBuilder> {
-UpdateModelSnapshotAction.RequestBuilder> {
 
     public static final UpdateModelSnapshotAction INSTANCE = new UpdateModelSnapshotAction();
     public static final String NAME = "cluster:admin/ml/job/model_snapshots/update";
@@ -70,7 +69,8 @@ UpdateModelSnapshotAction.RequestBuilder> {
         static {
             PARSER.declareString((request, jobId) -> request.jobId = jobId, Job.ID);
             PARSER.declareString((request, snapshotId) -> request.snapshotId = snapshotId, ModelSnapshot.SNAPSHOT_ID);
-            PARSER.declareString((request, description) -> request.description = description, ModelSnapshot.DESCRIPTION);
+            PARSER.declareString(Request::setDescription, ModelSnapshot.DESCRIPTION);
+            PARSER.declareBoolean(Request::setRetain, ModelSnapshot.RETAIN);
         }
 
         public static Request parseRequest(String jobId, String snapshotId, XContentParser parser) {
@@ -87,14 +87,14 @@ UpdateModelSnapshotAction.RequestBuilder> {
         private String jobId;
         private String snapshotId;
         private String description;
+        private Boolean retain;
 
         Request() {
         }
 
-        public Request(String jobId, String snapshotId, String description) {
+        public Request(String jobId, String snapshotId) {
             this.jobId = ExceptionsHelper.requireNonNull(jobId, Job.ID.getPreferredName());
             this.snapshotId = ExceptionsHelper.requireNonNull(snapshotId, ModelSnapshot.SNAPSHOT_ID.getPreferredName());
-            this.description = ExceptionsHelper.requireNonNull(description, ModelSnapshot.DESCRIPTION.getPreferredName());
         }
 
         public String getJobId() {
@@ -105,10 +105,22 @@ UpdateModelSnapshotAction.RequestBuilder> {
             return snapshotId;
         }
 
-        public String getDescriptionString() {
+        public String getDescription() {
             return description;
         }
 
+        public void setDescription(String description) {
+            this.description = description;
+        }
+
+        public Boolean getRetain() {
+            return retain;
+        }
+
+        public void setRetain(Boolean retain) {
+            this.retain = retain;
+        }
+
         @Override
         public ActionRequestValidationException validate() {
             return null;
@@ -119,7 +131,8 @@ UpdateModelSnapshotAction.RequestBuilder> {
             super.readFrom(in);
             jobId = in.readString();
             snapshotId = in.readString();
-            description = in.readString();
+            description = in.readOptionalString();
+            retain = in.readOptionalBoolean();
         }
 
         @Override
@@ -127,7 +140,8 @@ UpdateModelSnapshotAction.RequestBuilder> {
             super.writeTo(out);
             out.writeString(jobId);
             out.writeString(snapshotId);
-            out.writeString(description);
+            out.writeOptionalString(description);
+            out.writeOptionalBoolean(retain);
         }
 
         @Override
@@ -135,14 +149,19 @@ UpdateModelSnapshotAction.RequestBuilder> {
             builder.startObject();
             builder.field(Job.ID.getPreferredName(), jobId);
             builder.field(ModelSnapshot.SNAPSHOT_ID.getPreferredName(), snapshotId);
-            builder.field(ModelSnapshot.DESCRIPTION.getPreferredName(), description);
+            if (description != null) {
+                builder.field(ModelSnapshot.DESCRIPTION.getPreferredName(), description);
+            }
+            if (retain != null) {
+                builder.field(ModelSnapshot.RETAIN.getPreferredName(), retain);
+            }
             builder.endObject();
             return builder;
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(jobId, snapshotId, description);
+            return Objects.hash(jobId, snapshotId, description, retain);
         }
 
         @Override
@@ -154,8 +173,10 @@ UpdateModelSnapshotAction.RequestBuilder> {
                 return false;
             }
             Request other = (Request) obj;
-            return Objects.equals(jobId, other.jobId) && Objects.equals(snapshotId, other.snapshotId)
+            return Objects.equals(jobId, other.jobId)
-                    && Objects.equals(description, other.description);
+                    && Objects.equals(snapshotId, other.snapshotId)
+                    && Objects.equals(description, other.description)
+                    && Objects.equals(retain, other.retain);
         }
     }
 
@@ -250,17 +271,14 @@ UpdateModelSnapshotAction.RequestBuilder> {
 
         @Override
         protected void doExecute(Request request, ActionListener<Response> listener) {
-            logger.debug("Received request to change model snapshot description using '" + request.getDescriptionString()
+            logger.debug("Received request to update model snapshot [{}] for job [{}]", request.getSnapshotId(), request.getJobId());
-            + "' for snapshot ID '" + request.getSnapshotId() + "' for job '" + request.getJobId() + "'");
             getChangeCandidates(request, changeCandidates -> {
                 checkForClashes(request, aVoid -> {
                     if (changeCandidates.size() > 1) {
                         logger.warn("More than one model found for [{}: {}, {}: {}] tuple.", Job.ID.getPreferredName(), request.getJobId(),
                                 ModelSnapshot.SNAPSHOT_ID.getPreferredName(), request.getSnapshotId());
                     }
-                    ModelSnapshot.Builder updatedSnapshotBuilder = new ModelSnapshot.Builder(changeCandidates.get(0));
+                    ModelSnapshot updatedSnapshot = applyUpdate(request, changeCandidates.get(0));
-                    updatedSnapshotBuilder.setDescription(request.getDescriptionString());
-                    ModelSnapshot updatedSnapshot = updatedSnapshotBuilder.build();
                     jobManager.updateModelSnapshot(updatedSnapshot, b -> {
                         // The quantiles can be large, and totally dominate the output -
                         // it's clearer to remove them
@@ -283,10 +301,15 @@ UpdateModelSnapshotAction.RequestBuilder> {
         }
 
         private void checkForClashes(Request request, Consumer<Void> handler, Consumer<Exception> errorHandler) {
-            getModelSnapshots(request.getJobId(), null, request.getDescriptionString(), clashCandidates -> {
+            if (request.getDescription() == null) {
+                handler.accept(null);
+                return;
+            }
+
+            getModelSnapshots(request.getJobId(), null, request.getDescription(), clashCandidates -> {
                 if (clashCandidates != null && !clashCandidates.isEmpty()) {
                     errorHandler.accept(new IllegalArgumentException(Messages.getMessage(
-                            Messages.REST_DESCRIPTION_ALREADY_USED, request.getDescriptionString(), request.getJobId())));
+                            Messages.REST_DESCRIPTION_ALREADY_USED, request.getDescription(), request.getJobId())));
                 } else {
                     handler.accept(null);
                 }
@@ -299,6 +322,17 @@ UpdateModelSnapshotAction.RequestBuilder> {
                     page -> handler.accept(page.results()), errorHandler);
         }
 
+        private static ModelSnapshot applyUpdate(Request request, ModelSnapshot target) {
+            ModelSnapshot.Builder updatedSnapshotBuilder = new ModelSnapshot.Builder(target);
+            if (request.getDescription() != null) {
+                updatedSnapshotBuilder.setDescription(request.getDescription());
+            }
+            if (request.getRetain() != null) {
+                updatedSnapshotBuilder.setRetain(request.getRetain());
+            }
+            return updatedSnapshotBuilder.build();
+        }
+
     }
 
 }

plugin/src/main/java/org/elasticsearch/xpack/ml/job/persistence/ElasticsearchMappings.java

@@ -603,6 +603,9 @@ public class ElasticsearchMappings {
                             .startObject(ModelSnapshot.SNAPSHOT_DOC_COUNT.getPreferredName())
                                 .field(TYPE, INTEGER)
                             .endObject()
+                            .startObject(ModelSnapshot.RETAIN.getPreferredName())
+                                .field(TYPE, BOOLEAN)
+                            .endObject()
                             .startObject(ModelSizeStats.RESULT_TYPE_FIELD.getPreferredName())
                                 .startObject(PROPERTIES)
                                     .startObject(Job.ID.getPreferredName())

plugin/src/main/java/org/elasticsearch/xpack/ml/job/process/autodetect/state/ModelSnapshot.java

@@ -40,6 +40,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
     public static final ParseField SNAPSHOT_DOC_COUNT = new ParseField("snapshot_doc_count");
     public static final ParseField LATEST_RECORD_TIME = new ParseField("latest_record_time_stamp");
     public static final ParseField LATEST_RESULT_TIME = new ParseField("latest_result_time_stamp");
+    public static final ParseField RETAIN = new ParseField("retain");
 
     // Used for QueryPage
     public static final ParseField RESULTS_FIELD = new ParseField("model_snapshots");
@@ -84,6 +85,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
                     "unexpected token [" + p.currentToken() + "] for [" + LATEST_RESULT_TIME.getPreferredName() + "]");
         }, LATEST_RESULT_TIME, ValueType.VALUE);
         PARSER.declareObject(Builder::setQuantiles, Quantiles.PARSER, Quantiles.TYPE);
+        PARSER.declareBoolean(Builder::setRetain, RETAIN);
     }
 
     private final String jobId;
@@ -95,9 +97,11 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
     private final Date latestRecordTimeStamp;
     private final Date latestResultTimeStamp;
     private final Quantiles quantiles;
+    private final boolean retain;
 
     private ModelSnapshot(String jobId, Date timestamp, String description, String snapshotId, int snapshotDocCount,
-                          ModelSizeStats modelSizeStats, Date latestRecordTimeStamp, Date latestResultTimeStamp, Quantiles quantiles) {
+                          ModelSizeStats modelSizeStats, Date latestRecordTimeStamp, Date latestResultTimeStamp, Quantiles quantiles,
+                          boolean retain) {
         this.jobId = jobId;
         this.timestamp = timestamp;
         this.description = description;
@@ -107,6 +111,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
         this.latestRecordTimeStamp = latestRecordTimeStamp;
         this.latestResultTimeStamp = latestResultTimeStamp;
         this.quantiles = quantiles;
+        this.retain = retain;
     }
 
     public ModelSnapshot(StreamInput in) throws IOException {
@@ -119,6 +124,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
         latestRecordTimeStamp = in.readBoolean() ? new Date(in.readVLong()) : null;
         latestResultTimeStamp = in.readBoolean() ? new Date(in.readVLong()) : null;
         quantiles = in.readOptionalWriteable(Quantiles::new);
+        retain = in.readBoolean();
     }
 
     @Override
@@ -147,6 +153,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
             out.writeBoolean(false);
         }
         out.writeOptionalWriteable(quantiles);
+        out.writeBoolean(retain);
     }
 
     @Override
@@ -177,6 +184,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
         if (quantiles != null) {
             builder.field(Quantiles.TYPE.getPreferredName(), quantiles);
         }
+        builder.field(RETAIN.getPreferredName(), retain);
         builder.endObject();
         return builder;
     }
@@ -220,7 +228,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
     @Override
     public int hashCode() {
         return Objects.hash(jobId, timestamp, description, snapshotId, quantiles, snapshotDocCount, modelSizeStats, latestRecordTimeStamp,
-                latestResultTimeStamp);
+                latestResultTimeStamp, retain);
     }
 
     /**
@@ -246,7 +254,8 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
                 && Objects.equals(this.modelSizeStats, that.modelSizeStats)
                 && Objects.equals(this.quantiles, that.quantiles)
                 && Objects.equals(this.latestRecordTimeStamp, that.latestRecordTimeStamp)
-                && Objects.equals(this.latestResultTimeStamp, that.latestResultTimeStamp);
+                && Objects.equals(this.latestResultTimeStamp, that.latestResultTimeStamp)
+                && this.retain == that.retain;
     }
 
     public static String documentId(ModelSnapshot snapshot) {
@@ -275,6 +284,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
         private Date latestRecordTimeStamp;
         private Date latestResultTimeStamp;
         private Quantiles quantiles;
+        private boolean retain;
 
         public Builder() {
         }
@@ -294,6 +304,7 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
             this.latestRecordTimeStamp = modelSnapshot.latestRecordTimeStamp;
             this.latestResultTimeStamp = modelSnapshot.latestResultTimeStamp;
             this.quantiles = modelSnapshot.quantiles;
+            this.retain = modelSnapshot.retain;
         }
 
         public Builder setJobId(String jobId) {
@@ -346,9 +357,14 @@ public class ModelSnapshot extends ToXContentToBytes implements Writeable {
             return this;
         }
 
+        public Builder setRetain(boolean value) {
+            this.retain = value;
+            return this;
+        }
+
         public ModelSnapshot build() {
             return new ModelSnapshot(jobId, timestamp, description, snapshotId, snapshotDocCount, modelSizeStats, latestRecordTimeStamp,
-                    latestResultTimeStamp, quantiles);
+                    latestResultTimeStamp, quantiles, retain);
         }
     }
 }

plugin/src/main/java/org/elasticsearch/xpack/ml/job/results/ReservedFieldNames.java

@@ -143,6 +143,7 @@ public final class ReservedFieldNames {
             ModelSnapshot.SNAPSHOT_DOC_COUNT.getPreferredName(),
             ModelSnapshot.LATEST_RECORD_TIME.getPreferredName(),
             ModelSnapshot.LATEST_RESULT_TIME.getPreferredName(),
+            ModelSnapshot.RETAIN.getPreferredName(),
 
             PerPartitionMaxProbabilities.PER_PARTITION_MAX_PROBABILITIES.getPreferredName(),
             PerPartitionMaxProbabilities.MAX_RECORD_SCORE.getPreferredName(),

plugin/src/main/java/org/elasticsearch/xpack/ml/job/retention/ExpiredModelSnapshotsRemover.java

@@ -64,12 +64,20 @@ public class ExpiredModelSnapshotsRemover extends AbstractExpiredJobDataRemover
             return;
         }
         LOGGER.info("Removing model snapshots of job [{}] that have a timestamp before [{}]", job.getId(), cutoffEpochMs);
-        QueryBuilder excludeFilter = QueryBuilders.termQuery(ModelSnapshot.SNAPSHOT_ID.getPreferredName(), job.getModelSnapshotId());
+
         SearchRequest searchRequest = new SearchRequest();
         searchRequest.indices(AnomalyDetectorsIndex.jobResultsAliasedName(job.getId()));
         searchRequest.types(ModelSnapshot.TYPE.getPreferredName());
-        QueryBuilder query = createQuery(job.getId(), cutoffEpochMs).mustNot(excludeFilter);
+
+        QueryBuilder activeSnapshotFilter = QueryBuilders.termQuery(
+                ModelSnapshot.SNAPSHOT_ID.getPreferredName(), job.getModelSnapshotId());
+        QueryBuilder retainFilter = QueryBuilders.termQuery(ModelSnapshot.RETAIN.getPreferredName(), true);
+        QueryBuilder query = createQuery(job.getId(), cutoffEpochMs)
+                .mustNot(activeSnapshotFilter)
+                .mustNot(retainFilter);
+
         searchRequest.source(new SearchSourceBuilder().query(query).size(MODEL_SNAPSHOT_SEARCH_SIZE));
+
         client.execute(SearchAction.INSTANCE, searchRequest, new ActionListener<SearchResponse>() {
             @Override
             public void onResponse(SearchResponse searchResponse) {

plugin/src/test/java/org/elasticsearch/xpack/ml/action/UpdateModelSnapshotActionRequestTests.java

@@ -9,7 +9,7 @@ import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.xpack.ml.action.UpdateModelSnapshotAction.Request;
 import org.elasticsearch.xpack.ml.support.AbstractStreamableXContentTestCase;
 
-public class PutModelSnapshotDescriptionActionRequestTests
+public class UpdateModelSnapshotActionRequestTests
         extends AbstractStreamableXContentTestCase<UpdateModelSnapshotAction.Request> {
 
     @Override
@@ -19,12 +19,19 @@ public class PutModelSnapshotDescriptionActionRequestTests
 
     @Override
     protected Request createTestInstance() {
-        return new Request(randomAsciiOfLengthBetween(1, 20), randomAsciiOfLengthBetween(1, 20), randomAsciiOfLengthBetween(1, 20));
+        Request request = new Request(randomAsciiOfLengthBetween(1, 20),
+                randomAsciiOfLengthBetween(1, 20));
+        if (randomBoolean()) {
+            request.setDescription(randomAsciiOfLengthBetween(1, 20));
+        }
+        if (randomBoolean()) {
+            request.setRetain(randomBoolean());
+        }
+        return request;
     }
 
     @Override
     protected Request createBlankInstance() {
         return new Request();
     }
-
 }

plugin/src/test/java/org/elasticsearch/xpack/ml/action/UpdateModelSnapshotActionResponseTests.java

@@ -9,7 +9,8 @@ import org.elasticsearch.xpack.ml.action.UpdateModelSnapshotAction.Response;
 import org.elasticsearch.xpack.ml.job.process.autodetect.state.ModelSnapshotTests;
 import org.elasticsearch.xpack.ml.support.AbstractStreamableTestCase;
 
-public class PutModelSnapshotDescriptionActionResponseTests extends AbstractStreamableTestCase<UpdateModelSnapshotAction.Response> {
+public class UpdateModelSnapshotActionResponseTests
+        extends AbstractStreamableTestCase<UpdateModelSnapshotAction.Response> {
 
     @Override
     protected Response createTestInstance() {

plugin/src/test/java/org/elasticsearch/xpack/ml/integration/DeleteExpiredDataIT.java

@@ -29,6 +29,7 @@ import org.elasticsearch.xpack.ml.action.OpenJobAction;
 import org.elasticsearch.xpack.ml.action.PutDatafeedAction;
 import org.elasticsearch.xpack.ml.action.PutJobAction;
 import org.elasticsearch.xpack.ml.action.StartDatafeedAction;
+import org.elasticsearch.xpack.ml.action.UpdateModelSnapshotAction;
 import org.elasticsearch.xpack.ml.datafeed.DatafeedConfig;
 import org.elasticsearch.xpack.ml.job.config.AnalysisConfig;
 import org.elasticsearch.xpack.ml.job.config.DataDescription;
@@ -47,9 +48,7 @@ import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 import static org.hamcrest.Matchers.equalTo;
@@ -122,6 +121,7 @@ public class DeleteExpiredDataIT extends SecurityIntegTestCase {
         jobs.add(newJobBuilder("no-retention").setResultsRetentionDays(null).setModelSnapshotRetentionDays(null).build());
         jobs.add(newJobBuilder("results-retention").setResultsRetentionDays(1L).setModelSnapshotRetentionDays(null).build());
         jobs.add(newJobBuilder("snapshots-retention").setResultsRetentionDays(null).setModelSnapshotRetentionDays(2L).build());
+        jobs.add(newJobBuilder("snapshots-retention-with-retain").setResultsRetentionDays(null).setModelSnapshotRetentionDays(2L).build());
         jobs.add(newJobBuilder("results-and-snapshots-retention").setResultsRetentionDays(1L).setModelSnapshotRetentionDays(2L).build());
 
         long now = System.currentTimeMillis();
@@ -171,6 +171,8 @@ public class DeleteExpiredDataIT extends SecurityIntegTestCase {
             assertThat(modelSnapshots.size(), equalTo(2));
         }
 
+        retainAllSnapshots("snapshots-retention-with-retain");
+
         long totalModelSizeStatsBeforeDelete = client().prepareSearch("*").setTypes("result")
                 .setQuery(QueryBuilders.termQuery("result_type", "model_size_stats"))
                 .get().getHits().totalHits;
@@ -199,6 +201,10 @@ public class DeleteExpiredDataIT extends SecurityIntegTestCase {
         assertThat(getRecords("snapshots-retention").size(), equalTo(1));
         assertThat(getModelSnapshots("snapshots-retention").size(), equalTo(1));
 
+        assertThat(getBuckets("snapshots-retention-with-retain").size(), is(greaterThanOrEqualTo(70)));
+        assertThat(getRecords("snapshots-retention-with-retain").size(), equalTo(1));
+        assertThat(getModelSnapshots("snapshots-retention-with-retain").size(), equalTo(2));
+
         buckets = getBuckets("results-and-snapshots-retention");
         assertThat(buckets.size(), is(lessThanOrEqualTo(24)));
         assertThat(buckets.size(), is(greaterThanOrEqualTo(22)));
@@ -272,4 +278,16 @@ public class DeleteExpiredDataIT extends SecurityIntegTestCase {
     protected void ensureClusterStateConsistency() throws IOException {
         // this method in ESIntegTestCase is not plugin-friendly - it does not account for plugin NamedWritableRegistries
     }
+
+    private void retainAllSnapshots(String jobId) throws Exception {
+        List<ModelSnapshot> modelSnapshots = getModelSnapshots(jobId);
+        for (ModelSnapshot modelSnapshot : modelSnapshots) {
+            UpdateModelSnapshotAction.Request request = new UpdateModelSnapshotAction.Request(
+                    jobId, modelSnapshot.getSnapshotId());
+            request.setRetain(true);
+            client().execute(UpdateModelSnapshotAction.INSTANCE, request).get();
+        }
+        // We need to refresh to ensure the updates are visible
+        client().admin().indices().prepareRefresh("*").get();
+    }
 }

plugin/src/test/java/org/elasticsearch/xpack/ml/job/process/autodetect/state/ModelSnapshotTests.java

@@ -19,6 +19,7 @@ public class ModelSnapshotTests extends AbstractSerializingTestCase<ModelSnapsho
     private static final int DEFAULT_DOC_COUNT = 7;
     private static final Date DEFAULT_LATEST_RESULT_TIMESTAMP = new Date(12345678901234L);
     private static final Date DEFAULT_LATEST_RECORD_TIMESTAMP = new Date(12345678904321L);
+    private static final boolean DEFAULT_RETAIN = true;
 
     public void testCopyBuilder() {
         ModelSnapshot modelSnapshot1 = createFullyPopulated().build();
@@ -132,6 +133,7 @@ public class ModelSnapshotTests extends AbstractSerializingTestCase<ModelSnapsho
         modelSnapshot.setLatestResultTimeStamp(DEFAULT_LATEST_RESULT_TIMESTAMP);
         modelSnapshot.setLatestRecordTimeStamp(DEFAULT_LATEST_RECORD_TIMESTAMP);
         modelSnapshot.setQuantiles(new Quantiles("foo", DEFAULT_TIMESTAMP, "state"));
+        modelSnapshot.setRetain(DEFAULT_RETAIN);
         return modelSnapshot;
     }
 
@@ -152,6 +154,7 @@ public class ModelSnapshotTests extends AbstractSerializingTestCase<ModelSnapsho
         modelSnapshot.setLatestRecordTimeStamp(
                 new Date(TimeValue.parseTimeValue(randomTimeValue(), "test").millis()));
         modelSnapshot.setQuantiles(QuantilesTests.createRandomized());
+        modelSnapshot.setRetain(randomBoolean());
         return modelSnapshot.build();
     }
 

plugin/src/test/java/org/elasticsearch/xpack/ml/modelsnapshots/UpdateModelSnapshotActionTests.java

@@ -9,21 +9,15 @@ import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.ml.action.UpdateModelSnapshotAction;
 
 
-public class PutModelSnapshotDescriptionTests extends ESTestCase {
+public class UpdateModelSnapshotActionTests extends ESTestCase {
 
     public void testUpdateDescription_GivenMissingArg() {
         IllegalArgumentException e = expectThrows(IllegalArgumentException.class,
-                () -> new UpdateModelSnapshotAction.Request(null, "foo", "bar"));
+                () -> new UpdateModelSnapshotAction.Request(null, "foo"));
         assertEquals("[job_id] must not be null.", e.getMessage());
 
         e = expectThrows(IllegalArgumentException.class,
-                () -> new UpdateModelSnapshotAction.Request("foo", null, "bar"));
+                () -> new UpdateModelSnapshotAction.Request("foo", null));
         assertEquals("[snapshot_id] must not be null.", e.getMessage());
-
-        e = expectThrows(IllegalArgumentException.class,
-                () -> new UpdateModelSnapshotAction.Request("foo", "foo", null));
-        assertEquals("[description] must not be null.", e.getMessage());
     }
-
-
 }

plugin/src/test/resources/rest-api-spec/api/xpack.ml.update_model_snapshot.json

@@ -13,13 +13,13 @@
         "snapshot_id": {
           "type": "string",
           "required": true,
-          "description": "The ID of the snapshot whose description is to be updated"
+          "description": "The ID of the snapshot to update"
         }
       },
       "params": {}
     },
     "body": {
-      "description" : "A JSON object containing the description to update with",
+      "description" : "The model snapshot properties to update",
       "required" : true
     }
   }

plugin/src/test/resources/rest-api-spec/test/ml/update_model_snapshot.yaml

@@ -24,7 +24,8 @@ setup:
           {
             "job_id" : "foo",
             "timestamp": "2016-06-02T00:00:00Z",
-            "snapshot_id": "foo"
+            "snapshot_id": "foo",
+            "retain": false
           }
 
   - do:
@@ -40,25 +41,14 @@ setup:
             "job_id": "foo",
             "timestamp": "2016-06-01T00:00:00Z",
             "snapshot_id": "bar",
-            "description": "bar"
+            "description": "bar",
+            "retain": true
           }
 
   - do:
       indices.refresh:
         index: .ml-anomalies-foo
 
----
-"Test without description":
-  - do:
-      catch: request
-      xpack.ml.update_model_snapshot:
-        job_id: "foo"
-        snapshot_id: "foo"
-        body: >
-          {
-            "some_field": "foo"
-          }
-
 ---
 "Test with valid description":
   - do:
@@ -79,6 +69,7 @@ setup:
           }
 
   - match: { acknowledged: true }
+  - match: { model.retain: false }
   - match: { model.description: "new_description" }
 
   - do:
@@ -114,3 +105,46 @@ setup:
           {
             "description": "bar"
           }
+
+---
+"Test with retain":
+  - do:
+      xpack.ml.update_model_snapshot:
+        job_id: "foo"
+        snapshot_id: "foo"
+        body: >
+          {
+            "retain": true
+          }
+
+  - match: { acknowledged: true }
+  - match: { model.retain: true }
+
+  - do:
+      xpack.ml.update_model_snapshot:
+        job_id: "foo"
+        snapshot_id: "bar"
+        body: >
+          {
+            "retain": false
+          }
+
+  - match: { acknowledged: true }
+  - match: { model.retain: false }
+  - match: { model.description: "bar" }
+
+---
+"Test with all fields":
+  - do:
+      xpack.ml.update_model_snapshot:
+        job_id: "foo"
+        snapshot_id: "foo"
+        body: >
+          {
+            "description": "new foo",
+            "retain": true
+          }
+
+  - match: { acknowledged: true }
+  - match: { model.description: "new foo" }
+  - match: { model.retain: true }