[ML] add timeout parameter for DELETE trained_models API (elastic#79739)

DELETE trained models should allow timeout parameter to override the current default of 30s closes: elastic#77070
lockewritesdocs · Oct 28, 2021 · 422d061 · 422d061
1 parent 39d0d7b
commit 422d061
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/rest-api-spec/src/main/resources/rest-api-spec/api/ml.delete_trained_model.json b/rest-api-spec/src/main/resources/rest-api-spec/api/ml.delete_trained_model.json
@@ -24,6 +24,14 @@
           }
         }
       ]
+    },
+    "params":{
+      "timeout":{
+        "type":"time",
+        "required":false,
+        "description":"Controls the amount of time to wait for the model to be deleted.",
+        "default": "30s"
+      }
     }
   }
 }
diff --git a/...src/main/java/org/elasticsearch/xpack/ml/rest/inference/RestDeleteTrainedModelAction.java b/...src/main/java/org/elasticsearch/xpack/ml/rest/inference/RestDeleteTrainedModelAction.java
@@ -6,8 +6,10 @@
  */
 package org.elasticsearch.xpack.ml.rest.inference;
 
+import org.elasticsearch.action.support.master.AcknowledgedRequest;
 import org.elasticsearch.client.node.NodeClient;
 import org.elasticsearch.core.RestApiVersion;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.rest.BaseRestHandler;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.action.RestToXContentListener;
@@ -18,6 +20,7 @@
 import java.util.List;
 
 import static org.elasticsearch.rest.RestRequest.Method.DELETE;
+import static org.elasticsearch.xpack.core.ml.action.StartTrainedModelDeploymentAction.Request.TIMEOUT;
 import static org.elasticsearch.xpack.ml.MachineLearning.BASE_PATH;
 
 public class RestDeleteTrainedModelAction extends BaseRestHandler {
@@ -39,6 +42,10 @@ public String getName() {
     protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient client) throws IOException {
         String modelId = restRequest.param(TrainedModelConfig.MODEL_ID.getPreferredName());
         DeleteTrainedModelAction.Request request = new DeleteTrainedModelAction.Request(modelId);
+        if (restRequest.hasParam(TIMEOUT.getPreferredName())) {
+            TimeValue timeout = restRequest.paramAsTime(TIMEOUT.getPreferredName(), AcknowledgedRequest.DEFAULT_ACK_TIMEOUT);
+            request.timeout(timeout);
+        }
         return channel -> client.execute(DeleteTrainedModelAction.INSTANCE, request, new RestToXContentListener<>(channel));
     }
 }