Add Update QueryGroup API Logic (opensearch-project#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)
kaushalmahi12 · Sep 4, 2024 · 5b80e73 · 5b80e73
1 parent 7040df2
commit 5b80e73
Show file tree

Hide file tree

Showing 25 changed files with 3,070 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Remote Store] Rate limiter for remote store low priority uploads ([#14374](https://github.com/opensearch-project/OpenSearch/pull/14374/))
 - Apply the date histogram rewrite optimization to range aggregation ([#13865](https://github.com/opensearch-project/OpenSearch/pull/13865))
 - [Writable Warm] Add composite directory implementation and integrate it with FileCache ([12782](https://github.com/opensearch-project/OpenSearch/pull/12782))
+- [Workload Management] Add Update QueryGroup API Logic ([#14775](https://github.com/opensearch-project/OpenSearch/pull/14775))
 - Add batching supported processor base type AbstractBatchingProcessor ([#14554](https://github.com/opensearch-project/OpenSearch/pull/14554))
 - Fix race condition while parsing derived fields from search definition ([14445](https://github.com/opensearch-project/OpenSearch/pull/14445))
 - Add allowlist setting for ingest-common and search-pipeline-common processors ([#14439](https://github.com/opensearch-project/OpenSearch/issues/14439))

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
@@ -0,0 +1,90 @@
+/*
+ * 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;
+
+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.Setting;
+import org.opensearch.common.settings.Settings;
+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.DeleteQueryGroupAction;
+import org.opensearch.plugin.wlm.action.GetQueryGroupAction;
+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;
+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<>(CreateQueryGroupAction.INSTANCE, TransportCreateQueryGroupAction.class),
+            new ActionPlugin.ActionHandler<>(GetQueryGroupAction.INSTANCE, TransportGetQueryGroupAction.class),
+            new ActionPlugin.ActionHandler<>(DeleteQueryGroupAction.INSTANCE, TransportDeleteQueryGroupAction.class),
+            new ActionPlugin.ActionHandler<>(UpdateQueryGroupAction.INSTANCE, TransportUpdateQueryGroupAction.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 RestCreateQueryGroupAction(),
+            new RestGetQueryGroupAction(),
+            new RestDeleteQueryGroupAction(),
+            new RestUpdateQueryGroupAction()
+        );
+    }
+
+    @Override
+    public List<Setting<?>> getSettings() {
+        return List.of(QueryGroupPersistenceService.MAX_QUERY_GROUP_COUNT);
+    }
+
+    @Override
+    public Collection<Module> createGuiceModules() {
+        return List.of(new WorkloadManagementPluginModule());
+    }
+}
diff --git a/...ad-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java b/...ad-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.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.action.ActionRequest;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.cluster.metadata.QueryGroup;
+import org.opensearch.common.UUIDs;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.xcontent.XContentParser;
+import org.joda.time.Instant;
+
+import java.io.IOException;
+
+/**
+ * A request for create QueryGroup
+ * User input schema:
+ *  {
+ *    "name": "analytics",
+ *    "resiliency_mode": "enforced",
+ *    "resource_limits": {
+ *           "cpu" : 0.4,
+ *           "memory" : 0.2
+ *      }
+ *  }
+ *
+ * @opensearch.experimental
+ */
+public class CreateQueryGroupRequest extends ActionRequest {
+    private final QueryGroup queryGroup;
+
+    /**
+     * Constructor for CreateQueryGroupRequest
+     * @param queryGroup - A {@link QueryGroup} object
+     */
+    CreateQueryGroupRequest(QueryGroup queryGroup) {
+        this.queryGroup = queryGroup;
+    }
+
+    /**
+     * Constructor for CreateQueryGroupRequest
+     * @param in - A {@link StreamInput} object
+     */
+    CreateQueryGroupRequest(StreamInput in) throws IOException {
+        super(in);
+        queryGroup = new QueryGroup(in);
+    }
+
+    /**
+     * Generate a CreateQueryGroupRequest from XContent
+     * @param parser - A {@link XContentParser} object
+     */
+    public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws IOException {
+        QueryGroup.Builder builder = QueryGroup.Builder.fromXContent(parser);
+        return new CreateQueryGroupRequest(builder._id(UUIDs.randomBase64UUID()).updatedAt(Instant.now().getMillis()).build());
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        return null;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        queryGroup.writeTo(out);
+    }
+
+    /**
+     * QueryGroup getter
+     */
+    public QueryGroup getQueryGroup() {
+        return queryGroup;
+    }
+}
diff --git a/...ement/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateQueryGroupAction.java b/...ement/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);
+    }
+}
diff --git a/...oad-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupAction.java b/...oad-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);
+    }
+}
diff --git a/...ad-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateQueryGroupRequest.java b/...ad-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);
+    }
+}