bug fix

Signed-off-by: Ruirui Zhang <[email protected]>

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupAction.java

@@ -9,13 +9,14 @@
 package org.opensearch.plugin.wlm.action;
 
 import org.opensearch.action.ActionType;
+import org.opensearch.action.support.master.AcknowledgedResponse;
 
 /**
  * Transport action to create QueryGroup
  *
  * @opensearch.experimental
  */
-public class CreateQueryGroupAction extends ActionType<CreateQueryGroupResponse> {
+public class CreateQueryGroupAction extends ActionType<AcknowledgedResponse> {
 
     /**
      * An instance of CreateQueryGroupAction
@@ -31,6 +32,6 @@ public class CreateQueryGroupAction extends ActionType<CreateQueryGroupResponse>
      * Default constructor
      */
     private CreateQueryGroupAction() {
-        super(NAME, CreateQueryGroupResponse::new);
+        super(NAME, AcknowledgedResponse::new);
     }
 }

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java

@@ -8,9 +8,8 @@
 
 package org.opensearch.plugin.wlm.action;
 
-import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRequestValidationException;
-import org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest;
+import org.opensearch.action.support.master.AcknowledgedRequest;
 import org.opensearch.cluster.metadata.QueryGroup;
 import org.opensearch.common.UUIDs;
 import org.opensearch.core.common.io.stream.StreamInput;
@@ -34,14 +33,14 @@
  *
  * @opensearch.experimental
  */
-public class CreateQueryGroupRequest extends ClusterManagerNodeRequest<CreateQueryGroupRequest> {
+public class CreateQueryGroupRequest extends AcknowledgedRequest<CreateQueryGroupRequest> {
     private final QueryGroup queryGroup;
 
     /**
      * Constructor for CreateQueryGroupRequest
      * @param queryGroup - A {@link QueryGroup} object
      */
-    CreateQueryGroupRequest(QueryGroup queryGroup) {
+    public CreateQueryGroupRequest(QueryGroup queryGroup) {
         this.queryGroup = queryGroup;
     }
 

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java

@@ -9,8 +9,8 @@
 package org.opensearch.plugin.wlm.action;
 
 import org.opensearch.action.support.ActionFilters;
-import org.opensearch.action.support.HandledTransportAction;
 import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction;
+import org.opensearch.action.support.master.AcknowledgedResponse;
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.block.ClusterBlockException;
 import org.opensearch.cluster.block.ClusterBlockLevel;
@@ -20,7 +20,6 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService;
-import org.opensearch.tasks.Task;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportService;
 
@@ -31,7 +30,7 @@
  *
  * @opensearch.experimental
  */
-public class TransportCreateQueryGroupAction extends TransportClusterManagerNodeAction<CreateQueryGroupRequest, CreateQueryGroupResponse> {
+public class TransportCreateQueryGroupAction extends TransportClusterManagerNodeAction<CreateQueryGroupRequest, AcknowledgedResponse> {
 
     private final QueryGroupPersistenceService queryGroupPersistenceService;
 
@@ -70,9 +69,9 @@ public TransportCreateQueryGroupAction(
     protected void clusterManagerOperation(
         CreateQueryGroupRequest request,
         ClusterState state,
-        ActionListener<CreateQueryGroupResponse> listener
+        ActionListener<AcknowledgedResponse> listener
     ) throws Exception {
-        queryGroupPersistenceService.persistInClusterStateMetadata(request.getQueryGroup(), listener);
+        queryGroupPersistenceService.persistInClusterStateMetadata(request, listener);
     }
 
     @Override
@@ -81,8 +80,8 @@ protected String executor() {
     }
 
     @Override
-    protected CreateQueryGroupResponse read(StreamInput in) throws IOException {
-        return new CreateQueryGroupResponse(in);
+    protected AcknowledgedResponse read(StreamInput in) throws IOException {
+        return new AcknowledgedResponse(in);
     }
 
     @Override

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateQueryGroupAction.java

@@ -9,7 +9,6 @@
 package org.opensearch.plugin.wlm.action;
 
 import org.opensearch.action.support.ActionFilters;
-import org.opensearch.action.support.HandledTransportAction;
 import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction;
 import org.opensearch.action.support.master.AcknowledgedResponse;
 import org.opensearch.cluster.ClusterState;
@@ -21,7 +20,6 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService;
-import org.opensearch.tasks.Task;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportService;
 
@@ -32,7 +30,7 @@
  *
  * @opensearch.experimental
  */
-public class TransportUpdateQueryGroupAction extends TransportClusterManagerNodeAction<UpdateQueryGroupRequest, UpdateQueryGroupResponse> {
+public class TransportUpdateQueryGroupAction extends TransportClusterManagerNodeAction<UpdateQueryGroupRequest, AcknowledgedResponse> {
 
     private final QueryGroupPersistenceService queryGroupPersistenceService;
 
@@ -71,7 +69,7 @@ public TransportUpdateQueryGroupAction(
     protected void clusterManagerOperation(
         UpdateQueryGroupRequest request,
         ClusterState state,
-        ActionListener<UpdateQueryGroupResponse> listener
+        ActionListener<AcknowledgedResponse> listener
     ) throws Exception {
         queryGroupPersistenceService.updateInClusterStateMetadata(request, listener);
     }
@@ -82,8 +80,8 @@ protected String executor() {
     }
 
     @Override
-    protected UpdateQueryGroupResponse read(StreamInput in) throws IOException {
-        return new UpdateQueryGroupResponse(in);
+    protected AcknowledgedResponse read(StreamInput in) throws IOException {
+        return new AcknowledgedResponse(in);
     }
 
     @Override

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupAction.java

@@ -9,13 +9,14 @@
 package org.opensearch.plugin.wlm.action;
 
 import org.opensearch.action.ActionType;
+import org.opensearch.action.support.master.AcknowledgedResponse;
 
 /**
  * Transport action to update QueryGroup
  *
  * @opensearch.experimental
  */
-public class UpdateQueryGroupAction extends ActionType<UpdateQueryGroupResponse> {
+public class UpdateQueryGroupAction extends ActionType<AcknowledgedResponse> {
 
     /**
      * An instance of UpdateQueryGroupAction
@@ -31,6 +32,6 @@ public class UpdateQueryGroupAction extends ActionType<UpdateQueryGroupResponse>
      * Default constructor
      */
     private UpdateQueryGroupAction() {
-        super(NAME, UpdateQueryGroupResponse::new);
+        super(NAME, AcknowledgedResponse::new);
     }
 }

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java

@@ -8,9 +8,8 @@
 
 package org.opensearch.plugin.wlm.action;
 
-import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRequestValidationException;
-import org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest;
+import org.opensearch.action.support.master.AcknowledgedRequest;
 import org.opensearch.cluster.metadata.QueryGroup;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
@@ -24,7 +23,7 @@
  *
  * @opensearch.experimental
  */
-public class UpdateQueryGroupRequest extends ClusterManagerNodeRequest<UpdateQueryGroupRequest> {
+public class UpdateQueryGroupRequest extends AcknowledgedRequest<UpdateQueryGroupRequest> {
     private final String name;
     private final MutableQueryGroupFragment mutableQueryGroupFragment;
 

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateQueryGroupAction.java

@@ -9,18 +9,12 @@
 package org.opensearch.plugin.wlm.rest;
 
 import org.opensearch.client.node.NodeClient;
-import org.opensearch.core.rest.RestStatus;
-import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.plugin.wlm.action.CreateQueryGroupAction;
 import org.opensearch.plugin.wlm.action.CreateQueryGroupRequest;
-import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse;
 import org.opensearch.rest.BaseRestHandler;
-import org.opensearch.rest.BytesRestResponse;
-import org.opensearch.rest.RestChannel;
 import org.opensearch.rest.RestRequest;
-import org.opensearch.rest.RestResponse;
-import org.opensearch.rest.action.RestResponseListener;
+import org.opensearch.rest.action.RestToXContentListener;
 
 import java.io.IOException;
 import java.util.List;
@@ -57,16 +51,15 @@ public List<Route> routes() {
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
         try (XContentParser parser = request.contentParser()) {
             CreateQueryGroupRequest createQueryGroupRequest = CreateQueryGroupRequest.fromXContent(parser);
-            return channel -> client.execute(CreateQueryGroupAction.INSTANCE, createQueryGroupRequest, createQueryGroupResponse(channel));
+            createQueryGroupRequest.clusterManagerNodeTimeout(
+                request.paramAsTime("cluster_manager_timeout", createQueryGroupRequest.clusterManagerNodeTimeout())
+            );
+            createQueryGroupRequest.timeout(request.paramAsTime("timeout", createQueryGroupRequest.timeout()));
+            return channel -> client.execute(
+                CreateQueryGroupAction.INSTANCE,
+                createQueryGroupRequest,
+                new RestToXContentListener<>(channel)
+            );
         }
     }
-
-    private RestResponseListener<CreateQueryGroupResponse> createQueryGroupResponse(final RestChannel channel) {
-        return new RestResponseListener<>(channel) {
-            @Override
-            public RestResponse buildResponse(final CreateQueryGroupResponse response) throws Exception {
-                return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS));
-            }
-        };
-    }
 }