Skip to content

Commit 5b80e73

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 7040df2 commit 5b80e73

25 files changed

+3070
-0
lines changed

‎CHANGELOG.md

+1
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
99
- [Remote Store] Rate limiter for remote store low priority uploads ([#14374](https://github.com/opensearch-project/OpenSearch/pull/14374/))
1010
- Apply the date histogram rewrite optimization to range aggregation ([#13865](https://github.com/opensearch-project/OpenSearch/pull/13865))
1111
- [Writable Warm] Add composite directory implementation and integrate it with FileCache ([12782](https://github.com/opensearch-project/OpenSearch/pull/12782))
12+
- [Workload Management] Add Update QueryGroup API Logic ([#14775](https://github.com/opensearch-project/OpenSearch/pull/14775))
1213
- Add batching supported processor base type AbstractBatchingProcessor ([#14554](https://github.com/opensearch-project/OpenSearch/pull/14554))
1314
- Fix race condition while parsing derived fields from search definition ([14445](https://github.com/opensearch-project/OpenSearch/pull/14445))
1415
- Add allowlist setting for ingest-common and search-pipeline-common processors ([#14439](https://github.com/opensearch-project/OpenSearch/issues/14439))

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

+90
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
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;
10+
11+
import org.opensearch.action.ActionRequest;
12+
import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
13+
import org.opensearch.cluster.node.DiscoveryNodes;
14+
import org.opensearch.common.inject.Module;
15+
import org.opensearch.common.settings.ClusterSettings;
16+
import org.opensearch.common.settings.IndexScopedSettings;
17+
import org.opensearch.common.settings.Setting;
18+
import org.opensearch.common.settings.Settings;
19+
import org.opensearch.common.settings.SettingsFilter;
20+
import org.opensearch.core.action.ActionResponse;
21+
import org.opensearch.plugin.wlm.action.CreateQueryGroupAction;
22+
import org.opensearch.plugin.wlm.action.DeleteQueryGroupAction;
23+
import org.opensearch.plugin.wlm.action.GetQueryGroupAction;
24+
import org.opensearch.plugin.wlm.action.TransportCreateQueryGroupAction;
25+
import org.opensearch.plugin.wlm.action.TransportDeleteQueryGroupAction;
26+
import org.opensearch.plugin.wlm.action.TransportGetQueryGroupAction;
27+
import org.opensearch.plugin.wlm.action.TransportUpdateQueryGroupAction;
28+
import org.opensearch.plugin.wlm.action.UpdateQueryGroupAction;
29+
import org.opensearch.plugin.wlm.rest.RestCreateQueryGroupAction;
30+
import org.opensearch.plugin.wlm.rest.RestDeleteQueryGroupAction;
31+
import org.opensearch.plugin.wlm.rest.RestGetQueryGroupAction;
32+
import org.opensearch.plugin.wlm.rest.RestUpdateQueryGroupAction;
33+
import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService;
34+
import org.opensearch.plugins.ActionPlugin;
35+
import org.opensearch.plugins.Plugin;
36+
import org.opensearch.rest.RestController;
37+
import org.opensearch.rest.RestHandler;
38+
39+
import java.util.Collection;
40+
import java.util.List;
41+
import java.util.function.Supplier;
42+
43+
/**
44+
* Plugin class for WorkloadManagement
45+
*/
46+
public class WorkloadManagementPlugin extends Plugin implements ActionPlugin {
47+
48+
/**
49+
* Default constructor
50+
*/
51+
public WorkloadManagementPlugin() {}
52+
53+
@Override
54+
public List<ActionHandler<? extends ActionRequest, ? extends ActionResponse>> getActions() {
55+
return List.of(
56+
new ActionPlugin.ActionHandler<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class),
57+
new ActionPlugin.ActionHandler<>(GetQueryGroupAction.INSTANCE, TransportGetQueryGroupAction.class),
58+
new ActionPlugin.ActionHandler<>(DeleteQueryGroupAction.INSTANCE, TransportDeleteQueryGroupAction.class),
59+
new ActionPlugin.ActionHandler<>(UpdateQueryGroupAction.INSTANCE, TransportUpdateQueryGroupAction.class)
60+
);
61+
}
62+
63+
@Override
64+
public List<RestHandler> getRestHandlers(
65+
Settings settings,
66+
RestController restController,
67+
ClusterSettings clusterSettings,
68+
IndexScopedSettings indexScopedSettings,
69+
SettingsFilter settingsFilter,
70+
IndexNameExpressionResolver indexNameExpressionResolver,
71+
Supplier<DiscoveryNodes> nodesInCluster
72+
) {
73+
return List.of(
74+
new RestCreateQueryGroupAction(),
75+
new RestGetQueryGroupAction(),
76+
new RestDeleteQueryGroupAction(),
77+
new RestUpdateQueryGroupAction()
78+
);
79+
}
80+
81+
@Override
82+
public List<Setting<?>> getSettings() {
83+
return List.of(QueryGroupPersistenceService.MAX_QUERY_GROUP_COUNT);
84+
}
85+
86+
@Override
87+
public Collection<Module> createGuiceModules() {
88+
return List.of(new WorkloadManagementPluginModule());
89+
}
90+
}

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

+82
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
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.common.UUIDs;
15+
import org.opensearch.core.common.io.stream.StreamInput;
16+
import org.opensearch.core.common.io.stream.StreamOutput;
17+
import org.opensearch.core.xcontent.XContentParser;
18+
import org.joda.time.Instant;
19+
20+
import java.io.IOException;
21+
22+
/**
23+
* A request for create QueryGroup
24+
* User input schema:
25+
* {
26+
* "name": "analytics",
27+
* "resiliency_mode": "enforced",
28+
* "resource_limits": {
29+
* "cpu" : 0.4,
30+
* "memory" : 0.2
31+
* }
32+
* }
33+
*
34+
* @opensearch.experimental
35+
*/
36+
public class CreateQueryGroupRequest extends ActionRequest {
37+
private final QueryGroup queryGroup;
38+
39+
/**
40+
* Constructor for CreateQueryGroupRequest
41+
* @param queryGroup - A {@link QueryGroup} object
42+
*/
43+
CreateQueryGroupRequest(QueryGroup queryGroup) {
44+
this.queryGroup = queryGroup;
45+
}
46+
47+
/**
48+
* Constructor for CreateQueryGroupRequest
49+
* @param in - A {@link StreamInput} object
50+
*/
51+
CreateQueryGroupRequest(StreamInput in) throws IOException {
52+
super(in);
53+
queryGroup = new QueryGroup(in);
54+
}
55+
56+
/**
57+
* Generate a CreateQueryGroupRequest from XContent
58+
* @param parser - A {@link XContentParser} object
59+
*/
60+
public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws IOException {
61+
QueryGroup.Builder builder = QueryGroup.Builder.fromXContent(parser);
62+
return new CreateQueryGroupRequest(builder._id(UUIDs.randomBase64UUID()).updatedAt(Instant.now().getMillis()).build());
63+
}
64+
65+
@Override
66+
public ActionRequestValidationException validate() {
67+
return null;
68+
}
69+
70+
@Override
71+
public void writeTo(StreamOutput out) throws IOException {
72+
super.writeTo(out);
73+
queryGroup.writeTo(out);
74+
}
75+
76+
/**
77+
* QueryGroup getter
78+
*/
79+
public QueryGroup getQueryGroup() {
80+
return queryGroup;
81+
}
82+
}

‎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+
}

0 commit comments

Comments
 (0)