HLRC: Add ML Get Buckets API (#33056)

Relates #29827

Author: dimitris-athanasiou

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

@@ -28,6 +28,7 @@
 import org.elasticsearch.protocol.xpack.ml.CloseJobRequest;
 import org.elasticsearch.protocol.xpack.ml.DeleteJobRequest;
 import org.elasticsearch.protocol.xpack.ml.GetJobRequest;
+import org.elasticsearch.protocol.xpack.ml.GetBucketsRequest;
 import org.elasticsearch.protocol.xpack.ml.OpenJobRequest;
 import org.elasticsearch.protocol.xpack.ml.PutJobRequest;
 
@@ -69,7 +70,7 @@ static Request getJob(GetJobRequest getJobRequest) {
         return request;
     }
 
-    static Request openJob(OpenJobRequest openJobRequest) throws IOException {
+    static Request openJob(OpenJobRequest openJobRequest) {
         String endpoint = new EndpointBuilder()
                 .addPathPartAsIs("_xpack")
                 .addPathPartAsIs("ml")
@@ -109,4 +110,18 @@ static Request deleteJob(DeleteJobRequest deleteJobRequest) {
 
         return request;
     }
+
+    static Request getBuckets(GetBucketsRequest getBucketsRequest) throws IOException {
+        String endpoint = new EndpointBuilder()
+                .addPathPartAsIs("_xpack")
+                .addPathPartAsIs("ml")
+                .addPathPartAsIs("anomaly_detectors")
+                .addPathPart(getBucketsRequest.getJobId())
+                .addPathPartAsIs("results")
+                .addPathPartAsIs("buckets")
+                .build();
+        Request request = new Request(HttpGet.METHOD_NAME, endpoint);
+        request.setEntity(createEntity(getBucketsRequest, REQUEST_BODY_CONTENT_TYPE));
+        return request;
+    }
 }

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

@@ -23,6 +23,8 @@
 import org.elasticsearch.protocol.xpack.ml.CloseJobResponse;
 import org.elasticsearch.protocol.xpack.ml.DeleteJobRequest;
 import org.elasticsearch.protocol.xpack.ml.DeleteJobResponse;
+import org.elasticsearch.protocol.xpack.ml.GetBucketsRequest;
+import org.elasticsearch.protocol.xpack.ml.GetBucketsResponse;
 import org.elasticsearch.protocol.xpack.ml.GetJobRequest;
 import org.elasticsearch.protocol.xpack.ml.GetJobResponse;
 import org.elasticsearch.protocol.xpack.ml.OpenJobRequest;
@@ -54,7 +56,7 @@ public final class MachineLearningClient {
      * For additional info
      * see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/ml-put-job.html">ML PUT job documentation</a>
      *
-     * @param request the PutJobRequest containing the {@link org.elasticsearch.protocol.xpack.ml.job.config.Job} settings
+     * @param request The PutJobRequest containing the {@link org.elasticsearch.protocol.xpack.ml.job.config.Job} settings
      * @param options Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
      * @return PutJobResponse with enclosed {@link org.elasticsearch.protocol.xpack.ml.job.config.Job} object
      * @throws IOException when there is a serialization issue sending the request or receiving the response
@@ -73,7 +75,7 @@ public PutJobResponse putJob(PutJobRequest request, RequestOptions options) thro
      * For additional info
      * see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/ml-put-job.html">ML PUT job documentation</a>
      *
-     * @param request  the request containing the {@link org.elasticsearch.protocol.xpack.ml.job.config.Job} settings
+     * @param request  The request containing the {@link org.elasticsearch.protocol.xpack.ml.job.config.Job} settings
      * @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
      */
@@ -93,7 +95,7 @@ public void putJobAsync(PutJobRequest request, RequestOptions options, ActionLis
      *     For additional info
      *     see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/ml-get-job.html"></a>
      * </p>
-     * @param request {@link GetJobRequest} request containing a list of jobId(s) and additional options
+     * @param request {@link GetJobRequest} Request containing a list of jobId(s) and additional options
      * @param options  Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
      * @return {@link GetJobResponse} response object containing
      * the {@link org.elasticsearch.protocol.xpack.ml.job.config.Job} objects and the number of jobs found
@@ -114,7 +116,7 @@ public GetJobResponse getJob(GetJobRequest request, RequestOptions options) thro
      *     For additional info
      *     see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/ml-get-job.html"></a>
      * </p>
-     * @param request {@link GetJobRequest} request containing a list of jobId(s) and additional options
+     * @param request {@link GetJobRequest} Request containing a list of jobId(s) and additional options
      * @param options  Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
      * @param listener Listener to be notified with {@link GetJobResponse} upon request completion
      */
@@ -133,7 +135,7 @@ public void getJobAsync(GetJobRequest request, RequestOptions options, ActionLis
      *     For additional info
      *     see <a href="http://www.elastic.co/guide/en/elasticsearch/reference/current/ml-delete-job.html">ML Delete Job documentation</a>
      * </p>
-     * @param request the request to delete the job
+     * @param request The request to delete the job
      * @param options  Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
      * @return action acknowledgement
      * @throws IOException when there is a serialization issue sending the request or receiving the response
@@ -152,7 +154,7 @@ public DeleteJobResponse deleteJob(DeleteJobRequest request, RequestOptions opti
      *     For additional info
      *     see <a href="http://www.elastic.co/guide/en/elasticsearch/reference/current/ml-delete-job.html">ML Delete Job documentation</a>
      * </p>
-     * @param request the request to delete the job
+     * @param request The request to delete the job
      * @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
      */
@@ -176,7 +178,7 @@ public void deleteJobAsync(DeleteJobRequest request, RequestOptions options, Act
      *     For additional info
      *     see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/ml-open-job.html"></a>
      * </p>
-     * @param request request containing job_id and additional optional options
+     * @param request Request containing job_id and additional optional options
      * @param options  Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
      * @return response containing if the job was successfully opened or not.
      * @throws IOException when there is a serialization issue sending the request or receiving the response
@@ -199,7 +201,7 @@ public OpenJobResponse openJob(OpenJobRequest request, RequestOptions options) t
      *     For additional info
      *     see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/ml-open-job.html"></a>
      * </p>
-     * @param request request containing job_id and additional optional options
+     * @param request Request containing job_id and additional optional options
      * @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
      */
@@ -217,7 +219,7 @@ public void openJobAsync(OpenJobRequest request, RequestOptions options, ActionL
      *
      * A closed job cannot receive data or perform analysis operations, but you can still explore and navigate results.
      *
-     * @param request request containing job_ids and additional options. See {@link CloseJobRequest}
+     * @param request Request containing job_ids and additional options. See {@link CloseJobRequest}
      * @param options  Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
      * @return response containing if the job was successfully closed or not.
      * @throws IOException when there is a serialization issue sending the request or receiving the response
@@ -235,7 +237,7 @@ public CloseJobResponse closeJob(CloseJobRequest request, RequestOptions options
      *
      * A closed job cannot receive data or perform analysis operations, but you can still explore and navigate results.
      *
-     * @param request request containing job_ids and additional options. See {@link CloseJobRequest}
+     * @param request Request containing job_ids and additional options. See {@link CloseJobRequest}
      * @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
      */
@@ -247,4 +249,40 @@ public void closeJobAsync(CloseJobRequest request, RequestOptions options, Actio
             listener,
             Collections.emptySet());
     }
+
+    /**
+     * Gets the buckets for a Machine Learning Job.
+     * <p>
+     * For additional info
+     * see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/ml-get-bucket.html">ML GET buckets documentation</a>
+     *
+     * @param request  The request
+     * @param options  Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
+     */
+    public GetBucketsResponse getBuckets(GetBucketsRequest request, RequestOptions options) throws IOException {
+        return restHighLevelClient.performRequestAndParseEntity(request,
+                MLRequestConverters::getBuckets,
+                options,
+                GetBucketsResponse::fromXContent,
+                Collections.emptySet());
+    }
+
+    /**
+     * Gets the buckets for a Machine Learning Job, notifies listener once the requested buckets are retrieved.
+     * <p>
+     * For additional info
+     * see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/ml-get-bucket.html">ML GET buckets 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 getBucketsAsync(GetBucketsRequest request, RequestOptions options, ActionListener<GetBucketsResponse> listener) {
+        restHighLevelClient.performRequestAsyncAndParseEntity(request,
+                MLRequestConverters::getBuckets,
+                options,
+                GetBucketsResponse::fromXContent,
+                listener,
+                Collections.emptySet());
+     }
 }

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

@@ -22,17 +22,20 @@
 import org.apache.http.client.methods.HttpDelete;
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpPost;
+import org.apache.http.client.methods.HttpPut;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.json.JsonXContent;
 import org.elasticsearch.protocol.xpack.ml.CloseJobRequest;
 import org.elasticsearch.protocol.xpack.ml.DeleteJobRequest;
+import org.elasticsearch.protocol.xpack.ml.GetBucketsRequest;
 import org.elasticsearch.protocol.xpack.ml.GetJobRequest;
 import org.elasticsearch.protocol.xpack.ml.OpenJobRequest;
 import org.elasticsearch.protocol.xpack.ml.PutJobRequest;
 import org.elasticsearch.protocol.xpack.ml.job.config.AnalysisConfig;
 import org.elasticsearch.protocol.xpack.ml.job.config.Detector;
 import org.elasticsearch.protocol.xpack.ml.job.config.Job;
+import org.elasticsearch.protocol.xpack.ml.job.util.PageParams;
 import org.elasticsearch.test.ESTestCase;
 
 import java.io.ByteArrayOutputStream;
@@ -49,6 +52,7 @@ public void testPutJob() throws IOException {
 
         Request request = MLRequestConverters.putJob(putJobRequest);
 
+        assertEquals(HttpPut.METHOD_NAME, request.getMethod());
         assertThat(request.getEndpoint(), equalTo("/_xpack/ml/anomaly_detectors/foo"));
         try (XContentParser parser = createParser(JsonXContent.jsonXContent, request.getEntity().getContent())) {
             Job parsedJob = Job.PARSER.apply(parser, null).build();
@@ -118,6 +122,23 @@ public void testDeleteJob() {
         assertEquals(Boolean.toString(true), request.getParameters().get("force"));
     }
 
+    public void testGetBuckets() throws IOException {
+        String jobId = randomAlphaOfLength(10);
+        GetBucketsRequest getBucketsRequest = new GetBucketsRequest(jobId);
+        getBucketsRequest.setPageParams(new PageParams(100, 300));
+        getBucketsRequest.setAnomalyScore(75.0);
+        getBucketsRequest.setSort("anomaly_score");
+        getBucketsRequest.setDescending(true);
+
+        Request request = MLRequestConverters.getBuckets(getBucketsRequest);
+        assertEquals(HttpGet.METHOD_NAME, request.getMethod());
+        assertEquals("/_xpack/ml/anomaly_detectors/" + jobId + "/results/buckets", request.getEndpoint());
+        try (XContentParser parser = createParser(JsonXContent.jsonXContent, request.getEntity().getContent())) {
+            GetBucketsRequest parsedRequest = GetBucketsRequest.PARSER.apply(parser, null);
+            assertThat(parsedRequest, equalTo(getBucketsRequest));
+        }
+    }
+
     private static Job createValidJob(String jobId) {
         AnalysisConfig.Builder analysisConfig = AnalysisConfig.builder(Collections.singletonList(
                 Detector.builder().setFunction("count").build()));