Skip to content

Commit

Permalink
Add Get QueryGroup API Logic
Browse files Browse the repository at this point in the history
Signed-off-by: Ruirui Zhang <mariazrr@amazon.com>
  • Loading branch information
ruai0511 committed Jul 10, 2024
1 parent a04cf24 commit d7a452a
Show file tree
Hide file tree
Showing 21 changed files with 1,128 additions and 1 deletion.
18 changes: 18 additions & 0 deletions plugins/workload-management/build.gradle
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
/*
* 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.
*
* Modifications Copyright OpenSearch Contributors. See
* GitHub history for details.
*/

opensearchplugin {
description 'OpenSearch Workload Management Plugin.'
classname 'org.opensearch.plugin.wlm.action.WorkloadManagementPlugin'
}

dependencies {
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
/*
* 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 get QueryGroup
*
* @opensearch.api
*/
public class GetQueryGroupAction extends ActionType<GetQueryGroupResponse> {

/**
/**
* An instance of GetQueryGroupAction
*/
public static final GetQueryGroupAction INSTANCE = new GetQueryGroupAction();

/**
* Name for GetQueryGroupAction
*/
public static final String NAME = "cluster:admin/opensearch/query_group/wlm/_get";

/**
* Default constructor
*/
private GetQueryGroupAction() {
super(NAME, GetQueryGroupResponse::new);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
/*
* 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.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.common.io.stream.Writeable;

import java.io.IOException;

/**
* A request for get QueryGroup
*
* @opensearch.internal
*/
public class GetQueryGroupRequest extends ActionRequest implements Writeable.Reader<GetQueryGroupRequest> {
String name;

/**
* Default constructor for GetQueryGroupRequest
* @param name - name for the QueryGroup to get
*/
public GetQueryGroupRequest(String name) {
this.name = name;
}

/**
* Constructor for GetQueryGroupRequest
* @param in - A {@link StreamInput} object
*/
public GetQueryGroupRequest(StreamInput in) throws IOException {
super(in);
name = in.readOptionalString();
}

@Override
public GetQueryGroupRequest read(StreamInput in) throws IOException {
return new GetQueryGroupRequest(in);
}

@Override
public ActionRequestValidationException validate() {
return null;
}

/**
* Name getter
*/
public String getName() {
return name;
}

@Override
public void writeTo(StreamOutput out) throws IOException {
super.writeTo(out);
out.writeOptionalString(name);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
/*
* 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;
import java.util.List;

/**
* Response for the get API for QueryGroup
*
* @opensearch.internal
*/
public class GetQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject {
private final List<QueryGroup> queryGroups;
private RestStatus restStatus;

/**
* Constructor for GetQueryGroupResponse
* @param queryGroups - The QueryGroup list to be fetched
* @param restStatus - The rest status of the request
*/
public GetQueryGroupResponse(final List<QueryGroup> queryGroups, RestStatus restStatus) {
this.queryGroups = queryGroups;
this.restStatus = restStatus;
}

/**
* Constructor for GetQueryGroupResponse
* @param in - A {@link StreamInput} object
*/
public GetQueryGroupResponse(StreamInput in) throws IOException {
this.queryGroups = in.readList(QueryGroup::new);
restStatus = RestStatus.readFrom(in);
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeList(queryGroups);
RestStatus.writeTo(out, restStatus);
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject();
builder.startArray("query_groups");
for (QueryGroup group : queryGroups) {
group.toXContent(builder, params);
}
builder.endArray();
builder.endObject();
return builder;
}

/**
* queryGroups getter
*/
public List<QueryGroup> getQueryGroups() {
return queryGroups;
}

/**
* restStatus getter
*/
public RestStatus getRestStatus() {
return restStatus;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
/*
* 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.cluster.metadata.QueryGroup;
import org.opensearch.common.inject.Inject;
import org.opensearch.core.action.ActionListener;
import org.opensearch.plugin.wlm.action.service.Persistable;
import org.opensearch.tasks.Task;
import org.opensearch.threadpool.ThreadPool;
import org.opensearch.transport.TransportService;

/**
* Transport action for get QueryGroup
*
* @opensearch.internal
*/
public class TransportGetQueryGroupAction extends HandledTransportAction<GetQueryGroupRequest, GetQueryGroupResponse> {

private final ThreadPool threadPool;
private final Persistable<QueryGroup> queryGroupPersistenceService;

/**
* Constructor for TransportGetQueryGroupAction
*
* @param actionName - action name
* @param transportService - a {@link TransportService} object
* @param actionFilters - a {@link ActionFilters} object
* @param threadPool - a {@link ThreadPool} object
* @param queryGroupPersistenceService - a {@link Persistable} object
*/
@Inject
public TransportGetQueryGroupAction(
String actionName,
TransportService transportService,
ActionFilters actionFilters,
ThreadPool threadPool,
Persistable<QueryGroup> queryGroupPersistenceService
) {
super(GetQueryGroupAction.NAME, transportService, actionFilters, GetQueryGroupRequest::new);
this.threadPool = threadPool;
this.queryGroupPersistenceService = queryGroupPersistenceService;
}

@Override
protected void doExecute(Task task, GetQueryGroupRequest request, ActionListener<GetQueryGroupResponse> listener) {
String name = request.getName();
threadPool.executor(ThreadPool.Names.GENERIC).execute(() -> queryGroupPersistenceService.get(name, listener));
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
/*
* 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.cluster.metadata.IndexNameExpressionResolver;
import org.opensearch.cluster.node.DiscoveryNodes;
import org.opensearch.common.inject.Module;
import org.opensearch.common.settings.ClusterSettings;
import org.opensearch.common.settings.IndexScopedSettings;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.settings.SettingsFilter;
import org.opensearch.core.action.ActionResponse;
import org.opensearch.plugin.wlm.action.rest.RestGetQueryGroupAction;
import org.opensearch.plugins.ActionPlugin;
import org.opensearch.plugins.Plugin;
import org.opensearch.rest.RestController;
import org.opensearch.rest.RestHandler;

import java.util.Collection;
import java.util.List;
import java.util.function.Supplier;

/**
* Plugin class for WorkloadManagement
*/
public class WorkloadManagementPlugin extends Plugin implements ActionPlugin {

/**
* Default constructor
*/
public WorkloadManagementPlugin() {}

@Override
public List<ActionHandler<? extends ActionRequest, ? extends ActionResponse>> getActions() {
return List.of(new ActionPlugin.ActionHandler<>(GetQueryGroupAction.INSTANCE, TransportGetQueryGroupAction.class));
}

@Override
public List<RestHandler> getRestHandlers(
Settings settings,
RestController restController,
ClusterSettings clusterSettings,
IndexScopedSettings indexScopedSettings,
SettingsFilter settingsFilter,
IndexNameExpressionResolver indexNameExpressionResolver,
Supplier<DiscoveryNodes> nodesInCluster
) {
return List.of(new RestGetQueryGroupAction());
}

@Override
public Collection<Module> createGuiceModules() {
return List.of(new WorkloadManagementPluginModule());
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
/*
* 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.common.inject.AbstractModule;
import org.opensearch.common.inject.TypeLiteral;
import org.opensearch.plugin.wlm.action.service.Persistable;
import org.opensearch.plugin.wlm.action.service.QueryGroupPersistenceService;

/**
* Guice Module to manage WorkloadManagement related objects
*/
public class WorkloadManagementPluginModule extends AbstractModule {

/**
* Constructor for WorkloadManagementPluginModule
*/
public WorkloadManagementPluginModule() {}

@Override
protected void configure() {
bind(new TypeLiteral<Persistable<QueryGroup>>() {
}).to(QueryGroupPersistenceService.class).asEagerSingleton();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
/*
* 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.
*/

/**
* Base Package of CRUD API of QueryGroup
*/
package org.opensearch.plugin.wlm.action;
Loading

0 comments on commit d7a452a

Please sign in to comment.