[Backport 2.17] Add Update QueryGroup API Logic (#14775)

CHANGELOG.md

@@ -19,6 +19,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add `rangeQuery` and `regexpQuery` for `constant_keyword` field type ([#14711](https://github.com/opensearch-project/OpenSearch/pull/14711))
 - Add took time to request nodes stats ([#15054](https://github.com/opensearch-project/OpenSearch/pull/15054))
 - [Workload Management] Add Get QueryGroup API Logic ([14709](https://github.com/opensearch-project/OpenSearch/pull/14709))
+- [Workload Management] Add Update QueryGroup API Logic ([#14775](https://github.com/opensearch-project/OpenSearch/pull/14775))
 - [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897))
 - Support filtering on a large list encoded by bitmap ([#14774](https://github.com/opensearch-project/OpenSearch/pull/14774))
 - Add slice execution listeners to SearchOperationListener interface ([#15153](https://github.com/opensearch-project/OpenSearch/pull/15153))

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

@@ -24,9 +24,12 @@
 import org.opensearch.plugin.wlm.action.TransportCreateQueryGroupAction;
 import org.opensearch.plugin.wlm.action.TransportDeleteQueryGroupAction;
 import org.opensearch.plugin.wlm.action.TransportGetQueryGroupAction;
+import org.opensearch.plugin.wlm.action.TransportUpdateQueryGroupAction;
+import org.opensearch.plugin.wlm.action.UpdateQueryGroupAction;
 import org.opensearch.plugin.wlm.rest.RestCreateQueryGroupAction;
 import org.opensearch.plugin.wlm.rest.RestDeleteQueryGroupAction;
 import org.opensearch.plugin.wlm.rest.RestGetQueryGroupAction;
+import org.opensearch.plugin.wlm.rest.RestUpdateQueryGroupAction;
 import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService;
 import org.opensearch.plugins.ActionPlugin;
 import org.opensearch.plugins.Plugin;
@@ -52,7 +55,8 @@
         return List.of(
             new ActionPlugin.ActionHandler<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class),
             new ActionPlugin.ActionHandler<>(GetQueryGroupAction.INSTANCE, TransportGetQueryGroupAction.class),
-            new ActionPlugin.ActionHandler<>(DeleteQueryGroupAction.INSTANCE, TransportDeleteQueryGroupAction.class)
+            new ActionPlugin.ActionHandler<>(DeleteQueryGroupAction.INSTANCE, TransportDeleteQueryGroupAction.class),
+            new ActionPlugin.ActionHandler<>(UpdateQueryGroupAction.INSTANCE, TransportUpdateQueryGroupAction.class)
         );
     }
 
@@ -66,7 +70,12 @@
         IndexNameExpressionResolver indexNameExpressionResolver,
         Supplier<DiscoveryNodes> nodesInCluster
     ) {
-        return List.of(new RestCreateQueryGroupAction(), new RestGetQueryGroupAction(), new RestDeleteQueryGroupAction());
+        return List.of(
+            new RestCreateQueryGroupAction(),
+            new RestGetQueryGroupAction(),
+            new RestDeleteQueryGroupAction(),
+            new RestUpdateQueryGroupAction()
+        );
     }
 
     @Override

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

@@ -40,15 +40,15 @@ public class CreateQueryGroupRequest extends ActionRequest {
      * Constructor for CreateQueryGroupRequest
      * @param queryGroup - A {@link QueryGroup} object
      */
-    public CreateQueryGroupRequest(QueryGroup queryGroup) {
+    CreateQueryGroupRequest(QueryGroup queryGroup) {
         this.queryGroup = queryGroup;
     }
 
     /**
      * Constructor for CreateQueryGroupRequest
      * @param in - A {@link StreamInput} object
      */
-    public CreateQueryGroupRequest(StreamInput in) throws IOException {
+    CreateQueryGroupRequest(StreamInput in) throws IOException {
         super(in);
         queryGroup = new QueryGroup(in);
     }

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

@@ -0,0 +1,51 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.wlm.action;
+
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService;
+import org.opensearch.tasks.Task;
+import org.opensearch.transport.TransportService;
+
+/**
+ * Transport action to update QueryGroup
+ *
+ * @opensearch.experimental
+ */
+public class TransportUpdateQueryGroupAction extends HandledTransportAction<UpdateQueryGroupRequest, UpdateQueryGroupResponse> {
+
+    private final QueryGroupPersistenceService queryGroupPersistenceService;
+
+    /**
+     * Constructor for TransportUpdateQueryGroupAction
+     *
+     * @param actionName - action name
+     * @param transportService - a {@link TransportService} object
+     * @param actionFilters - a {@link ActionFilters} object
+     * @param queryGroupPersistenceService - a {@link QueryGroupPersistenceService} object
+     */
+    @Inject
+    public TransportUpdateQueryGroupAction(
+        String actionName,
+        TransportService transportService,
+        ActionFilters actionFilters,
+        QueryGroupPersistenceService queryGroupPersistenceService
+    ) {
+        super(UpdateQueryGroupAction.NAME, transportService, actionFilters, UpdateQueryGroupRequest::new);
+        this.queryGroupPersistenceService = queryGroupPersistenceService;
+    }
+
+    @Override
+    protected void doExecute(Task task, UpdateQueryGroupRequest request, ActionListener<UpdateQueryGroupResponse> listener) {
+        queryGroupPersistenceService.updateInClusterStateMetadata(request, listener);
+    }
+}

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

@@ -0,0 +1,36 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.wlm.action;
+
+import org.opensearch.action.ActionType;
+
+/**
+ * Transport action to update QueryGroup
+ *
+ * @opensearch.experimental
+ */
+public class UpdateQueryGroupAction extends ActionType<UpdateQueryGroupResponse> {
+
+    /**
+     * An instance of UpdateQueryGroupAction
+     */
+    public static final UpdateQueryGroupAction INSTANCE = new UpdateQueryGroupAction();
+
+    /**
+     * Name for UpdateQueryGroupAction
+     */
+    public static final String NAME = "cluster:admin/opensearch/wlm/query_group/_update";
+
+    /**
+     * Default constructor
+     */
+    private UpdateQueryGroupAction() {
+        super(NAME, UpdateQueryGroupResponse::new);
+    }
+}

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

@@ -0,0 +1,83 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.wlm.action;
+
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.cluster.metadata.QueryGroup;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.wlm.MutableQueryGroupFragment;
+
+import java.io.IOException;
+
+/**
+ * A request for update QueryGroup
+ *
+ * @opensearch.experimental
+ */
+public class UpdateQueryGroupRequest extends ActionRequest {
+    private final String name;
+    private final MutableQueryGroupFragment mutableQueryGroupFragment;
+
+    /**
+     * Constructor for UpdateQueryGroupRequest
+     * @param name - QueryGroup name for UpdateQueryGroupRequest
+     * @param mutableQueryGroupFragment - MutableQueryGroupFragment for UpdateQueryGroupRequest
+     */
+    UpdateQueryGroupRequest(String name, MutableQueryGroupFragment mutableQueryGroupFragment) {
+        this.name = name;
+        this.mutableQueryGroupFragment = mutableQueryGroupFragment;
+    }
+
+    /**
+     * Constructor for UpdateQueryGroupRequest
+     * @param in - A {@link StreamInput} object
+     */
+    UpdateQueryGroupRequest(StreamInput in) throws IOException {
+        this(in.readString(), new MutableQueryGroupFragment(in));
+    }
+
+    /**
+     * Generate a UpdateQueryGroupRequest from XContent
+     * @param parser - A {@link XContentParser} object
+     * @param name - name of the QueryGroup to be updated
+     */
+    public static UpdateQueryGroupRequest fromXContent(XContentParser parser, String name) throws IOException {
+        QueryGroup.Builder builder = QueryGroup.Builder.fromXContent(parser);
+        return new UpdateQueryGroupRequest(name, builder.getMutableQueryGroupFragment());
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        QueryGroup.validateName(name);
+        return null;
+    }
+
+    /**
+     * name getter
+     */
+    public String getName() {
+        return name;
+    }
+
+    /**
+     * mutableQueryGroupFragment getter
+     */
+    public MutableQueryGroupFragment getmMutableQueryGroupFragment() {
+        return mutableQueryGroupFragment;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeString(name);
+        mutableQueryGroupFragment.writeTo(out);
+    }
+}

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

@@ -0,0 +1,74 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.wlm.action;
+
+import org.opensearch.cluster.metadata.QueryGroup;
+import org.opensearch.core.action.ActionResponse;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.rest.RestStatus;
+import org.opensearch.core.xcontent.ToXContent;
+import org.opensearch.core.xcontent.ToXContentObject;
+import org.opensearch.core.xcontent.XContentBuilder;
+
+import java.io.IOException;
+
+/**
+ * Response for the update API for QueryGroup
+ *
+ * @opensearch.experimental
+ */
+public class UpdateQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject {
+    private final QueryGroup queryGroup;
+    private final RestStatus restStatus;
+
+    /**
+     * Constructor for UpdateQueryGroupResponse
+     * @param queryGroup - the QueryGroup to be updated
+     * @param restStatus - the rest status for the response
+     */
+    public UpdateQueryGroupResponse(final QueryGroup queryGroup, RestStatus restStatus) {
+        this.queryGroup = queryGroup;
+        this.restStatus = restStatus;
+    }
+
+    /**
+     * Constructor for UpdateQueryGroupResponse
+     * @param in - a {@link StreamInput} object
+     */
+    public UpdateQueryGroupResponse(StreamInput in) throws IOException {
+        queryGroup = new QueryGroup(in);
+        restStatus = RestStatus.readFrom(in);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        queryGroup.writeTo(out);
+        RestStatus.writeTo(out, restStatus);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        return queryGroup.toXContent(builder, params);
+    }
+
+    /**
+     * queryGroup getter
+     */
+    public QueryGroup getQueryGroup() {
+        return queryGroup;
+    }
+
+    /**
+     * restStatus getter
+     */
+    public RestStatus getRestStatus() {
+        return restStatus;
+    }
+}

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

@@ -27,7 +27,7 @@
 import static org.opensearch.rest.RestRequest.Method.GET;
 
 /**
- * Rest action to get a QueryGroup0
+ * Rest action to get a QueryGroup
  *
  * @opensearch.experimental
  */

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

@@ -0,0 +1,72 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+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.UpdateQueryGroupAction;
+import org.opensearch.plugin.wlm.action.UpdateQueryGroupRequest;
+import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse;
+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 java.io.IOException;
+import java.util.List;
+
+import static org.opensearch.rest.RestRequest.Method.POST;
+import static org.opensearch.rest.RestRequest.Method.PUT;
+
+/**
+ * Rest action to update a QueryGroup
+ *
+ * @opensearch.experimental
+ */
+public class RestUpdateQueryGroupAction extends BaseRestHandler {
+
+    /**
+     * Constructor for RestUpdateQueryGroupAction
+     */
+    public RestUpdateQueryGroupAction() {}
+
+    @Override
+    public String getName() {
+        return "update_query_group";
+    }
+
+    /**
+     * The list of {@link Route}s that this RestHandler is responsible for handling.
+     */
+    @Override
+    public List<Route> routes() {
+        return List.of(new Route(POST, "_wlm/query_group/{name}"), new Route(PUT, "_wlm/query_group/{name}"));
+    }
+
+    @Override
+    protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        try (XContentParser parser = request.contentParser()) {
+            UpdateQueryGroupRequest updateQueryGroupRequest = UpdateQueryGroupRequest.fromXContent(parser, request.param("name"));
+            return channel -> client.execute(UpdateQueryGroupAction.INSTANCE, updateQueryGroupRequest, updateQueryGroupResponse(channel));
+        }
+    }
+
+    private RestResponseListener<UpdateQueryGroupResponse> updateQueryGroupResponse(final RestChannel channel) {
+        return new RestResponseListener<>(channel) {
+            @Override
+            public RestResponse buildResponse(final UpdateQueryGroupResponse response) throws Exception {
+                return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS));
+            }
+        };
+    }
+}