Add Get QueryGroup API Logic (opensearch-project#14709)

* Add Get QueryGroup API Logic Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * add to changelog Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * fix javadoc Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * change GetQueryGroupAction NAME and add more tests Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * add more unit tests Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * fix spotlessapply Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * addressed comments Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * incorperate comments from create api PR Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * use clustermanager to get the most recent querygroups Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * address comments Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * rebase with main Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * add IT Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * address comments Signed-off-by: Ruirui Zhang <mariazrr@amazon.com> * fix IT Signed-off-by: Ruirui Zhang <mariazrr@amazon.com>
wdongyu · Aug 22, 2024 · 76bf9b7 · 76bf9b7
1 parent 9dfc996
commit 76bf9b7
Show file tree

Hide file tree

Showing 16 changed files with 738 additions and 9 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Concurrent Segment Search] Support composite aggregations with scripting ([#15072](https://github.com/opensearch-project/OpenSearch/pull/15072))
 - 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] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897))
 - Add slice execution listeners to SearchOperationListener interface ([#15153](https://github.com/opensearch-project/OpenSearch/pull/15153))
 

diff --git a/...workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/...workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java
@@ -18,8 +18,11 @@
 import org.opensearch.common.settings.SettingsFilter;
 import org.opensearch.core.action.ActionResponse;
 import org.opensearch.plugin.wlm.action.CreateQueryGroupAction;
+import org.opensearch.plugin.wlm.action.GetQueryGroupAction;
 import org.opensearch.plugin.wlm.action.TransportCreateQueryGroupAction;
+import org.opensearch.plugin.wlm.action.TransportGetQueryGroupAction;
 import org.opensearch.plugin.wlm.rest.RestCreateQueryGroupAction;
+import org.opensearch.plugin.wlm.rest.RestGetQueryGroupAction;
 import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService;
 import org.opensearch.plugins.ActionPlugin;
 import org.opensearch.plugins.Plugin;
@@ -41,7 +44,10 @@ public WorkloadManagementPlugin() {}
 
     @Override
     public List<ActionHandler<? extends ActionRequest, ? extends ActionResponse>> getActions() {
-        return List.of(new ActionPlugin.ActionHandler<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class));
+        return List.of(
+            new ActionPlugin.ActionHandler<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class),
+            new ActionPlugin.ActionHandler<>(GetQueryGroupAction.INSTANCE, TransportGetQueryGroupAction.class)
+        );
     }
 
     @Override
@@ -54,7 +60,7 @@ public List<RestHandler> getRestHandlers(
         IndexNameExpressionResolver indexNameExpressionResolver,
         Supplier<DiscoveryNodes> nodesInCluster
     ) {
-        return List.of(new RestCreateQueryGroupAction());
+        return List.of(new RestCreateQueryGroupAction(), new RestGetQueryGroupAction());
     }
 
     @Override

