Skip to content

Commit eaa2718

Browse files
ruai0511ruai0511
committed
Add Update QueryGroup API Logic (#14775)
Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * Add Update QueryGroup API Logic Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * append to changlog Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * add javadoc Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * rebase Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * address comments Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * address comments Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * fix UT Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * adress comments Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * address comments Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> (cherry picked from commit 0753461)
1 parent a3312b4 commit eaa2718

25 files changed

+1292
-201
lines changed

‎CHANGELOG.md

+1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
1919
- Add `rangeQuery` and `regexpQuery` for `constant_keyword` field type ([#14711](https://github.com/opensearch-project/OpenSearch/pull/14711))
2020
- Add took time to request nodes stats ([#15054](https://github.com/opensearch-project/OpenSearch/pull/15054))
2121
- [Workload Management] Add Get QueryGroup API Logic ([14709](https://github.com/opensearch-project/OpenSearch/pull/14709))
22+
- [Workload Management] Add Update QueryGroup API Logic ([#14775](https://github.com/opensearch-project/OpenSearch/pull/14775))
2223
- [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897))
2324
- Support filtering on a large list encoded by bitmap ([#14774](https://github.com/opensearch-project/OpenSearch/pull/14774))
2425
- 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

+11-2
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,12 @@
2424
import org.opensearch.plugin.wlm.action.TransportCreateQueryGroupAction;
2525
import org.opensearch.plugin.wlm.action.TransportDeleteQueryGroupAction;
2626
import org.opensearch.plugin.wlm.action.TransportGetQueryGroupAction;
27+
import org.opensearch.plugin.wlm.action.TransportUpdateQueryGroupAction;
28+
import org.opensearch.plugin.wlm.action.UpdateQueryGroupAction;
2729
import org.opensearch.plugin.wlm.rest.RestCreateQueryGroupAction;
2830
import org.opensearch.plugin.wlm.rest.RestDeleteQueryGroupAction;
2931
import org.opensearch.plugin.wlm.rest.RestGetQueryGroupAction;
32+
import org.opensearch.plugin.wlm.rest.RestUpdateQueryGroupAction;
3033
import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService;
3134
import org.opensearch.plugins.ActionPlugin;
3235
import org.opensearch.plugins.Plugin;
@@ -52,7 +55,8 @@ public WorkloadManagementPlugin() {}
5255
return List.of(
5356
new ActionPlugin.ActionHandler<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class),
5457
new ActionPlugin.ActionHandler<>(GetQueryGroupAction.INSTANCE, TransportGetQueryGroupAction.class),
55-
new ActionPlugin.ActionHandler<>(DeleteQueryGroupAction.INSTANCE, TransportDeleteQueryGroupAction.class)
58+
new ActionPlugin.ActionHandler<>(DeleteQueryGroupAction.INSTANCE, TransportDeleteQueryGroupAction.class),
59+
new ActionPlugin.ActionHandler<>(UpdateQueryGroupAction.INSTANCE, TransportUpdateQueryGroupAction.class)
5660
);
5761
}
5862

@@ -66,7 +70,12 @@ public List<RestHandler> getRestHandlers(
6670
IndexNameExpressionResolver indexNameExpressionResolver,
6771
Supplier<DiscoveryNodes> nodesInCluster
6872
) {
69-
return List.of(new RestCreateQueryGroupAction(), new RestGetQueryGroupAction(), new RestDeleteQueryGroupAction());
73+
return List.of(
74+
new RestCreateQueryGroupAction(),
75+
new RestGetQueryGroupAction(),
76+
new RestDeleteQueryGroupAction(),
77+
new RestUpdateQueryGroupAction()
78+
);
7079
}
7180

7281
@Override

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -40,15 +40,15 @@ public class CreateQueryGroupRequest extends ActionRequest {
4040
* Constructor for CreateQueryGroupRequest
4141
* @param queryGroup - A {@link QueryGroup} object
4242
*/
43-
public CreateQueryGroupRequest(QueryGroup queryGroup) {
43+
CreateQueryGroupRequest(QueryGroup queryGroup) {
4444
this.queryGroup = queryGroup;
4545
}
4646

4747
/**
4848
* Constructor for CreateQueryGroupRequest
4949
* @param in - A {@link StreamInput} object
5050
*/
51-
public CreateQueryGroupRequest(StreamInput in) throws IOException {
51+
CreateQueryGroupRequest(StreamInput in) throws IOException {
5252
super(in);
5353
queryGroup = new QueryGroup(in);
5454
}

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

+51
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.plugin.wlm.action;
10+
11+
import org.opensearch.action.support.ActionFilters;
12+
import org.opensearch.action.support.HandledTransportAction;
13+
import org.opensearch.common.inject.Inject;
14+
import org.opensearch.core.action.ActionListener;
15+
import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService;
16+
import org.opensearch.tasks.Task;
17+
import org.opensearch.transport.TransportService;
18+
19+
/**
20+
* Transport action to update QueryGroup
21+
*
22+
* @opensearch.experimental
23+
*/
24+
public class TransportUpdateQueryGroupAction extends HandledTransportAction<UpdateQueryGroupRequest, UpdateQueryGroupResponse> {
25+
26+
private final QueryGroupPersistenceService queryGroupPersistenceService;
27+
28+
/**
29+
* Constructor for TransportUpdateQueryGroupAction
30+
*
31+
* @param actionName - action name
32+
* @param transportService - a {@link TransportService} object
33+
* @param actionFilters - a {@link ActionFilters} object
34+
* @param queryGroupPersistenceService - a {@link QueryGroupPersistenceService} object
35+
*/
36+
@Inject
37+
public TransportUpdateQueryGroupAction(
38+
String actionName,
39+
TransportService transportService,
40+
ActionFilters actionFilters,
41+
QueryGroupPersistenceService queryGroupPersistenceService
42+
) {
43+
super(UpdateQueryGroupAction.NAME, transportService, actionFilters, UpdateQueryGroupRequest::new);
44+
this.queryGroupPersistenceService = queryGroupPersistenceService;
45+
}
46+
47+
@Override
48+
protected void doExecute(Task task, UpdateQueryGroupRequest request, ActionListener<UpdateQueryGroupResponse> listener) {
49+
queryGroupPersistenceService.updateInClusterStateMetadata(request, listener);
50+
}
51+
}

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

+36
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.plugin.wlm.action;
10+
11+
import org.opensearch.action.ActionType;
12+
13+
/**
14+
* Transport action to update QueryGroup
15+
*
16+
* @opensearch.experimental
17+
*/
18+
public class UpdateQueryGroupAction extends ActionType<UpdateQueryGroupResponse> {
19+
20+
/**
21+
* An instance of UpdateQueryGroupAction
22+
*/
23+
public static final UpdateQueryGroupAction INSTANCE = new UpdateQueryGroupAction();
24+
25+
/**
26+
* Name for UpdateQueryGroupAction
27+
*/
28+
public static final String NAME = "cluster:admin/opensearch/wlm/query_group/_update";
29+
30+
/**
31+
* Default constructor
32+
*/
33+
private UpdateQueryGroupAction() {
34+
super(NAME, UpdateQueryGroupResponse::new);
35+
}
36+
}

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

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.plugin.wlm.action;
10+
11+
import org.opensearch.action.ActionRequest;
12+
import org.opensearch.action.ActionRequestValidationException;
13+
import org.opensearch.cluster.metadata.QueryGroup;
14+
import org.opensearch.core.common.io.stream.StreamInput;
15+
import org.opensearch.core.common.io.stream.StreamOutput;
16+
import org.opensearch.core.xcontent.XContentParser;
17+
import org.opensearch.wlm.MutableQueryGroupFragment;
18+
19+
import java.io.IOException;
20+
21+
/**
22+
* A request for update QueryGroup
23+
*
24+
* @opensearch.experimental
25+
*/
26+
public class UpdateQueryGroupRequest extends ActionRequest {
27+
private final String name;
28+
private final MutableQueryGroupFragment mutableQueryGroupFragment;
29+
30+
/**
31+
* Constructor for UpdateQueryGroupRequest
32+
* @param name - QueryGroup name for UpdateQueryGroupRequest
33+
* @param mutableQueryGroupFragment - MutableQueryGroupFragment for UpdateQueryGroupRequest
34+
*/
35+
UpdateQueryGroupRequest(String name, MutableQueryGroupFragment mutableQueryGroupFragment) {
36+
this.name = name;
37+
this.mutableQueryGroupFragment = mutableQueryGroupFragment;
38+
}
39+
40+
/**
41+
* Constructor for UpdateQueryGroupRequest
42+
* @param in - A {@link StreamInput} object
43+
*/
44+
UpdateQueryGroupRequest(StreamInput in) throws IOException {
45+
this(in.readString(), new MutableQueryGroupFragment(in));
46+
}
47+
48+
/**
49+
* Generate a UpdateQueryGroupRequest from XContent
50+
* @param parser - A {@link XContentParser} object
51+
* @param name - name of the QueryGroup to be updated
52+
*/
53+
public static UpdateQueryGroupRequest fromXContent(XContentParser parser, String name) throws IOException {
54+
QueryGroup.Builder builder = QueryGroup.Builder.fromXContent(parser);
55+
return new UpdateQueryGroupRequest(name, builder.getMutableQueryGroupFragment());
56+
}
57+
58+
@Override
59+
public ActionRequestValidationException validate() {
60+
QueryGroup.validateName(name);
61+
return null;
62+
}
63+
64+
/**
65+
* name getter
66+
*/
67+
public String getName() {
68+
return name;
69+
}
70+
71+
/**
72+
* mutableQueryGroupFragment getter
73+
*/
74+
public MutableQueryGroupFragment getmMutableQueryGroupFragment() {
75+
return mutableQueryGroupFragment;
76+
}
77+
78+
@Override
79+
public void writeTo(StreamOutput out) throws IOException {
80+
out.writeString(name);
81+
mutableQueryGroupFragment.writeTo(out);
82+
}
83+
}

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

+74
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.plugin.wlm.action;
10+
11+
import org.opensearch.cluster.metadata.QueryGroup;
12+
import org.opensearch.core.action.ActionResponse;
13+
import org.opensearch.core.common.io.stream.StreamInput;
14+
import org.opensearch.core.common.io.stream.StreamOutput;
15+
import org.opensearch.core.rest.RestStatus;
16+
import org.opensearch.core.xcontent.ToXContent;
17+
import org.opensearch.core.xcontent.ToXContentObject;
18+
import org.opensearch.core.xcontent.XContentBuilder;
19+
20+
import java.io.IOException;
21+
22+
/**
23+
* Response for the update API for QueryGroup
24+
*
25+
* @opensearch.experimental
26+
*/
27+
public class UpdateQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject {
28+
private final QueryGroup queryGroup;
29+
private final RestStatus restStatus;
30+
31+
/**
32+
* Constructor for UpdateQueryGroupResponse
33+
* @param queryGroup - the QueryGroup to be updated
34+
* @param restStatus - the rest status for the response
35+
*/
36+
public UpdateQueryGroupResponse(final QueryGroup queryGroup, RestStatus restStatus) {
37+
this.queryGroup = queryGroup;
38+
this.restStatus = restStatus;
39+
}
40+
41+
/**
42+
* Constructor for UpdateQueryGroupResponse
43+
* @param in - a {@link StreamInput} object
44+
*/
45+
public UpdateQueryGroupResponse(StreamInput in) throws IOException {
46+
queryGroup = new QueryGroup(in);
47+
restStatus = RestStatus.readFrom(in);
48+
}
49+
50+
@Override
51+
public void writeTo(StreamOutput out) throws IOException {
52+
queryGroup.writeTo(out);
53+
RestStatus.writeTo(out, restStatus);
54+
}
55+
56+
@Override
57+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
58+
return queryGroup.toXContent(builder, params);
59+
}
60+
61+
/**
62+
* queryGroup getter
63+
*/
64+
public QueryGroup getQueryGroup() {
65+
return queryGroup;
66+
}
67+
68+
/**
69+
* restStatus getter
70+
*/
71+
public RestStatus getRestStatus() {
72+
return restStatus;
73+
}
74+
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
import static org.opensearch.rest.RestRequest.Method.GET;
2828

2929
/**
30-
* Rest action to get a QueryGroup0
30+
* Rest action to get a QueryGroup
3131
*
3232
* @opensearch.experimental
3333
*/

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

+72
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.plugin.wlm.rest;
10+
11+
import org.opensearch.client.node.NodeClient;
12+
import org.opensearch.core.rest.RestStatus;
13+
import org.opensearch.core.xcontent.ToXContent;
14+
import org.opensearch.core.xcontent.XContentParser;
15+
import org.opensearch.plugin.wlm.action.UpdateQueryGroupAction;
16+
import org.opensearch.plugin.wlm.action.UpdateQueryGroupRequest;
17+
import org.opensearch.plugin.wlm.action.UpdateQueryGroupResponse;
18+
import org.opensearch.rest.BaseRestHandler;
19+
import org.opensearch.rest.BytesRestResponse;
20+
import org.opensearch.rest.RestChannel;
21+
import org.opensearch.rest.RestRequest;
22+
import org.opensearch.rest.RestResponse;
23+
import org.opensearch.rest.action.RestResponseListener;
24+
25+
import java.io.IOException;
26+
import java.util.List;
27+
28+
import static org.opensearch.rest.RestRequest.Method.POST;
29+
import static org.opensearch.rest.RestRequest.Method.PUT;
30+
31+
/**
32+
* Rest action to update a QueryGroup
33+
*
34+
* @opensearch.experimental
35+
*/
36+
public class RestUpdateQueryGroupAction extends BaseRestHandler {
37+
38+
/**
39+
* Constructor for RestUpdateQueryGroupAction
40+
*/
41+
public RestUpdateQueryGroupAction() {}
42+
43+
@Override
44+
public String getName() {
45+
return "update_query_group";
46+
}
47+
48+
/**
49+
* The list of {@link Route}s that this RestHandler is responsible for handling.
50+
*/
51+
@Override
52+
public List<Route> routes() {
53+
return List.of(new Route(POST, "_wlm/query_group/{name}"), new Route(PUT, "_wlm/query_group/{name}"));
54+
}
55+
56+
@Override
57+
protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
58+
try (XContentParser parser = request.contentParser()) {
59+
UpdateQueryGroupRequest updateQueryGroupRequest = UpdateQueryGroupRequest.fromXContent(parser, request.param("name"));
60+
return channel -> client.execute(UpdateQueryGroupAction.INSTANCE, updateQueryGroupRequest, updateQueryGroupResponse(channel));
61+
}
62+
}
63+
64+
private RestResponseListener<UpdateQueryGroupResponse> updateQueryGroupResponse(final RestChannel channel) {
65+
return new RestResponseListener<>(channel) {
66+
@Override
67+
public RestResponse buildResponse(final UpdateQueryGroupResponse response) throws Exception {
68+
return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS));
69+
}
70+
};
71+
}
72+
}

0 commit comments

Comments
 (0)