diff --git a/...rkload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupAction.java b/...rkload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupAction.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 get QueryGroup
+ *
+ * @opensearch.experimental
+ */
+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/wlm/query_group/_get";
+
+    /**
+     * Default constructor
+     */
+    private GetQueryGroupAction() {
+        super(NAME, GetQueryGroupResponse::new);
+    }
+}
diff --git a/...kload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupRequest.java b/...kload-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupRequest.java
@@ -0,0 +1,64 @@
+/*
+ * 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.ActionRequestValidationException;
+import org.opensearch.action.support.clustermanager.ClusterManagerNodeReadRequest;
+import org.opensearch.cluster.metadata.QueryGroup;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+
+/**
+ * Request for get QueryGroup
+ *
+ * @opensearch.experimental
+ */
+public class GetQueryGroupRequest extends ClusterManagerNodeReadRequest<GetQueryGroupRequest> {
+    final 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 ActionRequestValidationException validate() {
+        if (name != null) {
+            QueryGroup.validateName(name);
+        }
+        return null;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeOptionalString(name);
+    }
+
+    /**
+     * Name getter
+     */
+    public String getName() {
+        return name;
+    }
+}
diff --git a/...load-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponse.java b/...load-management/src/main/java/org/opensearch/plugin/wlm/action/GetQueryGroupResponse.java
@@ -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.Collection;
+
+/**
+ * Response for the get API for QueryGroup
+ *
+ * @opensearch.experimental
+ */
+public class GetQueryGroupResponse extends ActionResponse implements ToXContent, ToXContentObject {
+    private final Collection<QueryGroup> queryGroups;
+    private final RestStatus restStatus;
+
+    /**
+     * Constructor for GetQueryGroupResponse
+     * @param queryGroups - The QueryGroup list to be fetched
+     * @param restStatus - The rest status of the request
+     */
+    public GetQueryGroupResponse(final Collection<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.writeCollection(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 Collection<QueryGroup> getQueryGroups() {
+        return queryGroups;
+    }
+
+    /**
+     * restStatus getter
+     */
+    public RestStatus getRestStatus() {
+        return restStatus;
+    }
+}
diff --git a/...nagement/src/main/java/org/opensearch/plugin/wlm/action/TransportGetQueryGroupAction.java b/...nagement/src/main/java/org/opensearch/plugin/wlm/action/TransportGetQueryGroupAction.java
@@ -0,0 +1,98 @@
+/*
+ * 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.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.ResourceNotFoundException;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeReadAction;
+import org.opensearch.cluster.ClusterState;
+import org.opensearch.cluster.block.ClusterBlockException;
+import org.opensearch.cluster.block.ClusterBlockLevel;
+import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
+import org.opensearch.cluster.metadata.QueryGroup;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.rest.RestStatus;
+import org.opensearch.plugin.wlm.service.QueryGroupPersistenceService;
+import org.opensearch.search.pipeline.SearchPipelineService;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportService;
+
+import java.io.IOException;
+import java.util.Collection;
+
+/**
+ * Transport action to get QueryGroup
+ *
+ * @opensearch.experimental
+ */
+public class TransportGetQueryGroupAction extends TransportClusterManagerNodeReadAction<GetQueryGroupRequest, GetQueryGroupResponse> {
+    private static final Logger logger = LogManager.getLogger(SearchPipelineService.class);
+
+    /**
+     * Constructor for TransportGetQueryGroupAction
+     *
+     * @param clusterService - a {@link ClusterService} object
+     * @param transportService - a {@link TransportService} object
+     * @param actionFilters - a {@link ActionFilters} object
+     * @param threadPool - a {@link ThreadPool} object
+     * @param indexNameExpressionResolver - a {@link IndexNameExpressionResolver} object
+     */
+    @Inject
+    public TransportGetQueryGroupAction(
+        ClusterService clusterService,
+        TransportService transportService,
+        ActionFilters actionFilters,
+        ThreadPool threadPool,
+        IndexNameExpressionResolver indexNameExpressionResolver
+    ) {
+        super(
+            GetQueryGroupAction.NAME,
+            transportService,
+            clusterService,
+            threadPool,
+            actionFilters,
+            GetQueryGroupRequest::new,
+            indexNameExpressionResolver,
+            true
+        );
+    }
+
+    @Override
+    protected String executor() {
+        return ThreadPool.Names.SAME;
+    }
+
+    @Override
+    protected GetQueryGroupResponse read(StreamInput in) throws IOException {
+        return new GetQueryGroupResponse(in);
+    }
+
+    @Override
+    protected ClusterBlockException checkBlock(GetQueryGroupRequest request, ClusterState state) {
+        return state.blocks().globalBlockedException(ClusterBlockLevel.METADATA_READ);
+    }
+
+    @Override
+    protected void clusterManagerOperation(GetQueryGroupRequest request, ClusterState state, ActionListener<GetQueryGroupResponse> listener)
+        throws Exception {
+        final String name = request.getName();
+        final Collection<QueryGroup> resultGroups = QueryGroupPersistenceService.getFromClusterStateMetadata(name, state);
+
+        if (resultGroups.isEmpty() && name != null && !name.isEmpty()) {
+            logger.warn("No QueryGroup exists with the provided name: {}", name);
+            throw new ResourceNotFoundException("No QueryGroup exists with the provided name: " + name);
+        }
+        listener.onResponse(new GetQueryGroupResponse(resultGroups, RestStatus.OK));
+    }
+}
diff --git a/...load-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetQueryGroupAction.java b/...load-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetQueryGroupAction.java
@@ -0,0 +1,68 @@
+/*
+ * 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.plugin.wlm.action.GetQueryGroupAction;
+import org.opensearch.plugin.wlm.action.GetQueryGroupRequest;
+import org.opensearch.plugin.wlm.action.GetQueryGroupResponse;
+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.GET;
+
+/**
+ * Rest action to get a QueryGroup0
+ *
+ * @opensearch.experimental
+ */
+public class RestGetQueryGroupAction extends BaseRestHandler {
+
+    /**
+     * Constructor for RestGetQueryGroupAction
+     */
+    public RestGetQueryGroupAction() {}
+
+    @Override
+    public String getName() {
+        return "get_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(GET, "_wlm/query_group/{name}"), new Route(GET, "_wlm/query_group/"));
+    }
+
+    @Override
+    protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        final GetQueryGroupRequest getQueryGroupRequest = new GetQueryGroupRequest(request.param("name"));
+        return channel -> client.execute(GetQueryGroupAction.INSTANCE, getQueryGroupRequest, getQueryGroupResponse(channel));
+    }
+
+    private RestResponseListener<GetQueryGroupResponse> getQueryGroupResponse(final RestChannel channel) {
+        return new RestResponseListener<>(channel) {
+            @Override
+            public RestResponse buildResponse(final GetQueryGroupResponse response) throws Exception {
+                return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS));
+            }
+        };
+    }
+}