From 7d6d91937f623d7fc86a2a68e7fa7f3ad376ec17 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Thu, 23 Jan 2025 13:35:31 -0800 Subject: [PATCH 01/20] add get rule api logic Signed-off-by: Ruirui Zhang --- .../plugin/wlm/WorkloadManagementPlugin.java | 8 +- .../wlm/WorkloadManagementPluginModule.java | 1 - .../action/CreateWorkloadGroupAction.java | 2 +- .../action/CreateWorkloadGroupRequest.java | 2 +- .../action/CreateWorkloadGroupResponse.java | 2 +- .../action/DeleteWorkloadGroupAction.java | 2 +- .../action/DeleteWorkloadGroupRequest.java | 2 +- .../action/GetWorkloadGroupAction.java | 2 +- .../action/GetWorkloadGroupRequest.java | 2 +- .../action/GetWorkloadGroupResponse.java | 2 +- .../TransportCreateWorkloadGroupAction.java | 6 +- .../TransportDeleteWorkloadGroupAction.java | 6 +- .../TransportGetWorkloadGroupAction.java | 6 +- .../TransportUpdateWorkloadGroupAction.java | 6 +- .../action/UpdateWorkloadGroupAction.java | 2 +- .../action/UpdateWorkloadGroupRequest.java | 0 .../action/UpdateWorkloadGroupResponse.java | 2 +- .../wlm/querygroup/action/package-info.java | 12 + .../rest/RestCreateWorkloadGroupAction.java | 8 +- .../rest/RestDeleteWorkloadGroupAction.java | 2 +- .../rest/RestGetWorkloadGroupAction.java | 8 +- .../rest/RestUpdateWorkloadGroupAction.java | 8 +- .../wlm/querygroup/rest/package-info.java | 12 + .../wlm/querygroup/service/package-info.java | 12 + .../plugin/wlm/rule/action/GetRuleAction.java | 35 ++ .../wlm/rule/action/GetRuleRequest.java | 75 ++++ .../wlm/rule/action/GetRuleResponse.java | 83 ++++ .../rule/action/TransportGetRuleAction.java | 47 +++ .../{rest => rule/action}/package-info.java | 2 +- .../wlm/rule/rest/RestGetRuleAction.java | 82 ++++ .../{action => rule/rest}/package-info.java | 2 +- .../rule/service/RulePersistenceService.java | 196 ++++++++++ .../wlm/{ => rule}/service/package-info.java | 2 +- .../WorkloadGroupPersistenceService.java | 10 +- .../opensearch/plugin/wlm/RuleTestUtils.java | 83 ++++ .../CreateWorkloadGroupRequestTests.java | 2 +- .../CreateWorkloadGroupResponseTests.java | 2 +- .../DeleteWorkloadGroupRequestTests.java | 2 +- .../action/GetWorkloadGroupRequestTests.java | 2 +- .../action/GetWorkloadGroupResponseTests.java | 2 +- ...ansportDeleteWorkloadGroupActionTests.java | 6 +- .../TransportGetWorkloadGroupActionTests.java | 2 +- .../UpdateWorkloadGroupRequestTests.java | 2 +- .../UpdateWorkloadGroupResponseTests.java | 2 +- .../action/WorkloadGroupActionTestUtils.java | 0 .../RestDeleteWorkloadGroupActionTests.java | 7 +- .../WorkloadGroupPersistenceServiceTests.java | 0 .../wlm/rule/action/GetRuleRequestTests.java | 53 +++ .../wlm/rule/action/GetRuleResponseTests.java | 118 ++++++ .../wlm/rule/rest/RestGetRuleActionTests.java | 35 ++ .../service/RulePersistenceServiceTests.java | 101 +++++ .../main/java/org/opensearch/wlm/Rule.java | 359 ++++++++++++++++++ .../java/org/opensearch/wlm/RuleTests.java | 163 ++++++++ 53 files changed, 1555 insertions(+), 35 deletions(-) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/CreateWorkloadGroupAction.java (94%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/CreateWorkloadGroupRequest.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/CreateWorkloadGroupResponse.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/DeleteWorkloadGroupAction.java (94%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/DeleteWorkloadGroupRequest.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/GetWorkloadGroupAction.java (94%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/GetWorkloadGroupRequest.java (96%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/GetWorkloadGroupResponse.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/TransportCreateWorkloadGroupAction.java (87%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/TransportDeleteWorkloadGroupAction.java (88%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/TransportGetWorkloadGroupAction.java (89%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/TransportUpdateWorkloadGroupAction.java (87%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/UpdateWorkloadGroupAction.java (94%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/UpdateWorkloadGroupRequest.java (100%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/action/UpdateWorkloadGroupResponse.java (97%) create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/package-info.java rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/rest/RestCreateWorkloadGroupAction.java (82%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/rest/RestDeleteWorkloadGroupAction.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/rest/RestGetWorkloadGroupAction.java (80%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => querygroup}/rest/RestUpdateWorkloadGroupAction.java (82%) create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/package-info.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/service/package-info.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{rest => rule/action}/package-info.java (86%) create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{action => rule/rest}/package-info.java (86%) create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{ => rule}/service/package-info.java (86%) create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/action/CreateWorkloadGroupRequestTests.java (96%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/action/CreateWorkloadGroupResponseTests.java (98%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/action/DeleteWorkloadGroupRequestTests.java (96%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/action/GetWorkloadGroupRequestTests.java (97%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/action/GetWorkloadGroupResponseTests.java (99%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/action/TransportDeleteWorkloadGroupActionTests.java (84%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/action/TransportGetWorkloadGroupActionTests.java (97%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/action/UpdateWorkloadGroupRequestTests.java (98%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/action/UpdateWorkloadGroupResponseTests.java (98%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/action/WorkloadGroupActionTestUtils.java (100%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/rest/RestDeleteWorkloadGroupActionTests.java (88%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{ => querygroup}/service/WorkloadGroupPersistenceServiceTests.java (100%) create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleActionTests.java create mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java create mode 100644 server/src/main/java/org/opensearch/wlm/Rule.java create mode 100644 server/src/test/java/org/opensearch/wlm/RuleTests.java diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index dedc02a06e0c6..21daded9e3069 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -63,6 +63,8 @@ import java.util.List; import java.util.function.Supplier; +import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.RULE_INDEX; + /** * Plugin class for WorkloadManagement */ @@ -125,7 +127,8 @@ public List getActionFilters() { new ActionPlugin.ActionHandler<>(CreateWorkloadGroupAction.INSTANCE, TransportCreateWorkloadGroupAction.class), new ActionPlugin.ActionHandler<>(GetWorkloadGroupAction.INSTANCE, TransportGetWorkloadGroupAction.class), new ActionPlugin.ActionHandler<>(DeleteWorkloadGroupAction.INSTANCE, TransportDeleteWorkloadGroupAction.class), - new ActionPlugin.ActionHandler<>(UpdateWorkloadGroupAction.INSTANCE, TransportUpdateWorkloadGroupAction.class) + new ActionPlugin.ActionHandler<>(UpdateWorkloadGroupAction.INSTANCE, TransportUpdateWorkloadGroupAction.class), + new ActionPlugin.ActionHandler<>(GetRuleAction.INSTANCE, TransportGetRuleAction.class) ); } @@ -148,7 +151,8 @@ public List getRestHandlers( new RestCreateWorkloadGroupAction(), new RestGetWorkloadGroupAction(), new RestDeleteWorkloadGroupAction(), - new RestUpdateWorkloadGroupAction() + new RestUpdateWorkloadGroupAction(), + new RestGetRuleAction() ); } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java index bb0f4c7e90122..0ed3781aae409 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java @@ -22,7 +22,6 @@ public class WorkloadManagementPluginModule extends AbstractModule { */ public WorkloadManagementPluginModule() {} - @Override protected void configure() { // Bind WorkloadGroupPersistenceService as a singleton to ensure a single instance is used, // preventing multiple throttling key registrations in the constructor. diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupAction.java similarity index 94% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupAction.java index ca9784ebc7e4b..433305f7e9bf6 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.ActionType; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequest.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequest.java index ad932667b25e8..b31cc6f0fb1d2 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequest.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponse.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponse.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponse.java index b33214e042398..d765b3ed14b65 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponse.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.core.action.ActionResponse; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupAction.java similarity index 94% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupAction.java index 39b47d69776f4..54755758a15a3 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.ActionType; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequest.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequest.java index 940a3815b1662..73a6cc0e1dcd7 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequest.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.action.support.clustermanager.AcknowledgedRequest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupAction.java similarity index 94% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupAction.java index ee1b40a2f9bbc..28f35ba2f18b4 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.ActionType; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequest.java similarity index 96% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequest.java index 4b8a5f85fd236..28ab53849824e 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequest.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.action.support.clustermanager.ClusterManagerNodeReadRequest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponse.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponse.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponse.java index ab8f773088a37..22b66bca9ae56 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponse.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.core.action.ActionResponse; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportCreateWorkloadGroupAction.java similarity index 87% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportCreateWorkloadGroupAction.java index 2039f1cb590ff..ae401af9769ee 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportCreateWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction; @@ -17,7 +17,11 @@ import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; +<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportCreateWorkloadGroupAction.java import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; +======== +import org.opensearch.plugin.wlm.querygroup.service.QueryGroupPersistenceService; +>>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportUpdateQueryGroupAction.java import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupAction.java similarity index 88% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupAction.java index 2bfbadba4d51d..40c82b9c0a56f 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; @@ -19,7 +19,11 @@ import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; +<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupAction.java import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; +======== +import org.opensearch.plugin.wlm.querygroup.service.QueryGroupPersistenceService; +>>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteQueryGroupAction.java import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupAction.java similarity index 89% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupAction.java index bb2fbab047343..ce64e737337cd 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -23,7 +23,11 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.rest.RestStatus; +<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupAction.java import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; +======== +import org.opensearch.plugin.wlm.querygroup.service.QueryGroupPersistenceService; +>>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetQueryGroupAction.java import org.opensearch.search.pipeline.SearchPipelineService; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportUpdateWorkloadGroupAction.java similarity index 87% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportUpdateWorkloadGroupAction.java index ef639d44b4155..b29896c0a0b1f 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportUpdateWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction; @@ -17,7 +17,11 @@ import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; +<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportUpdateWorkloadGroupAction.java import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; +======== +import org.opensearch.plugin.wlm.querygroup.service.QueryGroupPersistenceService; +>>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportCreateQueryGroupAction.java import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupAction.java similarity index 94% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupAction.java index b4f8e1ce90126..6e49008ff2fae 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.ActionType; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupRequest.java similarity index 100% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupRequest.java diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponse.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponse.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponse.java index 9b8fccbdb5346..2f9d543d5b1f9 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponse.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.core.action.ActionResponse; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/package-info.java new file mode 100644 index 0000000000000..472b41716b44d --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/package-info.java @@ -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. + */ + +/** + * Package for the action classes related to query groups in WorkloadManagementPlugin + */ +package org.opensearch.plugin.wlm.querygroup.action; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestCreateWorkloadGroupAction.java similarity index 82% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestCreateWorkloadGroupAction.java index 5ef59602f7893..94d4d82cfbb8e 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestCreateWorkloadGroupAction.java @@ -6,14 +6,20 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.rest; +package org.opensearch.plugin.wlm.querygroup.rest; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentParser; +<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestCreateWorkloadGroupAction.java import org.opensearch.plugin.wlm.action.CreateWorkloadGroupAction; import org.opensearch.plugin.wlm.action.CreateWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.CreateWorkloadGroupResponse; +======== +import org.opensearch.plugin.wlm.querygroup.action.UpdateQueryGroupAction; +import org.opensearch.plugin.wlm.querygroup.action.UpdateQueryGroupRequest; +import org.opensearch.plugin.wlm.querygroup.action.UpdateQueryGroupResponse; +>>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestUpdateQueryGroupAction.java import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupAction.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupAction.java index d0d82f43679fa..e3b5b8c30ea3b 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.rest; +package org.opensearch.plugin.wlm.querygroup.rest; import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupAction; import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupRequest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestGetWorkloadGroupAction.java similarity index 80% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestGetWorkloadGroupAction.java index 818531352f4d3..50dd7e0ea50db 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestGetWorkloadGroupAction.java @@ -6,13 +6,19 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.rest; +package org.opensearch.plugin.wlm.querygroup.rest; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; +<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestGetWorkloadGroupAction.java import org.opensearch.plugin.wlm.action.GetWorkloadGroupAction; import org.opensearch.plugin.wlm.action.GetWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.GetWorkloadGroupResponse; +======== +import org.opensearch.plugin.wlm.querygroup.action.GetQueryGroupAction; +import org.opensearch.plugin.wlm.querygroup.action.GetQueryGroupRequest; +import org.opensearch.plugin.wlm.querygroup.action.GetQueryGroupResponse; +>>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestGetQueryGroupAction.java import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestUpdateWorkloadGroupAction.java similarity index 82% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestUpdateWorkloadGroupAction.java index db77dc5963037..1294529d857cb 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestUpdateWorkloadGroupAction.java @@ -6,14 +6,20 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.rest; +package org.opensearch.plugin.wlm.querygroup.rest; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentParser; +<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestUpdateWorkloadGroupAction.java import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupAction; import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupResponse; +======== +import org.opensearch.plugin.wlm.querygroup.action.CreateQueryGroupAction; +import org.opensearch.plugin.wlm.querygroup.action.CreateQueryGroupRequest; +import org.opensearch.plugin.wlm.querygroup.action.CreateQueryGroupResponse; +>>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestCreateQueryGroupAction.java import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/package-info.java new file mode 100644 index 0000000000000..a51d67a6cb3b7 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/package-info.java @@ -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. + */ + +/** + * Package for the rest classes related to query groups in WorkloadManagementPlugin + */ +package org.opensearch.plugin.wlm.querygroup.rest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/service/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/service/package-info.java new file mode 100644 index 0000000000000..3758c9fcd9b81 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/service/package-info.java @@ -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. + */ + +/** + * Package for the service classes related to query groups in WorkloadManagementPlugin + */ +package org.opensearch.plugin.wlm.querygroup.service; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java new file mode 100644 index 0000000000000..f334b064acc62 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java @@ -0,0 +1,35 @@ +/* + * 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.rule.action; + +import org.opensearch.action.ActionType; + +/** + * Transport action to get Rule + * @opensearch.experimental + */ +public class GetRuleAction extends ActionType { + + /** + * An instance of GetRuleAction + */ + public static final GetRuleAction INSTANCE = new GetRuleAction(); + + /** + * Name for GetRuleAction + */ + public static final String NAME = "cluster:admin/opensearch/wlm/rule/_get"; + + /** + * Default constructor + */ + private GetRuleAction() { + super(NAME, GetRuleResponse::new); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java new file mode 100644 index 0000000000000..e821007b1d13b --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java @@ -0,0 +1,75 @@ +/* + * 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.rule.action; + +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.wlm.Rule.RuleAttribute; + +import java.io.IOException; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +/** + * A request for get Rule + * @opensearch.experimental + */ +public class GetRuleRequest extends ClusterManagerNodeRequest { + private final String _id; + private final Map> attributeFilters; + + /** + * Constructor for GetRuleRequest + * @param _id - Rule _id that we want to get + * @param attributeFilters - Attributes that we want to filter on + */ + public GetRuleRequest(String _id, Map> attributeFilters) { + this._id = _id; + this.attributeFilters = attributeFilters; + } + + /** + * Constructor for GetRuleRequest + * @param in - A {@link StreamInput} object + */ + public GetRuleRequest(StreamInput in) throws IOException { + super(in); + _id = in.readOptionalString(); + attributeFilters = in.readMap((i) -> RuleAttribute.fromName(i.readString()), i -> new HashSet<>(i.readStringList())); + } + + @Override + public ActionRequestValidationException validate() { + return null; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeOptionalString(_id); + out.writeMap(attributeFilters, RuleAttribute::writeTo, StreamOutput::writeStringCollection); + } + + /** + * _id getter + */ + public String get_id() { + return _id; + } + + /** + * attributeFilters getter + */ + public Map> getAttributeFilters() { + return attributeFilters; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java new file mode 100644 index 0000000000000..b8be558c7e567 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.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.rule.action; + +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 org.opensearch.wlm.Rule; + +import java.io.IOException; +import java.util.Map; + +import static org.opensearch.wlm.Rule._ID_STRING; + +/** + * Response for the get API for Rule + * @opensearch.experimental + */ +public class GetRuleResponse extends ActionResponse implements ToXContent, ToXContentObject { + private final Map rules; + private final RestStatus restStatus; + + /** + * Constructor for GetRuleResponse + * @param rules - The Map of Rules to be included in the response + * @param restStatus - The restStatus for the response + */ + public GetRuleResponse(final Map rules, RestStatus restStatus) { + this.rules = rules; + this.restStatus = restStatus; + } + + /** + * Constructor for GetRuleResponse + * @param in - A {@link StreamInput} object + */ + public GetRuleResponse(StreamInput in) throws IOException { + this.rules = in.readMap(StreamInput::readString, Rule::new); + this.restStatus = RestStatus.readFrom(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeMap(rules, StreamOutput::writeString, (outStream, rule) -> rule.writeTo(outStream)); + RestStatus.writeTo(out, restStatus); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.startArray("rules"); + for (Map.Entry entry : rules.entrySet()) { + entry.getValue().toXContent(builder, new MapParams(Map.of(_ID_STRING, entry.getKey()))); + } + builder.endArray(); + builder.endObject(); + return builder; + } + + /** + * rules getter + */ + public Map getRules() { + return rules; + } + + /** + * restStatus getter + */ + public RestStatus getRestStatus() { + return restStatus; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java new file mode 100644 index 0000000000000..66b243ff30e64 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java @@ -0,0 +1,47 @@ +/* + * 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.rule.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.rule.service.RulePersistenceService; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +/** + * Transport action to get Rule + * @opensearch.experimental + */ +public class TransportGetRuleAction extends HandledTransportAction { + + private final RulePersistenceService rulePersistenceService; + + /** + * Constructor for TransportGetRuleAction + * @param transportService - a {@link TransportService} object + * @param actionFilters - a {@link ActionFilters} object + * @param rulePersistenceService - a {@link RulePersistenceService} object + */ + @Inject + public TransportGetRuleAction( + TransportService transportService, + ActionFilters actionFilters, + RulePersistenceService rulePersistenceService + ) { + super(GetRuleAction.NAME, transportService, actionFilters, GetRuleRequest::new); + this.rulePersistenceService = rulePersistenceService; + } + + @Override + protected void doExecute(Task task, GetRuleRequest request, ActionListener listener) { + rulePersistenceService.getRule(request.get_id(), request.getAttributeFilters(), listener); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/package-info.java similarity index 86% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/package-info.java index 889f3e107db07..4fd4eb7ac8eb3 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/package-info.java @@ -9,4 +9,4 @@ /** * Package for the rest classes related to query groups in WorkloadManagementPlugin */ -package org.opensearch.plugin.wlm.rest; +package org.opensearch.plugin.wlm.rule.action; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java new file mode 100644 index 0000000000000..685fd01f200d5 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.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.rule.rest; + +import org.opensearch.client.node.NodeClient; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.plugin.wlm.rule.action.GetRuleAction; +import org.opensearch.plugin.wlm.rule.action.GetRuleRequest; +import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; +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 org.opensearch.wlm.Rule.RuleAttribute; + +import java.io.IOException; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.opensearch.rest.RestRequest.Method.GET; +import static org.opensearch.wlm.Rule._ID_STRING; + +/** + * Rest action to get a Rule + * @opensearch.experimental + */ +public class RestGetRuleAction extends BaseRestHandler { + + /** + * Constructor for RestGetRuleAction + */ + public RestGetRuleAction() {} + + @Override + public String getName() { + return "get_rule"; + } + + /** + * The list of {@link Route}s that this RestHandler is responsible for handling. + */ + @Override + public List routes() { + return List.of(new Route(GET, "_wlm/rule/"), new Route(GET, "_wlm/rule/{_id}")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + final Map> attributeFilters = new HashMap<>(); + for (String attributeName : request.params().keySet()) { + if (attributeName.equals(_ID_STRING)) { + continue; + } + String[] valuesArray = request.param(attributeName).split(","); + attributeFilters.put(RuleAttribute.fromName(attributeName), new HashSet<>(Arrays.asList(valuesArray))); + } + final GetRuleRequest getRuleRequest = new GetRuleRequest(request.param(_ID_STRING), attributeFilters); + return channel -> client.execute(GetRuleAction.INSTANCE, getRuleRequest, getRuleResponse(channel)); + } + + private RestResponseListener getRuleResponse(final RestChannel channel) { + return new RestResponseListener<>(channel) { + @Override + public RestResponse buildResponse(final GetRuleResponse response) throws Exception { + return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS)); + } + }; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/package-info.java similarity index 86% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/package-info.java index df30f55a99b3c..68d60df54d1b4 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/package-info.java @@ -9,4 +9,4 @@ /** * Package for the action classes related to query groups in WorkloadManagementPlugin */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.rule.rest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java new file mode 100644 index 0000000000000..4d00cb0decd27 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java @@ -0,0 +1,196 @@ +/* + * 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.rule.service; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.get.GetResponse; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.client.Client; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.DeprecationHandler; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; +import org.opensearch.wlm.Rule; +import org.opensearch.wlm.Rule.Builder; +import org.opensearch.wlm.Rule.RuleAttribute; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * This class defines the functions for Rule persistence + * @opensearch.experimental + */ +public class RulePersistenceService { + public static final String RULE_INDEX = ".rule"; + private final Client client; + private static final Logger logger = LogManager.getLogger(RulePersistenceService.class); + private static final int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 50; + + /** + * Constructor for RulePersistenceService + * @param client {@link Client} - The client to be used by RulePersistenceService + */ + @Inject + public RulePersistenceService(final Client client) { + this.client = client; + } + + /** + * Entry point for the get rule api logic in persistence service. + * @param id - The id of the rule to get. Get all matching rules when id is null + * @param attributeFilters - A map containing the attributes that user want to filter on + * @param listener - ActionListener for GetRuleResponse + */ + public void getRule(String id, Map> attributeFilters, ActionListener listener) { + if (id != null) { + fetchRuleById(id, listener); + } else { + fetchAllRules(attributeFilters, listener); + } + } + + /** + * Fetch a single rule from system index using id + * @param id - The id of the rule to get. + * @param listener - ActionListener for GetRuleResponse + */ + private void fetchRuleById(String id, ActionListener listener) { + ThreadContext.StoredContext storedContext = client.threadPool().getThreadContext().stashContext(); + client.prepareGet(RULE_INDEX, id) + .execute(ActionListener.wrap( + getResponse -> { + try (ThreadContext.StoredContext context = storedContext) { + handleGetOneRuleResponse(id, getResponse, listener); + } + }, + e -> { + try (ThreadContext.StoredContext context = storedContext) { + logger.error("Failed to fetch rule with ID {}: {}", id, e.getMessage()); + listener.onFailure(e); + } + } + )); + } + + /** + * Process getResponse from index and send a GetRuleResponse + * @param id - The id of the rule to get + * @param getResponse - Response received from index + * @param listener - ActionListener for GetRuleResponse + */ + private void handleGetOneRuleResponse(String id, GetResponse getResponse, ActionListener listener) { + if (getResponse.isExists()) { + try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { + XContentParser parser = MediaTypeRegistry.JSON.xContent() + .createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.THROW_UNSUPPORTED_OPERATION, + getResponse.getSourceAsString() + ); + listener.onResponse(new GetRuleResponse(Map.of(id, Builder.fromXContent(parser).build()), RestStatus.OK)); + } catch (IOException e) { + logger.error("Error parsing rule with ID {}: {}", id, e.getMessage()); + listener.onFailure(e); + } + } else { + listener.onFailure(new ResourceNotFoundException("Rule with ID " + id + " not found.")); + } + } + + /** + * Fetch all rule from system index based on attributeFilters + * @param attributeFilters - A map containing the attributes that user want to filter on + * @param listener - ActionListener for GetRuleResponse + */ + private void fetchAllRules(Map> attributeFilters, ActionListener listener) { + ThreadContext.StoredContext storedContext = client.threadPool().getThreadContext().stashContext(); + client.prepareSearch(RULE_INDEX) + .setQuery(QueryBuilders.matchAllQuery()) + .setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST) + .execute(ActionListener.wrap( + searchResponse -> { + try (ThreadContext.StoredContext context = storedContext) { + handleGetAllRuleResponse(searchResponse, attributeFilters, listener); + } + }, + e -> { + try (ThreadContext.StoredContext context = storedContext) { + logger.error("Failed to fetch all rules: {}", e.getMessage()); + listener.onFailure(e); + } + } + )); + } + + /** + * Process searchResponse from index and send a GetRuleResponse + * @param searchResponse - Response received from index + * @param attributeFilters - A map containing the attributes that user want to filter on + * @param listener - ActionListener for GetRuleResponse + */ + private void handleGetAllRuleResponse( + SearchResponse searchResponse, + Map> attributeFilters, + ActionListener listener + ) { + Map ruleMap = Arrays.stream(searchResponse.getHits().getHits()).map(hit -> { + try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { + XContentParser parser = MediaTypeRegistry.JSON.xContent() + .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, hit.getSourceAsString()); + Rule currRule = Rule.Builder.fromXContent(parser).build(); + if (matchesFilters(currRule, attributeFilters)) { + return Map.entry(hit.getId(), currRule); + } + return null; + } catch (IOException e) { + logger.error("Failed to parse rule from hit: {}", e.getMessage()); + listener.onFailure(e); + return null; + } + }).filter(Objects::nonNull).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + listener.onResponse(new GetRuleResponse(ruleMap, RestStatus.OK)); + } + + /** + * Returns true if the rule matches the attributeFilters and should be included in the response + * @param rule - the rule to be checked against the attribute filters + * @param attributeFilters - A map containing the attributes that user want to filter on + */ + private boolean matchesFilters(Rule rule, Map> attributeFilters) { + for (Map.Entry> entry : attributeFilters.entrySet()) { + RuleAttribute attribute = entry.getKey(); + Set expectedValues = entry.getValue(); + Set ruleValues = rule.getAttributeMap().get(attribute); + if (ruleValues == null || ruleValues.stream().noneMatch(expectedValues::contains)) { + return false; + } + } + return true; + } + + /** + * client getter + */ + public Client getClient() { + return client; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/package-info.java similarity index 86% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/package-info.java index e8c88ee656dc7..603145e617cf9 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/package-info.java @@ -9,4 +9,4 @@ /** * Package for the service classes related to query groups in WorkloadManagementPlugin */ -package org.opensearch.plugin.wlm.service; +package org.opensearch.plugin.wlm.rule.service; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java index 8fd5fe5dfcfed..02d4b5d7c97d2 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.service; +package org.opensearch.plugin.wlm.querygroup.service; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -27,11 +27,19 @@ import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; +<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java import org.opensearch.plugin.wlm.action.CreateWorkloadGroupResponse; import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupResponse; import org.opensearch.wlm.MutableWorkloadGroupFragment; +======== +import org.opensearch.plugin.wlm.querygroup.action.CreateQueryGroupResponse; +import org.opensearch.plugin.wlm.querygroup.action.DeleteQueryGroupRequest; +import org.opensearch.plugin.wlm.querygroup.action.UpdateQueryGroupRequest; +import org.opensearch.plugin.wlm.querygroup.action.UpdateQueryGroupResponse; +import org.opensearch.wlm.MutableQueryGroupFragment; +>>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/service/QueryGroupPersistenceService.java import org.opensearch.wlm.ResourceType; import java.util.Collection; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java new file mode 100644 index 0000000000000..d7ccb01e32e59 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.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; + +import org.opensearch.client.Client; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.plugin.wlm.rule.service.RulePersistenceService; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.Rule; +import org.opensearch.wlm.Rule.RuleAttribute; + +import java.util.Map; +import java.util.Set; + +import static org.opensearch.wlm.Rule.builder; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class RuleTestUtils { + public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; + public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; + public static final String LABEL_ONE = "label_one"; + public static final String LABEL_TWO = "label_two"; + public static final String PATTERN_ONE = "pattern_1"; + public static final String PATTERN_TWO = "pattern_2"; + public static final String QUERY_GROUP = "query_group"; + public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; + public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; + public static final Rule ruleOne = builder().feature(QUERY_GROUP) + .label(LABEL_ONE) + .attributeMap(Map.of(RuleAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE))) + .updatedAt(TIMESTAMP_ONE) + .build(); + + public static final Rule ruleTwo = builder().feature(QUERY_GROUP) + .label(LABEL_TWO) + .attributeMap(Map.of(RuleAttribute.INDEX_PATTERN, Set.of(PATTERN_TWO))) + .updatedAt(TIMESTAMP_TWO) + .build(); + + public static Map ruleMap() { + return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); + } + + public static RulePersistenceService setUpRulePersistenceService() { + Client client = mock(Client.class); + ThreadPool threadPool = mock(ThreadPool.class); + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + when(client.threadPool()).thenReturn(threadPool); + when(threadPool.getThreadContext()).thenReturn(threadContext); + return new RulePersistenceService(client); + } + + public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { + assertEquals(mapOne.size(), mapTwo.size()); + for (Map.Entry entry : mapOne.entrySet()) { + String id = entry.getKey(); + assertTrue(mapTwo.containsKey(id)); + Rule one = mapOne.get(id); + Rule two = mapTwo.get(id); + assertEqualRule(one, two, ruleUpdated); + } + } + + public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { + if (ruleUpdated) { + assertEquals(one.getFeature(), two.getFeature()); + assertEquals(one.getLabel(), two.getLabel()); + assertEquals(one.getAttributeMap(), two.getAttributeMap()); + } else { + assertEquals(one, two); + } + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequestTests.java similarity index 96% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequestTests.java index 31d3ea00b7bda..272f4db09109e 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequestTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.common.io.stream.BytesStreamOutput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponseTests.java similarity index 98% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponseTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponseTests.java index d25050341f997..033e1a67cc012 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponseTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.common.io.stream.BytesStreamOutput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequestTests.java similarity index 96% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequestTests.java index a7fa0939583c5..2df70252f9f55 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequestTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.common.io.stream.BytesStreamOutput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequestTests.java similarity index 97% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequestTests.java index 832761d5084bb..52e1880cf9a4a 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequestTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponseTests.java similarity index 99% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponseTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponseTests.java index dc0aeabc7a033..f9d77de883e73 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponseTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.common.io.stream.BytesStreamOutput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupActionTests.java similarity index 84% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupActionTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupActionTests.java index 7ffa33aa8a80a..00aa957371e74 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupActionTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupActionTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; @@ -14,7 +14,11 @@ import org.opensearch.cluster.metadata.IndexNameExpressionResolver; import org.opensearch.cluster.service.ClusterService; import org.opensearch.core.action.ActionListener; +<<<<<<<< HEAD:plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupActionTests.java import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; +======== +import org.opensearch.plugin.wlm.querygroup.service.QueryGroupPersistenceService; +>>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteQueryGroupActionTests.java import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupActionTests.java similarity index 97% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupActionTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupActionTests.java index cf12d9f6408cf..b3aa0dceee78e 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupActionTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupActionTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.ResourceNotFoundException; import org.opensearch.action.support.ActionFilters; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupRequestTests.java similarity index 98% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupRequestTests.java index e8d883da5c6eb..c982d5837d8eb 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupRequestTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponseTests.java similarity index 98% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponseTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponseTests.java index 97b9b9029373f..ff27c74becc5b 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponseTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.action; +package org.opensearch.plugin.wlm.querygroup.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.common.io.stream.BytesStreamOutput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/WorkloadGroupActionTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/WorkloadGroupActionTestUtils.java similarity index 100% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/WorkloadGroupActionTestUtils.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/WorkloadGroupActionTestUtils.java diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java similarity index 88% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupActionTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java index 8ce5c869f4481..4269c6d64ab97 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupActionTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java @@ -6,13 +6,18 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.rest; +package org.opensearch.plugin.wlm.querygroup.rest; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; import org.opensearch.common.CheckedConsumer; import org.opensearch.common.unit.TimeValue; +<<<<<<<< HEAD:plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupAction; import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupRequest; +======== +import org.opensearch.plugin.wlm.querygroup.action.DeleteQueryGroupAction; +import org.opensearch.plugin.wlm.querygroup.action.DeleteQueryGroupRequest; +>>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteQueryGroupActionTests.java import org.opensearch.rest.RestChannel; import org.opensearch.rest.RestHandler; import org.opensearch.rest.RestRequest; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/service/WorkloadGroupPersistenceServiceTests.java similarity index 100% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceServiceTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/service/WorkloadGroupPersistenceServiceTests.java diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java new file mode 100644 index 0000000000000..e9abb281df8b2 --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java @@ -0,0 +1,53 @@ +/* + * 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.rule.action; + +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.wlm.Rule; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; + +import static org.opensearch.plugin.wlm.RuleTestUtils.PATTERN_ONE; +import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; + +public class GetRuleRequestTests extends OpenSearchTestCase { + + /** + * Test case to verify the serialization and deserialization of GetRuleRequest + */ + public void testSerialization() throws IOException { + GetRuleRequest request = new GetRuleRequest(_ID_ONE, Map.of(Rule.RuleAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE))); + assertEquals(_ID_ONE, request.get_id()); + BytesStreamOutput out = new BytesStreamOutput(); + request.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + GetRuleRequest otherRequest = new GetRuleRequest(streamInput); + assertEquals(request.get_id(), otherRequest.get_id()); + assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters()); + } + + /** + * Test case to verify the serialization and deserialization of GetRuleRequest when name is null + */ + public void testSerializationWithNull() throws IOException { + GetRuleRequest request = new GetRuleRequest((String) null, new HashMap<>()); + assertNull(request.get_id()); + BytesStreamOutput out = new BytesStreamOutput(); + request.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + GetRuleRequest otherRequest = new GetRuleRequest(streamInput); + assertEquals(request.get_id(), otherRequest.get_id()); + assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters()); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java new file mode 100644 index 0000000000000..47545e554803e --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java @@ -0,0 +1,118 @@ +/* + * 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.rule.action; + +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.wlm.Rule; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; +import static org.opensearch.plugin.wlm.RuleTestUtils.assertEqualRules; +import static org.opensearch.plugin.wlm.RuleTestUtils.ruleMap; +import static org.opensearch.plugin.wlm.RuleTestUtils.ruleOne; +import static org.mockito.Mockito.mock; + +public class GetRuleResponseTests extends OpenSearchTestCase { + + /** + * Test case to verify the serialization and deserialization of GetRuleResponse + */ + public void testSerializationSingleRule() throws IOException { + Map map = new HashMap<>(); + map.put(_ID_ONE, ruleOne); + GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), RestStatus.OK); + assertEquals(response.getRules(), map); + + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + + GetRuleResponse otherResponse = new GetRuleResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); + assertEqualRules(response.getRules(), otherResponse.getRules(), false); + } + + /** + * Test case to verify the serialization and deserialization of GetRuleResponse when the result contains multiple rules + */ + public void testSerializationMultipleRule() throws IOException { + GetRuleResponse response = new GetRuleResponse(ruleMap(), RestStatus.OK); + assertEquals(response.getRules(), ruleMap()); + + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + + GetRuleResponse otherResponse = new GetRuleResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); + assertEquals(2, otherResponse.getRules().size()); + assertEqualRules(response.getRules(), otherResponse.getRules(), false); + } + + /** + * Test case to verify the serialization and deserialization of GetRuleResponse when the result is empty + */ + public void testSerializationNull() throws IOException { + Map map = new HashMap<>(); + GetRuleResponse response = new GetRuleResponse(map, RestStatus.OK); + assertEquals(response.getRules(), map); + + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + + GetRuleResponse otherResponse = new GetRuleResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); + assertEquals(0, otherResponse.getRules().size()); + } + + /** + * Test case to verify the toXContent of GetRuleResponse + */ + public void testToXContentGetSingleRule() throws IOException { + Map map = new HashMap<>(); + map.put(_ID_ONE, ruleOne); + GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), RestStatus.OK); + XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); + String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); + String expected = "{\n" + + " \"rules\" : [\n" + + " {\n" + + " \"_id\" : \"AgfUO5Ja9yfvhdONlYi3TQ==\",\n" + + " \"index_pattern\" : [\n" + + " \"pattern_1\"\n" + + " ],\n" + + " \"query_group\" : \"label_one\",\n" + + " \"updated_at\" : \"2024-01-26T08:58:57.558Z\"\n" + + " }\n" + + " ]\n" + + "}"; + assertEquals(expected, actual); + } + + /** + * Test case to verify toXContent of GetRuleResponse when the result contains zero Rule + */ + public void testToXContentGetZeroRule() throws IOException { + XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); + GetRuleResponse otherResponse = new GetRuleResponse(new HashMap<>(), RestStatus.OK); + String actual = otherResponse.toXContent(builder, mock(ToXContent.Params.class)).toString(); + String expected = "{\n" + " \"rules\" : [ ]\n" + "}"; + assertEquals(expected, actual); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleActionTests.java new file mode 100644 index 0000000000000..c09600be6fc6e --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleActionTests.java @@ -0,0 +1,35 @@ +/* + * 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.rule.rest; + +import org.opensearch.rest.RestHandler; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.List; + +import static org.opensearch.rest.RestRequest.Method.GET; + +public class RestGetRuleActionTests extends OpenSearchTestCase { + /** + * Test case to validate the construction for RestGetRuleAction + */ + public void testConstruction() { + RestGetRuleAction action = new RestGetRuleAction(); + assertNotNull(action); + assertEquals("get_rule", action.getName()); + List routes = action.routes(); + assertEquals(2, routes.size()); + RestHandler.Route route = routes.get(0); + assertEquals(GET, route.getMethod()); + assertEquals("_wlm/rule/", route.getPath()); + route = routes.get(1); + assertEquals(GET, route.getMethod()); + assertEquals("_wlm/rule/{_id}", route.getPath()); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java new file mode 100644 index 0000000000000..fa8225cf3552e --- /dev/null +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java @@ -0,0 +1,101 @@ +/* + * 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.rule.service; + +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.get.GetRequestBuilder; +import org.opensearch.action.get.GetResponse; +import org.opensearch.client.Client; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +import org.mockito.ArgumentCaptor; + +import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; +import static org.opensearch.plugin.wlm.RuleTestUtils.assertEqualRules; +import static org.opensearch.plugin.wlm.RuleTestUtils.ruleOne; +import static org.opensearch.plugin.wlm.RuleTestUtils.setUpRulePersistenceService; +import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.RULE_INDEX; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.clearInvocations; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@SuppressWarnings("unchecked") +public class RulePersistenceServiceTests extends OpenSearchTestCase { + + /** + * Test case to validate the logic for get Rule by id + */ + public void testGetRuleById() throws IOException { + String ruleSource = ruleOne.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString(); + ActionListener listener = mock(ActionListener.class); + RulePersistenceService rulePersistenceService = setUpRulePersistenceService(); + Client client = rulePersistenceService.getClient(); + GetRequestBuilder getRequestBuilder = mock(GetRequestBuilder.class); + GetResponse getResponse = mock(GetResponse.class); + + when(getResponse.isExists()).thenReturn(true); + when(getResponse.getSourceAsString()).thenReturn(ruleSource); + when(client.prepareGet(eq(RULE_INDEX), eq(_ID_ONE))).thenReturn(getRequestBuilder); + doAnswer(invocation -> { + ActionListener actionListener = invocation.getArgument(0); + actionListener.onResponse(getResponse); + return null; + }).when(getRequestBuilder).execute(any(ActionListener.class)); + + rulePersistenceService.getRule(_ID_ONE, new HashMap<>(), listener); + + ArgumentCaptor captor = ArgumentCaptor.forClass(GetRuleResponse.class); + verify(listener).onResponse(captor.capture()); + GetRuleResponse response = captor.getValue(); + assertNotNull(response); + assertEqualRules(Map.of(_ID_ONE, ruleOne), response.getRules(), false); + clearInvocations(client, getRequestBuilder, getResponse, listener); + } + + /** + * Test case to validate the get rule for an invalid id + */ + public void testGetRuleByIdNotFound() { + String nonExistentRuleId = "non-existent-rule"; + RulePersistenceService rulePersistenceService = setUpRulePersistenceService(); + Client client = rulePersistenceService.getClient(); + GetRequestBuilder getRequestBuilder = mock(GetRequestBuilder.class); + GetResponse getResponse = mock(GetResponse.class); + ActionListener listener = mock(ActionListener.class); + + when(client.prepareGet(RULE_INDEX, nonExistentRuleId)).thenReturn(getRequestBuilder); + when(getResponse.isExists()).thenReturn(false); + + doAnswer(invocation -> { + ActionListener actionListener = invocation.getArgument(0); + actionListener.onResponse(getResponse); + return null; + }).when(getRequestBuilder).execute(any(ActionListener.class)); + + rulePersistenceService.getRule(nonExistentRuleId, new HashMap<>(), listener); + + ArgumentCaptor captor = ArgumentCaptor.forClass(Exception.class); + verify(listener).onFailure(captor.capture()); + Exception exception = captor.getValue(); + assertTrue(exception instanceof ResourceNotFoundException); + clearInvocations(client, getRequestBuilder, getResponse, listener); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/Rule.java b/server/src/main/java/org/opensearch/wlm/Rule.java new file mode 100644 index 0000000000000..273abc455c7e3 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/Rule.java @@ -0,0 +1,359 @@ +/* + * 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.wlm; + +import org.opensearch.common.ValidationException; +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 org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.core.xcontent.XContentParseException; +import org.opensearch.core.xcontent.XContentParser; +import org.joda.time.Instant; + +import java.io.IOException; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Objects; +import java.util.Set; + +import static org.opensearch.cluster.metadata.QueryGroup.isValid; + +/** + * Represents a rule schema used for automatic query tagging in the system. + * This class encapsulates the criteria (defined through attributes) for automatically applying relevant + * tags to queries based on matching attribute patterns. This class provides an in-memory representation + * of a rule. The indexed view may differ in representation. + * { + * "_id": "fwehf8302582mglfio349==", + * "index_pattern": ["logs123", "user*"], + * "query_group": "dev_query_group_id", + * "updated_at": "01-10-2025T21:23:21.456Z" + * } + * @opensearch.experimental + */ +public class Rule implements Writeable, ToXContentObject { + private final Map> attributeMap; + private final Feature feature; + private final String label; + private final String updatedAt; + public static final String _ID_STRING = "_id"; + public static final String UPDATED_AT_STRING = "updated_at"; + public static final int MAX_NUMBER_OF_VALUES_PER_ATTRIBUTE = 10; + public static final int MAX_CHARACTER_LENGTH_PER_ATTRIBUTE_VALUE_STRING = 100; + + public Rule(Map> attributeMap, String label, String updatedAt, Feature feature) { + ValidationException validationException = new ValidationException(); + validateRuleInputs(attributeMap, label, updatedAt, feature, validationException); + if (!validationException.validationErrors().isEmpty()) { + throw new IllegalArgumentException(validationException); + } + + this.attributeMap = attributeMap; + this.feature = feature; + this.label = label; + this.updatedAt = updatedAt; + } + + public Rule(StreamInput in) throws IOException { + this( + in.readMap((i) -> RuleAttribute.fromName(i.readString()), i -> new HashSet<>(i.readStringList())), + in.readString(), + in.readString(), + Feature.fromName(in.readString()) + ); + } + + public static void requireNonNullOrEmpty(String value, String message, ValidationException validationException) { + if (value == null || value.isEmpty()) { + validationException.addValidationError(message); + } + } + + public static void validateRuleInputs( + Map> attributeMap, + String label, + String updatedAt, + Feature feature, + ValidationException validationException + ) { + if (feature == null) { + validationException.addValidationError("Couldn't identify which feature the rule belongs to. Rule feature name can't be null."); + } + requireNonNullOrEmpty(label, "Rule label can't be null or empty", validationException); + requireNonNullOrEmpty(updatedAt, "Rule update time can't be null or empty", validationException); + if (attributeMap == null || attributeMap.isEmpty()) { + validationException.addValidationError("Rule should have at least 1 attribute requirement"); + } + if (updatedAt != null && !isValid(Instant.parse(updatedAt).getMillis())) { + validationException.addValidationError("Rule update time is not a valid epoch"); + } + if (attributeMap != null && feature != null) { + validateAttributeMap(attributeMap, feature, validationException); + } + } + + public static void validateAttributeMap( + Map> attributeMap, + Feature feature, + ValidationException validationException + ) { + for (Map.Entry> entry : attributeMap.entrySet()) { + RuleAttribute ruleAttribute = entry.getKey(); + Set attributeValues = entry.getValue(); + if (!feature.isValidAttribute(ruleAttribute)) { + validationException.addValidationError( + ruleAttribute.getName() + " is not a valid attribute name under the feature: " + feature.getName() + ); + } + if (attributeValues.size() > MAX_NUMBER_OF_VALUES_PER_ATTRIBUTE) { + validationException.addValidationError( + "Each attribute can only have a maximum of 10 values. The input attribute " + ruleAttribute + " exceeds this limit." + ); + } + for (String attributeValue : attributeValues) { + if (attributeValue.isEmpty() || attributeValue.length() > MAX_CHARACTER_LENGTH_PER_ATTRIBUTE_VALUE_STRING) { + validationException.addValidationError( + "Attribute value [" + attributeValue + "] is invalid (empty or exceeds 100 characters)" + ); + } + } + } + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeMap(attributeMap, RuleAttribute::writeTo, StreamOutput::writeStringCollection); + out.writeString(label); + out.writeString(updatedAt); + out.writeString(feature.getName()); + } + + public static Rule fromXContent(final XContentParser parser) throws IOException { + return Builder.fromXContent(parser).build(); + } + + public String getLabel() { + return label; + } + + public String getUpdatedAt() { + return updatedAt; + } + + public Feature getFeature() { + return feature; + } + + public Map> getAttributeMap() { + return attributeMap; + } + + /** + * This enum enumerates the features that can use the Rule Based Auto-tagging + * @opensearch.experimental + */ + public enum Feature { + QUERY_GROUP("query_group", Set.of(RuleAttribute.INDEX_PATTERN)); + + private final String name; + private final Set allowedAttributes; + + Feature(String name, Set allowedAttributes) { + this.name = name; + this.allowedAttributes = allowedAttributes; + } + + public String getName() { + return name; + } + + public Set getAllowedAttributes() { + return allowedAttributes; + } + + public boolean isValidAttribute(RuleAttribute attribute) { + return allowedAttributes.contains(attribute); + } + + public static boolean isValidFeature(String s) { + return Arrays.stream(values()).anyMatch(feature -> feature.getName().equalsIgnoreCase(s)); + } + + public static Feature fromName(String s) { + return Arrays.stream(values()) + .filter(feature -> feature.getName().equalsIgnoreCase(s)) + .findFirst() + .orElseThrow(() -> new IllegalArgumentException("Invalid value for Feature: " + s)); + } + } + + /** + * This RuleAttribute enum contains the attribute names for a rule. + * @opensearch.experimental + */ + public enum RuleAttribute { + INDEX_PATTERN("index_pattern"); + + private final String name; + + RuleAttribute(String name) { + this.name = name; + } + + public String getName() { + return name; + } + + public static void writeTo(StreamOutput out, RuleAttribute ruleAttribute) throws IOException { + out.writeString(ruleAttribute.getName()); + } + + public static RuleAttribute fromName(String s) { + for (RuleAttribute attribute : values()) { + if (attribute.getName().equalsIgnoreCase(s)) return attribute; + + } + throw new IllegalArgumentException("Invalid value for RuleAttribute: " + s); + } + } + + @Override + public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { + builder.startObject(); + String id = params.param(_ID_STRING); + if (id != null) { + builder.field(_ID_STRING, id); + } + for (Map.Entry> entry : attributeMap.entrySet()) { + builder.array(entry.getKey().getName(), entry.getValue().toArray(new String[0])); + } + builder.field(feature.getName(), label); + builder.field(UPDATED_AT_STRING, updatedAt); + builder.endObject(); + return builder; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + Rule that = (Rule) o; + return Objects.equals(label, that.label) + && Objects.equals(feature, that.feature) + && Objects.equals(attributeMap, that.attributeMap) + && Objects.equals(updatedAt, that.updatedAt); + } + + @Override + public int hashCode() { + return Objects.hash(label, feature, attributeMap, updatedAt); + } + + /** + * builder method for the {@link Rule} + * @return Builder object + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Builder class for {@link Rule} + * @opensearch.experimental + */ + public static class Builder { + private Map> attributeMap; + private Feature feature; + private String label; + private String updatedAt; + + private Builder() {} + + public static Builder fromXContent(XContentParser parser) throws IOException { + if (parser.currentToken() == null) { + parser.nextToken(); + } + Builder builder = builder(); + XContentParser.Token token = parser.currentToken(); + + if (token != XContentParser.Token.START_OBJECT) { + throw new IllegalArgumentException("Expected START_OBJECT token but found [" + parser.currentName() + "]"); + } + Map> attributeMap1 = new HashMap<>(); + String fieldName = ""; + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + fieldName = parser.currentName(); + } else if (token.isValue()) { + if (Feature.isValidFeature(fieldName)) { + builder.feature(fieldName); + builder.label(parser.text()); + } else if (fieldName.equals(UPDATED_AT_STRING)) { + builder.updatedAt(parser.text()); + } else { + throw new IllegalArgumentException(fieldName + " is not a valid field in Rule"); + } + } else if (token == XContentParser.Token.START_ARRAY) { + fromXContentParseArray(parser, fieldName, attributeMap1); + } + } + return builder.attributeMap(attributeMap1); + } + + public static void fromXContentParseArray(XContentParser parser, String fieldName, Map> attributeMap) + throws IOException { + RuleAttribute ruleAttribute = RuleAttribute.fromName(fieldName); + Set attributeValueSet = new HashSet<>(); + while (parser.nextToken() != XContentParser.Token.END_ARRAY) { + if (parser.currentToken() == XContentParser.Token.VALUE_STRING) { + attributeValueSet.add(parser.text()); + } else { + throw new XContentParseException("Unexpected token in array: " + parser.currentToken()); + } + } + attributeMap.put(ruleAttribute, attributeValueSet); + } + + public Builder label(String label) { + this.label = label; + return this; + } + + public Builder attributeMap(Map> attributeMap) { + this.attributeMap = attributeMap; + return this; + } + + public Builder feature(String feature) { + this.feature = Feature.fromName(feature); + return this; + } + + public Builder updatedAt(String updatedAt) { + this.updatedAt = updatedAt; + return this; + } + + public Rule build() { + return new Rule(attributeMap, label, updatedAt, feature); + } + + public String getLabel() { + return label; + } + + public Map> getAttributeMap() { + return attributeMap; + } + } +} diff --git a/server/src/test/java/org/opensearch/wlm/RuleTests.java b/server/src/test/java/org/opensearch/wlm/RuleTests.java new file mode 100644 index 0000000000000..72a4abf1f3877 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/RuleTests.java @@ -0,0 +1,163 @@ +/* + * 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.wlm; + +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.test.AbstractSerializingTestCase; +import org.opensearch.wlm.Rule.Feature; +import org.opensearch.wlm.Rule.RuleAttribute; +import org.joda.time.Instant; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.opensearch.wlm.Rule.MAX_CHARACTER_LENGTH_PER_ATTRIBUTE_VALUE_STRING; +import static org.opensearch.wlm.Rule.MAX_NUMBER_OF_VALUES_PER_ATTRIBUTE; +import static org.opensearch.wlm.Rule._ID_STRING; + +public class RuleTests extends AbstractSerializingTestCase { + public static final String _ID = "AgfUfjw039vhdONlYi3TQ=="; + public static final String LABEL = "label"; + + static Rule createRandomRule(String label) { + Feature feature = randomFeature(); + return Rule.builder() + .label(label) + .feature(feature.getName()) + .attributeMap(randomAttributeMaps(feature)) + .updatedAt(Instant.now().toString()) + .build(); + } + + private static Feature randomFeature() { + return Feature.values()[randomIntBetween(0, Feature.values().length - 1)]; + } + + private static Map> randomAttributeMaps(Feature feature) { + Map> attributeMap = new HashMap<>(); + if (feature == null) { + return attributeMap; + } + List allowedAttributes = new ArrayList<>(feature.getAllowedAttributes()); + do { + attributeMap.clear(); + for (RuleAttribute currAttribute : allowedAttributes) { + if (randomBoolean()) { + attributeMap.put(currAttribute, randomAttributeValues()); + } + } + } while (attributeMap.isEmpty()); + return attributeMap; + } + + private static Set randomAttributeValues() { + Set res = new HashSet<>(); + int numberOfValues = randomIntBetween(1, MAX_NUMBER_OF_VALUES_PER_ATTRIBUTE); + for (int i = 0; i < numberOfValues; i++) { + res.add(randomAlphaOfLength(randomIntBetween(1, MAX_CHARACTER_LENGTH_PER_ATTRIBUTE_VALUE_STRING))); + } + return res; + } + + @Override + protected Rule doParseInstance(XContentParser parser) throws IOException { + return Rule.fromXContent(parser); + } + + @Override + protected Writeable.Reader instanceReader() { + return Rule::new; + } + + @Override + protected Rule createTestInstance() { + return createRandomRule(LABEL); + } + + static Rule buildRule(String label, String feature, Map> attributeListMap, String updatedAt) { + return Rule.builder().label(label).feature(feature).attributeMap(attributeListMap).updatedAt(updatedAt).build(); + } + + public void testInvalidFeature() { + assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, null, randomAttributeMaps(null), Instant.now().toString())); + assertThrows( + IllegalArgumentException.class, + () -> buildRule(LABEL, "invalid", randomAttributeMaps(null), Instant.now().toString()) + ); + } + + public void testInvalidLabel() { + assertThrows(IllegalArgumentException.class, () -> createRandomRule(null)); + assertThrows(IllegalArgumentException.class, () -> createRandomRule("")); + } + + public void testInvalidUpdateTime() { + Feature feature = randomFeature(); + assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, feature.toString(), randomAttributeMaps(feature), null)); + } + + public void testNullOrEmptyAttributeMap() { + Feature feature = randomFeature(); + assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, feature.toString(), new HashMap<>(), Instant.now().toString())); + assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, feature.toString(), null, Instant.now().toString())); + } + + public void testInvalidAttributeMap() { + Map> map = new HashMap<>(); + map.put(RuleAttribute.INDEX_PATTERN, Set.of("")); + assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, randomFeature().getName(), map, Instant.now().toString())); + + map.put(RuleAttribute.INDEX_PATTERN, Set.of(randomAlphaOfLength(MAX_CHARACTER_LENGTH_PER_ATTRIBUTE_VALUE_STRING + 1))); + assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, randomFeature().getName(), map, Instant.now().toString())); + + map.put(RuleAttribute.INDEX_PATTERN, new HashSet<>()); + for (int i = 0; i < MAX_NUMBER_OF_VALUES_PER_ATTRIBUTE + 1; i++) { + map.get(RuleAttribute.INDEX_PATTERN).add(String.valueOf(i)); + } + assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, randomFeature().getName(), map, Instant.now().toString())); + } + + public void testValidRule() { + Map> map = Map.of(RuleAttribute.INDEX_PATTERN, Set.of("index*", "log*")); + String updatedAt = Instant.now().toString(); + Rule rule = buildRule(LABEL, Feature.QUERY_GROUP.getName(), map, updatedAt); + assertNotNull(rule.getLabel()); + assertEquals(LABEL, rule.getLabel()); + assertNotNull(updatedAt); + assertEquals(updatedAt, rule.getUpdatedAt()); + Map> resultMap = rule.getAttributeMap(); + assertNotNull(resultMap); + assertFalse(resultMap.isEmpty()); + assertNotNull(rule.getFeature()); + assertEquals(Feature.QUERY_GROUP, rule.getFeature()); + } + + public void testToXContent() throws IOException { + Map> map = Map.of(RuleAttribute.INDEX_PATTERN, Set.of("log*")); + String updatedAt = Instant.now().toString(); + Rule rule = buildRule(LABEL, Feature.QUERY_GROUP.getName(), map, updatedAt); + + XContentBuilder builder = JsonXContent.contentBuilder(); + rule.toXContent(builder, new ToXContent.MapParams(Map.of(_ID_STRING, _ID))); + + assertEquals( + "{\"_id\":\"" + _ID + "\",\"index_pattern\":[\"log*\"],\"query_group\":\"label\",\"updated_at\":\"" + updatedAt + "\"}", + builder.toString() + ); + } +} From 4be74e40f6ec38e365cabbcd1b2a0c74afc70c57 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 18 Feb 2025 16:58:03 -0800 Subject: [PATCH 02/20] modify based on comments Signed-off-by: Ruirui Zhang --- .../plugin/wlm/WorkloadManagementPlugin.java | 2 +- .../plugin/wlm/rule/QueryGroupAttribute.java | 51 +++ .../wlm/rule/QueryGroupFeatureType.java | 58 +++ .../plugin/wlm/rule/action/GetRuleAction.java | 11 +- .../wlm/rule/action/GetRuleRequest.java | 32 +- .../wlm/rule/action/GetRuleResponse.java | 47 ++- .../rule/action/TransportGetRuleAction.java | 2 +- .../wlm/rule/rest/RestGetRuleAction.java | 18 +- .../rule/service/RulePersistenceService.java | 174 +++++---- .../opensearch/plugin/wlm/RuleTestUtils.java | 47 ++- .../wlm/rule/action/GetRuleRequestTests.java | 7 +- .../wlm/rule/action/GetRuleResponseTests.java | 24 +- .../service/RulePersistenceServiceTests.java | 53 ++- .../main/java/org/opensearch/wlm/Rule.java | 359 ------------------ .../java/org/opensearch/wlm/RuleTests.java | 163 -------- 15 files changed, 377 insertions(+), 671 deletions(-) create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java delete mode 100644 server/src/main/java/org/opensearch/wlm/Rule.java delete mode 100644 server/src/test/java/org/opensearch/wlm/RuleTests.java diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index 21daded9e3069..a82a0404eeb75 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -63,7 +63,7 @@ import java.util.List; import java.util.function.Supplier; -import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.RULE_INDEX; +import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.RULES_INDEX; /** * Plugin class for WorkloadManagement diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java new file mode 100644 index 0000000000000..906b68e37efff --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.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.rule; + +import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.AutoTaggingRegistry; + +/** + * Attributes specific to the query group feature. + * @opensearch.experimental + */ +public enum QueryGroupAttribute implements Attribute { + INDEX_PATTERN("index_pattern"); + + private final String name; + + QueryGroupAttribute(String name) { + this.name = name; + } + + static { + for (QueryGroupAttribute attr : QueryGroupAttribute.values()) { + attr.registerAttribute(); + } + } + + @Override + public String getName() { + return name; + } + + @Override + public void registerAttribute() { + AutoTaggingRegistry.registerAttribute(this); + } + + public static QueryGroupAttribute fromName(String name) { + for (QueryGroupAttribute attr : QueryGroupAttribute.values()) { + if (attr.getName().equals(name)) { + return attr; + } + } + throw new IllegalArgumentException("Unknown QueryGroupAttribute: " + name); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java new file mode 100644 index 0000000000000..18e227d23e888 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java @@ -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.rule; + +import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.AutoTaggingRegistry; +import org.opensearch.autotagging.FeatureType; + +import java.util.Set; + +/** + * Represents the feature type for "query_group" in the OpenSearch workload management plugin. + * @opensearch.experimental + */ +public class QueryGroupFeatureType implements FeatureType { + public static final QueryGroupFeatureType INSTANCE = new QueryGroupFeatureType(); + public static final String NAME = "query_group"; + private static final int MAX_ATTRIBUTE_VALUES = 10; + private static final int MAX_ATTRIBUTE_VALUE_LENGTH = 100; + private static final Set ALLOWED_ATTRIBUTES = Set.of(QueryGroupAttribute.values()); + + private QueryGroupFeatureType() {} + + static { + INSTANCE.registerFeatureType(); + } + + @Override + public String getName() { + return NAME; + } + + @Override + public int getMaxNumberOfValuesPerAttribute() { + return MAX_ATTRIBUTE_VALUES; + } + + @Override + public int getMaxCharLengthPerAttributeValue() { + return MAX_ATTRIBUTE_VALUE_LENGTH; + } + + @Override + public Set getAllowedAttributes() { + return ALLOWED_ATTRIBUTES; + } + + @Override + public void registerFeatureType() { + AutoTaggingRegistry.registerFeatureType(INSTANCE); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java index f334b064acc62..2ee030576c18e 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java @@ -11,24 +11,15 @@ import org.opensearch.action.ActionType; /** - * Transport action to get Rule + * Action type for getting Rules in workload management * @opensearch.experimental */ public class GetRuleAction extends ActionType { - /** - * An instance of GetRuleAction - */ public static final GetRuleAction INSTANCE = new GetRuleAction(); - /** - * Name for GetRuleAction - */ public static final String NAME = "cluster:admin/opensearch/wlm/rule/_get"; - /** - * Default constructor - */ private GetRuleAction() { super(NAME, GetRuleResponse::new); } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java index e821007b1d13b..0f9b4cef43925 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java @@ -8,11 +8,11 @@ package org.opensearch.plugin.wlm.rule.action; +import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionRequestValidationException; -import org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest; +import org.opensearch.autotagging.Attribute; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; -import org.opensearch.wlm.Rule.RuleAttribute; import java.io.IOException; import java.util.HashSet; @@ -21,20 +21,27 @@ /** * A request for get Rule + * Example Request: + * curl -X GET "localhost:9200/_wlm/rule" - get all rules + * curl -X GET "localhost:9200/_wlm/rule/{_id}" - get single rule by id + * curl -X GET "localhost:9200/_wlm/rule?index_pattern=a,b" - get all rules containing index_pattern as a or b * @opensearch.experimental */ -public class GetRuleRequest extends ClusterManagerNodeRequest { +public class GetRuleRequest extends ActionRequest { private final String _id; - private final Map> attributeFilters; + private final Map> attributeFilters; + private final String searchAfter; /** * Constructor for GetRuleRequest * @param _id - Rule _id that we want to get * @param attributeFilters - Attributes that we want to filter on + * @param searchAfter - The sort values from the last document of the previous page, used for pagination */ - public GetRuleRequest(String _id, Map> attributeFilters) { + public GetRuleRequest(String _id, Map> attributeFilters, String searchAfter) { this._id = _id; this.attributeFilters = attributeFilters; + this.searchAfter = searchAfter; } /** @@ -44,7 +51,8 @@ public GetRuleRequest(String _id, Map> attributeFilte public GetRuleRequest(StreamInput in) throws IOException { super(in); _id = in.readOptionalString(); - attributeFilters = in.readMap((i) -> RuleAttribute.fromName(i.readString()), i -> new HashSet<>(i.readStringList())); + attributeFilters = in.readMap(Attribute::from, i -> new HashSet<>(i.readStringList())); + searchAfter = in.readOptionalString(); } @Override @@ -56,7 +64,8 @@ public ActionRequestValidationException validate() { public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeOptionalString(_id); - out.writeMap(attributeFilters, RuleAttribute::writeTo, StreamOutput::writeStringCollection); + out.writeMap(attributeFilters, (output, attribute) -> attribute.writeTo(output), StreamOutput::writeStringCollection); + out.writeOptionalString(searchAfter); } /** @@ -69,7 +78,14 @@ public String get_id() { /** * attributeFilters getter */ - public Map> getAttributeFilters() { + public Map> getAttributeFilters() { return attributeFilters; } + + /** + * searchAfter getter + */ + public String getSearchAfter() { + return searchAfter; + } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java index b8be558c7e567..2b4ffd6b4f431 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java @@ -8,6 +8,7 @@ package org.opensearch.plugin.wlm.rule.action; +import org.opensearch.autotagging.Rule; import org.opensearch.core.action.ActionResponse; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; @@ -15,28 +16,46 @@ import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.wlm.Rule; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import java.io.IOException; import java.util.Map; -import static org.opensearch.wlm.Rule._ID_STRING; +import static org.opensearch.autotagging.Rule._ID_STRING; +import static org.opensearch.plugin.wlm.rule.rest.RestGetRuleAction.SEARCH_AFTER_STRING; /** - * Response for the get API for Rule + * Response for the get API for Rule. + * Example response: + * { + * "rules": [ + * { + * "_id": "z1MJApUB0zgMcDmz-UQq", + * "description": "Rule for tagging query_group_id to index123" + * "index_pattern": ["index123"], + * "query_group": "query_group_id", + * "updated_at": "2025-02-14T01:19:22.589Z" + * }, + * ... + * ], + * "search_after": ["z1MJApUB0zgMcDmz-UQq"] + * } * @opensearch.experimental */ public class GetRuleResponse extends ActionResponse implements ToXContent, ToXContentObject { - private final Map rules; + private final Map> rules; + private final String searchAfter; private final RestStatus restStatus; /** * Constructor for GetRuleResponse * @param rules - The Map of Rules to be included in the response + * @param searchAfter - The searchAfter field for the response * @param restStatus - The restStatus for the response */ - public GetRuleResponse(final Map rules, RestStatus restStatus) { + public GetRuleResponse(final Map> rules, String searchAfter, RestStatus restStatus) { this.rules = rules; + this.searchAfter = searchAfter; this.restStatus = restStatus; } @@ -45,13 +64,13 @@ public GetRuleResponse(final Map rules, RestStatus restStatus) { * @param in - A {@link StreamInput} object */ public GetRuleResponse(StreamInput in) throws IOException { - this.rules = in.readMap(StreamInput::readString, Rule::new); - this.restStatus = RestStatus.readFrom(in); + this(in.readMap(StreamInput::readString, Rule::new), in.readOptionalString(), RestStatus.readFrom(in)); } @Override public void writeTo(StreamOutput out) throws IOException { out.writeMap(rules, StreamOutput::writeString, (outStream, rule) -> rule.writeTo(outStream)); + out.writeOptionalString(searchAfter); RestStatus.writeTo(out, restStatus); } @@ -59,10 +78,13 @@ public void writeTo(StreamOutput out) throws IOException { public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); builder.startArray("rules"); - for (Map.Entry entry : rules.entrySet()) { + for (Map.Entry> entry : rules.entrySet()) { entry.getValue().toXContent(builder, new MapParams(Map.of(_ID_STRING, entry.getKey()))); } builder.endArray(); + if (searchAfter != null && !searchAfter.isEmpty()) { + builder.field(SEARCH_AFTER_STRING, new Object[] { searchAfter }); + } builder.endObject(); return builder; } @@ -70,7 +92,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws /** * rules getter */ - public Map getRules() { + public Map> getRules() { return rules; } @@ -80,4 +102,11 @@ public Map getRules() { public RestStatus getRestStatus() { return restStatus; } + + /** + * searchAfter getter + */ + public String getSearchAfter() { + return searchAfter; + } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java index 66b243ff30e64..a7c2e4cac2e04 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java @@ -42,6 +42,6 @@ public TransportGetRuleAction( @Override protected void doExecute(Task task, GetRuleRequest request, ActionListener listener) { - rulePersistenceService.getRule(request.get_id(), request.getAttributeFilters(), listener); + rulePersistenceService.getRule(request.get_id(), request.getAttributeFilters(), request.getSearchAfter(), listener); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java index 685fd01f200d5..01f734e12d91e 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java @@ -8,9 +8,11 @@ package org.opensearch.plugin.wlm.rule.rest; +import org.opensearch.autotagging.Attribute; import org.opensearch.client.node.NodeClient; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.plugin.wlm.rule.QueryGroupAttribute; import org.opensearch.plugin.wlm.rule.action.GetRuleAction; import org.opensearch.plugin.wlm.rule.action.GetRuleRequest; import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; @@ -20,7 +22,6 @@ import org.opensearch.rest.RestRequest; import org.opensearch.rest.RestResponse; import org.opensearch.rest.action.RestResponseListener; -import org.opensearch.wlm.Rule.RuleAttribute; import java.io.IOException; import java.util.Arrays; @@ -30,14 +31,15 @@ import java.util.Map; import java.util.Set; +import static org.opensearch.autotagging.Rule._ID_STRING; import static org.opensearch.rest.RestRequest.Method.GET; -import static org.opensearch.wlm.Rule._ID_STRING; /** * Rest action to get a Rule * @opensearch.experimental */ public class RestGetRuleAction extends BaseRestHandler { + public static final String SEARCH_AFTER_STRING = "search_after"; /** * Constructor for RestGetRuleAction @@ -59,15 +61,19 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - final Map> attributeFilters = new HashMap<>(); + final Map> attributeFilters = new HashMap<>(); for (String attributeName : request.params().keySet()) { - if (attributeName.equals(_ID_STRING)) { + if (attributeName.equals(_ID_STRING) || attributeName.equals(SEARCH_AFTER_STRING)) { continue; } String[] valuesArray = request.param(attributeName).split(","); - attributeFilters.put(RuleAttribute.fromName(attributeName), new HashSet<>(Arrays.asList(valuesArray))); + attributeFilters.put(QueryGroupAttribute.fromName(attributeName), new HashSet<>(Arrays.asList(valuesArray))); } - final GetRuleRequest getRuleRequest = new GetRuleRequest(request.param(_ID_STRING), attributeFilters); + final GetRuleRequest getRuleRequest = new GetRuleRequest( + request.param(_ID_STRING), + attributeFilters, + request.param(SEARCH_AFTER_STRING) + ); return channel -> client.execute(GetRuleAction.INSTANCE, getRuleRequest, getRuleResponse(channel)); } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java index 4d00cb0decd27..aa45dda3ba210 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java @@ -12,8 +12,13 @@ import org.apache.logging.log4j.Logger; import org.opensearch.ResourceNotFoundException; import org.opensearch.action.get.GetResponse; +import org.opensearch.action.search.SearchRequestBuilder; import org.opensearch.action.search.SearchResponse; +import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.Rule; +import org.opensearch.autotagging.Rule.Builder; import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; @@ -22,35 +27,41 @@ import org.opensearch.core.xcontent.MediaTypeRegistry; import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilders; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; -import org.opensearch.wlm.Rule; -import org.opensearch.wlm.Rule.Builder; -import org.opensearch.wlm.Rule.RuleAttribute; +import org.opensearch.search.SearchHit; +import org.opensearch.search.sort.SortOrder; import java.io.IOException; import java.util.Arrays; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; import java.util.stream.Collectors; +import static org.opensearch.autotagging.Rule._ID_STRING; + /** - * This class defines the functions for Rule persistence + * This class encapsulates the logic to manage the lifecycle of rules at index level * @opensearch.experimental */ public class RulePersistenceService { - public static final String RULE_INDEX = ".rule"; + public static final String RULES_INDEX = ".rules"; private final Client client; + private final ClusterService clusterService; private static final Logger logger = LogManager.getLogger(RulePersistenceService.class); - private static final int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 50; + public static final int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 100; /** * Constructor for RulePersistenceService * @param client {@link Client} - The client to be used by RulePersistenceService */ @Inject - public RulePersistenceService(final Client client) { + public RulePersistenceService(final ClusterService clusterService, final Client client) { + this.clusterService = clusterService; this.client = client; } @@ -58,37 +69,39 @@ public RulePersistenceService(final Client client) { * Entry point for the get rule api logic in persistence service. * @param id - The id of the rule to get. Get all matching rules when id is null * @param attributeFilters - A map containing the attributes that user want to filter on + * @param searchAfter - The sort values from the last document of the previous page, used for pagination * @param listener - ActionListener for GetRuleResponse */ - public void getRule(String id, Map> attributeFilters, ActionListener listener) { + public void getRule( + String id, + Map> attributeFilters, + String searchAfter, + ActionListener listener + ) { if (id != null) { fetchRuleById(id, listener); } else { - fetchAllRules(attributeFilters, listener); + fetchAllRules(attributeFilters, searchAfter, listener); } } /** * Fetch a single rule from system index using id - * @param id - The id of the rule to get. + * @param id - The id of the rule to get * @param listener - ActionListener for GetRuleResponse */ private void fetchRuleById(String id, ActionListener listener) { ThreadContext.StoredContext storedContext = client.threadPool().getThreadContext().stashContext(); - client.prepareGet(RULE_INDEX, id) - .execute(ActionListener.wrap( - getResponse -> { - try (ThreadContext.StoredContext context = storedContext) { - handleGetOneRuleResponse(id, getResponse, listener); - } - }, - e -> { - try (ThreadContext.StoredContext context = storedContext) { - logger.error("Failed to fetch rule with ID {}: {}", id, e.getMessage()); - listener.onFailure(e); - } - } - )); + client.prepareGet(RULES_INDEX, id).execute(ActionListener.wrap(getResponse -> { + try (ThreadContext.StoredContext context = storedContext) { + handleGetOneRuleResponse(id, getResponse, listener); + } + }, e -> { + try (ThreadContext.StoredContext context = storedContext) { + logger.error("Failed to fetch rule with ID {}: {}", id, e.getMessage()); + listener.onFailure(e); + } + })); } /** @@ -106,7 +119,13 @@ private void handleGetOneRuleResponse(String id, GetResponse getResponse, Action DeprecationHandler.THROW_UNSUPPORTED_OPERATION, getResponse.getSourceAsString() ); - listener.onResponse(new GetRuleResponse(Map.of(id, Builder.fromXContent(parser).build()), RestStatus.OK)); + listener.onResponse( + new GetRuleResponse( + Map.of(id, Builder.fromXContent(parser, QueryGroupFeatureType.INSTANCE).build()), + null, + RestStatus.OK + ) + ); } catch (IOException e) { logger.error("Error parsing rule with ID {}: {}", id, e.getMessage()); listener.onFailure(e); @@ -117,74 +136,79 @@ private void handleGetOneRuleResponse(String id, GetResponse getResponse, Action } /** - * Fetch all rule from system index based on attributeFilters + * Fetch all rule from system index based on attributeFilters. * @param attributeFilters - A map containing the attributes that user want to filter on + * @param searchAfter - The sort values from the last document of the previous page, used for pagination * @param listener - ActionListener for GetRuleResponse */ - private void fetchAllRules(Map> attributeFilters, ActionListener listener) { + private void fetchAllRules(Map> attributeFilters, String searchAfter, ActionListener listener) { ThreadContext.StoredContext storedContext = client.threadPool().getThreadContext().stashContext(); - client.prepareSearch(RULE_INDEX) - .setQuery(QueryBuilders.matchAllQuery()) - .setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST) - .execute(ActionListener.wrap( - searchResponse -> { - try (ThreadContext.StoredContext context = storedContext) { - handleGetAllRuleResponse(searchResponse, attributeFilters, listener); - } - }, - e -> { - try (ThreadContext.StoredContext context = storedContext) { - logger.error("Failed to fetch all rules: {}", e.getMessage()); - listener.onFailure(e); - } + SearchRequestBuilder searchRequest = buildGetAllRuleSearchRequest(attributeFilters, searchAfter); + searchRequest.execute(ActionListener.wrap(searchResponse -> { + try (ThreadContext.StoredContext context = storedContext) { + listener.onResponse(handleGetAllRuleResponse(searchResponse)); + } + }, e -> { + try (ThreadContext.StoredContext context = storedContext) { + logger.error("Failed to fetch all rules: {}", e.getMessage()); + listener.onFailure(e); + } + })); + } + + /** + * Builds a search request to retrieve all rules from the rules index, applying attribute-based filters + * and ensuring that the rules are associated with the query group feature type. + * @param attributeFilters A map of attributes to their associated set of values used to filter the rules. + * @param searchAfter A cursor to enable pagination, used to fetch results after a specific document. + */ + SearchRequestBuilder buildGetAllRuleSearchRequest(Map> attributeFilters, String searchAfter) { + BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); + for (Map.Entry> entry : attributeFilters.entrySet()) { + Attribute attribute = entry.getKey(); + Set values = entry.getValue(); + if (values != null && !values.isEmpty()) { + BoolQueryBuilder attributeQuery = QueryBuilders.boolQuery(); + for (String value : values) { + attributeQuery.should(QueryBuilders.matchQuery(attribute.getName(), value)); } - )); + boolQuery.must(attributeQuery); + } + } + boolQuery.filter(QueryBuilders.existsQuery(QueryGroupFeatureType.NAME)); + SearchRequestBuilder searchRequest = client.prepareSearch(RULES_INDEX) + .setQuery(boolQuery) + .setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST) + .addSort(_ID_STRING, SortOrder.ASC); + if (searchAfter != null) { + searchRequest.searchAfter(new Object[] { searchAfter }); + } + return searchRequest; } /** * Process searchResponse from index and send a GetRuleResponse * @param searchResponse - Response received from index - * @param attributeFilters - A map containing the attributes that user want to filter on - * @param listener - ActionListener for GetRuleResponse */ - private void handleGetAllRuleResponse( - SearchResponse searchResponse, - Map> attributeFilters, - ActionListener listener - ) { - Map ruleMap = Arrays.stream(searchResponse.getHits().getHits()).map(hit -> { + GetRuleResponse handleGetAllRuleResponse(SearchResponse searchResponse) { + List hits = Arrays.asList(searchResponse.getHits().getHits()); + Map> ruleMap = hits.stream().map(hit -> { + String hitId = hit.getId(); try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { XContentParser parser = MediaTypeRegistry.JSON.xContent() .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, hit.getSourceAsString()); - Rule currRule = Rule.Builder.fromXContent(parser).build(); - if (matchesFilters(currRule, attributeFilters)) { - return Map.entry(hit.getId(), currRule); - } - return null; + return Map.entry(hitId, Builder.fromXContent(parser, QueryGroupFeatureType.INSTANCE).build()); } catch (IOException e) { - logger.error("Failed to parse rule from hit: {}", e.getMessage()); - listener.onFailure(e); + logger.info( + "Issue met when parsing rule from hit, the feature type for rule id {} is probably not query_group: {}", + hitId, + e.getMessage() + ); return null; } }).filter(Objects::nonNull).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - listener.onResponse(new GetRuleResponse(ruleMap, RestStatus.OK)); - } - - /** - * Returns true if the rule matches the attributeFilters and should be included in the response - * @param rule - the rule to be checked against the attribute filters - * @param attributeFilters - A map containing the attributes that user want to filter on - */ - private boolean matchesFilters(Rule rule, Map> attributeFilters) { - for (Map.Entry> entry : attributeFilters.entrySet()) { - RuleAttribute attribute = entry.getKey(); - Set expectedValues = entry.getValue(); - Set ruleValues = rule.getAttributeMap().get(attribute); - if (ruleValues == null || ruleValues.stream().noneMatch(expectedValues::contains)) { - return false; - } - } - return true; + String nextSearchAfter = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); + return new GetRuleResponse(ruleMap, nextSearchAfter, RestStatus.OK); } /** diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java index d7ccb01e32e59..387cdb1bae6e7 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java @@ -8,18 +8,20 @@ package org.opensearch.plugin.wlm; +import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.Rule; import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.plugin.wlm.rule.QueryGroupAttribute; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import org.opensearch.plugin.wlm.rule.service.RulePersistenceService; import org.opensearch.threadpool.ThreadPool; -import org.opensearch.wlm.Rule; -import org.opensearch.wlm.Rule.RuleAttribute; import java.util.Map; import java.util.Set; -import static org.opensearch.wlm.Rule.builder; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; @@ -32,48 +34,61 @@ public class RuleTestUtils { public static final String LABEL_TWO = "label_two"; public static final String PATTERN_ONE = "pattern_1"; public static final String PATTERN_TWO = "pattern_2"; - public static final String QUERY_GROUP = "query_group"; + public static final String DESCRIPTION_ONE = "description_1"; + public static final String DESCRIPTION_TWO = "description_2"; public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; - public static final Rule ruleOne = builder().feature(QUERY_GROUP) + public static final String SEARCH_AFTER = "search_after_id"; + public static final Map> attributes = Map.of(QueryGroupAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE)); + public static final Rule ruleOne = Rule.builder() + .description(DESCRIPTION_ONE) + .featureType(QueryGroupFeatureType.INSTANCE) .label(LABEL_ONE) - .attributeMap(Map.of(RuleAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE))) + .attributeMap(Map.of(QueryGroupAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE))) .updatedAt(TIMESTAMP_ONE) .build(); - public static final Rule ruleTwo = builder().feature(QUERY_GROUP) + public static final Rule ruleTwo = Rule.builder() + .description(DESCRIPTION_TWO) + .featureType(QueryGroupFeatureType.INSTANCE) .label(LABEL_TWO) - .attributeMap(Map.of(RuleAttribute.INDEX_PATTERN, Set.of(PATTERN_TWO))) + .attributeMap(Map.of(QueryGroupAttribute.INDEX_PATTERN, Set.of(PATTERN_TWO))) .updatedAt(TIMESTAMP_TWO) .build(); - public static Map ruleMap() { + public static Map> ruleMap() { return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); } public static RulePersistenceService setUpRulePersistenceService() { Client client = mock(Client.class); + ClusterService clusterService = mock(ClusterService.class); ThreadPool threadPool = mock(ThreadPool.class); ThreadContext threadContext = new ThreadContext(Settings.EMPTY); when(client.threadPool()).thenReturn(threadPool); when(threadPool.getThreadContext()).thenReturn(threadContext); - return new RulePersistenceService(client); + return new RulePersistenceService(clusterService, client); } - public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { + public static void assertEqualRules( + Map> mapOne, + Map> mapTwo, + boolean ruleUpdated + ) { assertEquals(mapOne.size(), mapTwo.size()); - for (Map.Entry entry : mapOne.entrySet()) { + for (Map.Entry> entry : mapOne.entrySet()) { String id = entry.getKey(); assertTrue(mapTwo.containsKey(id)); - Rule one = mapOne.get(id); - Rule two = mapTwo.get(id); + Rule one = mapOne.get(id); + Rule two = mapTwo.get(id); assertEqualRule(one, two, ruleUpdated); } } - public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { + public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { if (ruleUpdated) { - assertEquals(one.getFeature(), two.getFeature()); + assertEquals(one.getDescription(), two.getDescription()); + assertEquals(one.getFeatureType(), two.getFeatureType()); assertEquals(one.getLabel(), two.getLabel()); assertEquals(one.getAttributeMap(), two.getAttributeMap()); } else { diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java index e9abb281df8b2..0b8e9d186a6cb 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java @@ -11,7 +11,6 @@ import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.wlm.Rule; import java.io.IOException; import java.util.HashMap; @@ -19,7 +18,9 @@ import java.util.Set; import static org.opensearch.plugin.wlm.RuleTestUtils.PATTERN_ONE; +import static org.opensearch.plugin.wlm.RuleTestUtils.SEARCH_AFTER; import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; +import static org.opensearch.plugin.wlm.rule.QueryGroupAttribute.INDEX_PATTERN; public class GetRuleRequestTests extends OpenSearchTestCase { @@ -27,7 +28,7 @@ public class GetRuleRequestTests extends OpenSearchTestCase { * Test case to verify the serialization and deserialization of GetRuleRequest */ public void testSerialization() throws IOException { - GetRuleRequest request = new GetRuleRequest(_ID_ONE, Map.of(Rule.RuleAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE))); + GetRuleRequest request = new GetRuleRequest(_ID_ONE, Map.of(INDEX_PATTERN, Set.of(PATTERN_ONE)), null); assertEquals(_ID_ONE, request.get_id()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); @@ -41,7 +42,7 @@ public void testSerialization() throws IOException { * Test case to verify the serialization and deserialization of GetRuleRequest when name is null */ public void testSerializationWithNull() throws IOException { - GetRuleRequest request = new GetRuleRequest((String) null, new HashMap<>()); + GetRuleRequest request = new GetRuleRequest((String) null, new HashMap<>(), SEARCH_AFTER); assertNull(request.get_id()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java index 47545e554803e..9c5198704058c 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java @@ -8,19 +8,21 @@ package org.opensearch.plugin.wlm.rule.action; +import org.opensearch.autotagging.Rule; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.wlm.Rule; import java.io.IOException; import java.util.HashMap; import java.util.Map; +import static org.opensearch.plugin.wlm.RuleTestUtils.SEARCH_AFTER; import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; import static org.opensearch.plugin.wlm.RuleTestUtils.assertEqualRules; import static org.opensearch.plugin.wlm.RuleTestUtils.ruleMap; @@ -33,9 +35,9 @@ public class GetRuleResponseTests extends OpenSearchTestCase { * Test case to verify the serialization and deserialization of GetRuleResponse */ public void testSerializationSingleRule() throws IOException { - Map map = new HashMap<>(); + Map> map = new HashMap<>(); map.put(_ID_ONE, ruleOne); - GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), RestStatus.OK); + GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), null, RestStatus.OK); assertEquals(response.getRules(), map); BytesStreamOutput out = new BytesStreamOutput(); @@ -51,7 +53,7 @@ public void testSerializationSingleRule() throws IOException { * Test case to verify the serialization and deserialization of GetRuleResponse when the result contains multiple rules */ public void testSerializationMultipleRule() throws IOException { - GetRuleResponse response = new GetRuleResponse(ruleMap(), RestStatus.OK); + GetRuleResponse response = new GetRuleResponse(ruleMap(), SEARCH_AFTER, RestStatus.OK); assertEquals(response.getRules(), ruleMap()); BytesStreamOutput out = new BytesStreamOutput(); @@ -68,8 +70,8 @@ public void testSerializationMultipleRule() throws IOException { * Test case to verify the serialization and deserialization of GetRuleResponse when the result is empty */ public void testSerializationNull() throws IOException { - Map map = new HashMap<>(); - GetRuleResponse response = new GetRuleResponse(map, RestStatus.OK); + Map> map = new HashMap<>(); + GetRuleResponse response = new GetRuleResponse(map, SEARCH_AFTER, RestStatus.OK); assertEquals(response.getRules(), map); BytesStreamOutput out = new BytesStreamOutput(); @@ -85,21 +87,25 @@ public void testSerializationNull() throws IOException { * Test case to verify the toXContent of GetRuleResponse */ public void testToXContentGetSingleRule() throws IOException { - Map map = new HashMap<>(); + Map> map = new HashMap<>(); map.put(_ID_ONE, ruleOne); - GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), RestStatus.OK); + GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), SEARCH_AFTER, RestStatus.OK); XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); String expected = "{\n" + " \"rules\" : [\n" + " {\n" + " \"_id\" : \"AgfUO5Ja9yfvhdONlYi3TQ==\",\n" + + " \"description\" : \"description_1\",\n" + " \"index_pattern\" : [\n" + " \"pattern_1\"\n" + " ],\n" + " \"query_group\" : \"label_one\",\n" + " \"updated_at\" : \"2024-01-26T08:58:57.558Z\"\n" + " }\n" + + " ],\n" + + " \"search_after\" : [\n" + + " \"search_after_id\"\n" + " ]\n" + "}"; assertEquals(expected, actual); @@ -110,7 +116,7 @@ public void testToXContentGetSingleRule() throws IOException { */ public void testToXContentGetZeroRule() throws IOException { XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); - GetRuleResponse otherResponse = new GetRuleResponse(new HashMap<>(), RestStatus.OK); + GetRuleResponse otherResponse = new GetRuleResponse(new HashMap<>(), null, RestStatus.OK); String actual = otherResponse.toXContent(builder, mock(ToXContent.Params.class)).toString(); String expected = "{\n" + " \"rules\" : [ ]\n" + "}"; assertEquals(expected, actual); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java index fa8225cf3552e..da426760086b7 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java @@ -11,11 +11,14 @@ import org.opensearch.ResourceNotFoundException; import org.opensearch.action.get.GetRequestBuilder; import org.opensearch.action.get.GetResponse; +import org.opensearch.action.search.SearchRequestBuilder; import org.opensearch.client.Client; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.action.ActionListener; import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; +import org.opensearch.search.sort.SortOrder; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; @@ -24,25 +27,26 @@ import org.mockito.ArgumentCaptor; +import static org.opensearch.autotagging.Rule._ID_STRING; import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; import static org.opensearch.plugin.wlm.RuleTestUtils.assertEqualRules; import static org.opensearch.plugin.wlm.RuleTestUtils.ruleOne; import static org.opensearch.plugin.wlm.RuleTestUtils.setUpRulePersistenceService; -import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.RULE_INDEX; +import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST; +import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.RULES_INDEX; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.clearInvocations; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @SuppressWarnings("unchecked") public class RulePersistenceServiceTests extends OpenSearchTestCase { - /** - * Test case to validate the logic for get Rule by id - */ public void testGetRuleById() throws IOException { String ruleSource = ruleOne.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString(); ActionListener listener = mock(ActionListener.class); @@ -53,14 +57,14 @@ public void testGetRuleById() throws IOException { when(getResponse.isExists()).thenReturn(true); when(getResponse.getSourceAsString()).thenReturn(ruleSource); - when(client.prepareGet(eq(RULE_INDEX), eq(_ID_ONE))).thenReturn(getRequestBuilder); + when(client.prepareGet(eq(RULES_INDEX), eq(_ID_ONE))).thenReturn(getRequestBuilder); doAnswer(invocation -> { ActionListener actionListener = invocation.getArgument(0); actionListener.onResponse(getResponse); return null; }).when(getRequestBuilder).execute(any(ActionListener.class)); - rulePersistenceService.getRule(_ID_ONE, new HashMap<>(), listener); + rulePersistenceService.getRule(_ID_ONE, new HashMap<>(), null, listener); ArgumentCaptor captor = ArgumentCaptor.forClass(GetRuleResponse.class); verify(listener).onResponse(captor.capture()); @@ -70,9 +74,6 @@ public void testGetRuleById() throws IOException { clearInvocations(client, getRequestBuilder, getResponse, listener); } - /** - * Test case to validate the get rule for an invalid id - */ public void testGetRuleByIdNotFound() { String nonExistentRuleId = "non-existent-rule"; RulePersistenceService rulePersistenceService = setUpRulePersistenceService(); @@ -81,7 +82,7 @@ public void testGetRuleByIdNotFound() { GetResponse getResponse = mock(GetResponse.class); ActionListener listener = mock(ActionListener.class); - when(client.prepareGet(RULE_INDEX, nonExistentRuleId)).thenReturn(getRequestBuilder); + when(client.prepareGet(RULES_INDEX, nonExistentRuleId)).thenReturn(getRequestBuilder); when(getResponse.isExists()).thenReturn(false); doAnswer(invocation -> { @@ -90,7 +91,7 @@ public void testGetRuleByIdNotFound() { return null; }).when(getRequestBuilder).execute(any(ActionListener.class)); - rulePersistenceService.getRule(nonExistentRuleId, new HashMap<>(), listener); + rulePersistenceService.getRule(nonExistentRuleId, new HashMap<>(), null, listener); ArgumentCaptor captor = ArgumentCaptor.forClass(Exception.class); verify(listener).onFailure(captor.capture()); @@ -98,4 +99,34 @@ public void testGetRuleByIdNotFound() { assertTrue(exception instanceof ResourceNotFoundException); clearInvocations(client, getRequestBuilder, getResponse, listener); } + + public void testBuildGetAllRuleSearchRequest() { + RulePersistenceService rulePersistenceService = setUpRulePersistenceService(); + Client client = rulePersistenceService.getClient(); + SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + + when(client.prepareSearch(anyString())).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.setQuery(any(BoolQueryBuilder.class))).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST)).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.addSort(eq(_ID_STRING), eq(SortOrder.ASC))).thenReturn(searchRequestBuilder); + + rulePersistenceService.buildGetAllRuleSearchRequest(ruleOne.getAttributeMap(), _ID_ONE); + ArgumentCaptor captor = ArgumentCaptor.forClass(Object[].class); + verify(searchRequestBuilder).searchAfter(captor.capture()); + assertEquals(_ID_ONE, captor.getValue()[0]); + } + + public void testBuildGetAllRuleSearchRequest_noSearchAfter() { + RulePersistenceService rulePersistenceService = setUpRulePersistenceService(); + Client client = rulePersistenceService.getClient(); + SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + + when(client.prepareSearch(anyString())).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.setQuery(any(BoolQueryBuilder.class))).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST)).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.addSort(eq(_ID_STRING), eq(SortOrder.ASC))).thenReturn(searchRequestBuilder); + + rulePersistenceService.buildGetAllRuleSearchRequest(ruleOne.getAttributeMap(), null); + verify(searchRequestBuilder, times(0)).searchAfter(any()); + } } diff --git a/server/src/main/java/org/opensearch/wlm/Rule.java b/server/src/main/java/org/opensearch/wlm/Rule.java deleted file mode 100644 index 273abc455c7e3..0000000000000 --- a/server/src/main/java/org/opensearch/wlm/Rule.java +++ /dev/null @@ -1,359 +0,0 @@ -/* - * 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.wlm; - -import org.opensearch.common.ValidationException; -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 org.opensearch.core.xcontent.ToXContentObject; -import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.core.xcontent.XContentParseException; -import org.opensearch.core.xcontent.XContentParser; -import org.joda.time.Instant; - -import java.io.IOException; -import java.util.Arrays; -import java.util.HashMap; -import java.util.HashSet; -import java.util.Map; -import java.util.Objects; -import java.util.Set; - -import static org.opensearch.cluster.metadata.QueryGroup.isValid; - -/** - * Represents a rule schema used for automatic query tagging in the system. - * This class encapsulates the criteria (defined through attributes) for automatically applying relevant - * tags to queries based on matching attribute patterns. This class provides an in-memory representation - * of a rule. The indexed view may differ in representation. - * { - * "_id": "fwehf8302582mglfio349==", - * "index_pattern": ["logs123", "user*"], - * "query_group": "dev_query_group_id", - * "updated_at": "01-10-2025T21:23:21.456Z" - * } - * @opensearch.experimental - */ -public class Rule implements Writeable, ToXContentObject { - private final Map> attributeMap; - private final Feature feature; - private final String label; - private final String updatedAt; - public static final String _ID_STRING = "_id"; - public static final String UPDATED_AT_STRING = "updated_at"; - public static final int MAX_NUMBER_OF_VALUES_PER_ATTRIBUTE = 10; - public static final int MAX_CHARACTER_LENGTH_PER_ATTRIBUTE_VALUE_STRING = 100; - - public Rule(Map> attributeMap, String label, String updatedAt, Feature feature) { - ValidationException validationException = new ValidationException(); - validateRuleInputs(attributeMap, label, updatedAt, feature, validationException); - if (!validationException.validationErrors().isEmpty()) { - throw new IllegalArgumentException(validationException); - } - - this.attributeMap = attributeMap; - this.feature = feature; - this.label = label; - this.updatedAt = updatedAt; - } - - public Rule(StreamInput in) throws IOException { - this( - in.readMap((i) -> RuleAttribute.fromName(i.readString()), i -> new HashSet<>(i.readStringList())), - in.readString(), - in.readString(), - Feature.fromName(in.readString()) - ); - } - - public static void requireNonNullOrEmpty(String value, String message, ValidationException validationException) { - if (value == null || value.isEmpty()) { - validationException.addValidationError(message); - } - } - - public static void validateRuleInputs( - Map> attributeMap, - String label, - String updatedAt, - Feature feature, - ValidationException validationException - ) { - if (feature == null) { - validationException.addValidationError("Couldn't identify which feature the rule belongs to. Rule feature name can't be null."); - } - requireNonNullOrEmpty(label, "Rule label can't be null or empty", validationException); - requireNonNullOrEmpty(updatedAt, "Rule update time can't be null or empty", validationException); - if (attributeMap == null || attributeMap.isEmpty()) { - validationException.addValidationError("Rule should have at least 1 attribute requirement"); - } - if (updatedAt != null && !isValid(Instant.parse(updatedAt).getMillis())) { - validationException.addValidationError("Rule update time is not a valid epoch"); - } - if (attributeMap != null && feature != null) { - validateAttributeMap(attributeMap, feature, validationException); - } - } - - public static void validateAttributeMap( - Map> attributeMap, - Feature feature, - ValidationException validationException - ) { - for (Map.Entry> entry : attributeMap.entrySet()) { - RuleAttribute ruleAttribute = entry.getKey(); - Set attributeValues = entry.getValue(); - if (!feature.isValidAttribute(ruleAttribute)) { - validationException.addValidationError( - ruleAttribute.getName() + " is not a valid attribute name under the feature: " + feature.getName() - ); - } - if (attributeValues.size() > MAX_NUMBER_OF_VALUES_PER_ATTRIBUTE) { - validationException.addValidationError( - "Each attribute can only have a maximum of 10 values. The input attribute " + ruleAttribute + " exceeds this limit." - ); - } - for (String attributeValue : attributeValues) { - if (attributeValue.isEmpty() || attributeValue.length() > MAX_CHARACTER_LENGTH_PER_ATTRIBUTE_VALUE_STRING) { - validationException.addValidationError( - "Attribute value [" + attributeValue + "] is invalid (empty or exceeds 100 characters)" - ); - } - } - } - } - - @Override - public void writeTo(StreamOutput out) throws IOException { - out.writeMap(attributeMap, RuleAttribute::writeTo, StreamOutput::writeStringCollection); - out.writeString(label); - out.writeString(updatedAt); - out.writeString(feature.getName()); - } - - public static Rule fromXContent(final XContentParser parser) throws IOException { - return Builder.fromXContent(parser).build(); - } - - public String getLabel() { - return label; - } - - public String getUpdatedAt() { - return updatedAt; - } - - public Feature getFeature() { - return feature; - } - - public Map> getAttributeMap() { - return attributeMap; - } - - /** - * This enum enumerates the features that can use the Rule Based Auto-tagging - * @opensearch.experimental - */ - public enum Feature { - QUERY_GROUP("query_group", Set.of(RuleAttribute.INDEX_PATTERN)); - - private final String name; - private final Set allowedAttributes; - - Feature(String name, Set allowedAttributes) { - this.name = name; - this.allowedAttributes = allowedAttributes; - } - - public String getName() { - return name; - } - - public Set getAllowedAttributes() { - return allowedAttributes; - } - - public boolean isValidAttribute(RuleAttribute attribute) { - return allowedAttributes.contains(attribute); - } - - public static boolean isValidFeature(String s) { - return Arrays.stream(values()).anyMatch(feature -> feature.getName().equalsIgnoreCase(s)); - } - - public static Feature fromName(String s) { - return Arrays.stream(values()) - .filter(feature -> feature.getName().equalsIgnoreCase(s)) - .findFirst() - .orElseThrow(() -> new IllegalArgumentException("Invalid value for Feature: " + s)); - } - } - - /** - * This RuleAttribute enum contains the attribute names for a rule. - * @opensearch.experimental - */ - public enum RuleAttribute { - INDEX_PATTERN("index_pattern"); - - private final String name; - - RuleAttribute(String name) { - this.name = name; - } - - public String getName() { - return name; - } - - public static void writeTo(StreamOutput out, RuleAttribute ruleAttribute) throws IOException { - out.writeString(ruleAttribute.getName()); - } - - public static RuleAttribute fromName(String s) { - for (RuleAttribute attribute : values()) { - if (attribute.getName().equalsIgnoreCase(s)) return attribute; - - } - throw new IllegalArgumentException("Invalid value for RuleAttribute: " + s); - } - } - - @Override - public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { - builder.startObject(); - String id = params.param(_ID_STRING); - if (id != null) { - builder.field(_ID_STRING, id); - } - for (Map.Entry> entry : attributeMap.entrySet()) { - builder.array(entry.getKey().getName(), entry.getValue().toArray(new String[0])); - } - builder.field(feature.getName(), label); - builder.field(UPDATED_AT_STRING, updatedAt); - builder.endObject(); - return builder; - } - - @Override - public boolean equals(Object o) { - if (this == o) return true; - if (o == null || getClass() != o.getClass()) return false; - Rule that = (Rule) o; - return Objects.equals(label, that.label) - && Objects.equals(feature, that.feature) - && Objects.equals(attributeMap, that.attributeMap) - && Objects.equals(updatedAt, that.updatedAt); - } - - @Override - public int hashCode() { - return Objects.hash(label, feature, attributeMap, updatedAt); - } - - /** - * builder method for the {@link Rule} - * @return Builder object - */ - public static Builder builder() { - return new Builder(); - } - - /** - * Builder class for {@link Rule} - * @opensearch.experimental - */ - public static class Builder { - private Map> attributeMap; - private Feature feature; - private String label; - private String updatedAt; - - private Builder() {} - - public static Builder fromXContent(XContentParser parser) throws IOException { - if (parser.currentToken() == null) { - parser.nextToken(); - } - Builder builder = builder(); - XContentParser.Token token = parser.currentToken(); - - if (token != XContentParser.Token.START_OBJECT) { - throw new IllegalArgumentException("Expected START_OBJECT token but found [" + parser.currentName() + "]"); - } - Map> attributeMap1 = new HashMap<>(); - String fieldName = ""; - while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { - if (token == XContentParser.Token.FIELD_NAME) { - fieldName = parser.currentName(); - } else if (token.isValue()) { - if (Feature.isValidFeature(fieldName)) { - builder.feature(fieldName); - builder.label(parser.text()); - } else if (fieldName.equals(UPDATED_AT_STRING)) { - builder.updatedAt(parser.text()); - } else { - throw new IllegalArgumentException(fieldName + " is not a valid field in Rule"); - } - } else if (token == XContentParser.Token.START_ARRAY) { - fromXContentParseArray(parser, fieldName, attributeMap1); - } - } - return builder.attributeMap(attributeMap1); - } - - public static void fromXContentParseArray(XContentParser parser, String fieldName, Map> attributeMap) - throws IOException { - RuleAttribute ruleAttribute = RuleAttribute.fromName(fieldName); - Set attributeValueSet = new HashSet<>(); - while (parser.nextToken() != XContentParser.Token.END_ARRAY) { - if (parser.currentToken() == XContentParser.Token.VALUE_STRING) { - attributeValueSet.add(parser.text()); - } else { - throw new XContentParseException("Unexpected token in array: " + parser.currentToken()); - } - } - attributeMap.put(ruleAttribute, attributeValueSet); - } - - public Builder label(String label) { - this.label = label; - return this; - } - - public Builder attributeMap(Map> attributeMap) { - this.attributeMap = attributeMap; - return this; - } - - public Builder feature(String feature) { - this.feature = Feature.fromName(feature); - return this; - } - - public Builder updatedAt(String updatedAt) { - this.updatedAt = updatedAt; - return this; - } - - public Rule build() { - return new Rule(attributeMap, label, updatedAt, feature); - } - - public String getLabel() { - return label; - } - - public Map> getAttributeMap() { - return attributeMap; - } - } -} diff --git a/server/src/test/java/org/opensearch/wlm/RuleTests.java b/server/src/test/java/org/opensearch/wlm/RuleTests.java deleted file mode 100644 index 72a4abf1f3877..0000000000000 --- a/server/src/test/java/org/opensearch/wlm/RuleTests.java +++ /dev/null @@ -1,163 +0,0 @@ -/* - * 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.wlm; - -import org.opensearch.common.xcontent.json.JsonXContent; -import org.opensearch.core.common.io.stream.Writeable; -import org.opensearch.core.xcontent.ToXContent; -import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.test.AbstractSerializingTestCase; -import org.opensearch.wlm.Rule.Feature; -import org.opensearch.wlm.Rule.RuleAttribute; -import org.joda.time.Instant; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; - -import static org.opensearch.wlm.Rule.MAX_CHARACTER_LENGTH_PER_ATTRIBUTE_VALUE_STRING; -import static org.opensearch.wlm.Rule.MAX_NUMBER_OF_VALUES_PER_ATTRIBUTE; -import static org.opensearch.wlm.Rule._ID_STRING; - -public class RuleTests extends AbstractSerializingTestCase { - public static final String _ID = "AgfUfjw039vhdONlYi3TQ=="; - public static final String LABEL = "label"; - - static Rule createRandomRule(String label) { - Feature feature = randomFeature(); - return Rule.builder() - .label(label) - .feature(feature.getName()) - .attributeMap(randomAttributeMaps(feature)) - .updatedAt(Instant.now().toString()) - .build(); - } - - private static Feature randomFeature() { - return Feature.values()[randomIntBetween(0, Feature.values().length - 1)]; - } - - private static Map> randomAttributeMaps(Feature feature) { - Map> attributeMap = new HashMap<>(); - if (feature == null) { - return attributeMap; - } - List allowedAttributes = new ArrayList<>(feature.getAllowedAttributes()); - do { - attributeMap.clear(); - for (RuleAttribute currAttribute : allowedAttributes) { - if (randomBoolean()) { - attributeMap.put(currAttribute, randomAttributeValues()); - } - } - } while (attributeMap.isEmpty()); - return attributeMap; - } - - private static Set randomAttributeValues() { - Set res = new HashSet<>(); - int numberOfValues = randomIntBetween(1, MAX_NUMBER_OF_VALUES_PER_ATTRIBUTE); - for (int i = 0; i < numberOfValues; i++) { - res.add(randomAlphaOfLength(randomIntBetween(1, MAX_CHARACTER_LENGTH_PER_ATTRIBUTE_VALUE_STRING))); - } - return res; - } - - @Override - protected Rule doParseInstance(XContentParser parser) throws IOException { - return Rule.fromXContent(parser); - } - - @Override - protected Writeable.Reader instanceReader() { - return Rule::new; - } - - @Override - protected Rule createTestInstance() { - return createRandomRule(LABEL); - } - - static Rule buildRule(String label, String feature, Map> attributeListMap, String updatedAt) { - return Rule.builder().label(label).feature(feature).attributeMap(attributeListMap).updatedAt(updatedAt).build(); - } - - public void testInvalidFeature() { - assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, null, randomAttributeMaps(null), Instant.now().toString())); - assertThrows( - IllegalArgumentException.class, - () -> buildRule(LABEL, "invalid", randomAttributeMaps(null), Instant.now().toString()) - ); - } - - public void testInvalidLabel() { - assertThrows(IllegalArgumentException.class, () -> createRandomRule(null)); - assertThrows(IllegalArgumentException.class, () -> createRandomRule("")); - } - - public void testInvalidUpdateTime() { - Feature feature = randomFeature(); - assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, feature.toString(), randomAttributeMaps(feature), null)); - } - - public void testNullOrEmptyAttributeMap() { - Feature feature = randomFeature(); - assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, feature.toString(), new HashMap<>(), Instant.now().toString())); - assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, feature.toString(), null, Instant.now().toString())); - } - - public void testInvalidAttributeMap() { - Map> map = new HashMap<>(); - map.put(RuleAttribute.INDEX_PATTERN, Set.of("")); - assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, randomFeature().getName(), map, Instant.now().toString())); - - map.put(RuleAttribute.INDEX_PATTERN, Set.of(randomAlphaOfLength(MAX_CHARACTER_LENGTH_PER_ATTRIBUTE_VALUE_STRING + 1))); - assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, randomFeature().getName(), map, Instant.now().toString())); - - map.put(RuleAttribute.INDEX_PATTERN, new HashSet<>()); - for (int i = 0; i < MAX_NUMBER_OF_VALUES_PER_ATTRIBUTE + 1; i++) { - map.get(RuleAttribute.INDEX_PATTERN).add(String.valueOf(i)); - } - assertThrows(IllegalArgumentException.class, () -> buildRule(LABEL, randomFeature().getName(), map, Instant.now().toString())); - } - - public void testValidRule() { - Map> map = Map.of(RuleAttribute.INDEX_PATTERN, Set.of("index*", "log*")); - String updatedAt = Instant.now().toString(); - Rule rule = buildRule(LABEL, Feature.QUERY_GROUP.getName(), map, updatedAt); - assertNotNull(rule.getLabel()); - assertEquals(LABEL, rule.getLabel()); - assertNotNull(updatedAt); - assertEquals(updatedAt, rule.getUpdatedAt()); - Map> resultMap = rule.getAttributeMap(); - assertNotNull(resultMap); - assertFalse(resultMap.isEmpty()); - assertNotNull(rule.getFeature()); - assertEquals(Feature.QUERY_GROUP, rule.getFeature()); - } - - public void testToXContent() throws IOException { - Map> map = Map.of(RuleAttribute.INDEX_PATTERN, Set.of("log*")); - String updatedAt = Instant.now().toString(); - Rule rule = buildRule(LABEL, Feature.QUERY_GROUP.getName(), map, updatedAt); - - XContentBuilder builder = JsonXContent.contentBuilder(); - rule.toXContent(builder, new ToXContent.MapParams(Map.of(_ID_STRING, _ID))); - - assertEquals( - "{\"_id\":\"" + _ID + "\",\"index_pattern\":[\"log*\"],\"query_group\":\"label\",\"updated_at\":\"" + updatedAt + "\"}", - builder.toString() - ); - } -} From 1816b0bda394bf85bf4cc56e6146cf5593ed4162 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Fri, 21 Mar 2025 17:31:55 -0700 Subject: [PATCH 03/20] rebase from main after the schema merged Signed-off-by: Ruirui Zhang --- .../plugin/wlm/rule/QueryGroupAttribute.java | 25 ++- .../wlm/rule/QueryGroupFeatureType.java | 10 +- .../wlm/rule/action/GetRuleRequest.java | 21 +-- .../wlm/rule/action/GetRuleResponse.java | 9 +- .../rule/action/TransportGetRuleAction.java | 2 +- .../wlm/rule/rest/RestGetRuleAction.java | 2 +- .../rule/service/RulePersistenceService.java | 167 +++++++++++------- .../opensearch/plugin/wlm/RuleTestUtils.java | 46 ++--- .../wlm/rule/action/GetRuleRequestTests.java | 8 +- .../wlm/rule/action/GetRuleResponseTests.java | 9 +- .../service/RulePersistenceServiceTests.java | 44 +---- 11 files changed, 169 insertions(+), 174 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java index 906b68e37efff..f57a92bce5c9d 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java @@ -9,11 +9,12 @@ package org.opensearch.plugin.wlm.rule; import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.AutoTaggingRegistry; + +import java.util.HashMap; +import java.util.Map; /** * Attributes specific to the query group feature. - * @opensearch.experimental */ public enum QueryGroupAttribute implements Attribute { INDEX_PATTERN("index_pattern"); @@ -22,12 +23,7 @@ public enum QueryGroupAttribute implements Attribute { QueryGroupAttribute(String name) { this.name = name; - } - - static { - for (QueryGroupAttribute attr : QueryGroupAttribute.values()) { - attr.registerAttribute(); - } + validateAttribute(); } @Override @@ -35,11 +31,6 @@ public String getName() { return name; } - @Override - public void registerAttribute() { - AutoTaggingRegistry.registerAttribute(this); - } - public static QueryGroupAttribute fromName(String name) { for (QueryGroupAttribute attr : QueryGroupAttribute.values()) { if (attr.getName().equals(name)) { @@ -48,4 +39,12 @@ public static QueryGroupAttribute fromName(String name) { } throw new IllegalArgumentException("Unknown QueryGroupAttribute: " + name); } + + public static Map toMap() { + Map map = new HashMap<>(); + for (QueryGroupAttribute attr : QueryGroupAttribute.values()) { + map.put(attr.getName(), attr); + } + return map; + } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java index 18e227d23e888..feb762b8ad99b 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java @@ -12,18 +12,14 @@ import org.opensearch.autotagging.AutoTaggingRegistry; import org.opensearch.autotagging.FeatureType; -import java.util.Set; +import java.util.Map; -/** - * Represents the feature type for "query_group" in the OpenSearch workload management plugin. - * @opensearch.experimental - */ public class QueryGroupFeatureType implements FeatureType { public static final QueryGroupFeatureType INSTANCE = new QueryGroupFeatureType(); public static final String NAME = "query_group"; private static final int MAX_ATTRIBUTE_VALUES = 10; private static final int MAX_ATTRIBUTE_VALUE_LENGTH = 100; - private static final Set ALLOWED_ATTRIBUTES = Set.of(QueryGroupAttribute.values()); + private static final Map ALLOWED_ATTRIBUTES = QueryGroupAttribute.toMap(); private QueryGroupFeatureType() {} @@ -47,7 +43,7 @@ public int getMaxCharLengthPerAttributeValue() { } @Override - public Set getAllowedAttributes() { + public Map getAllowedAttributesRegistry() { return ALLOWED_ATTRIBUTES; } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java index 0f9b4cef43925..181bd9583ce34 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java @@ -13,6 +13,7 @@ import org.opensearch.autotagging.Attribute; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import java.io.IOException; import java.util.HashSet; @@ -28,18 +29,18 @@ * @opensearch.experimental */ public class GetRuleRequest extends ActionRequest { - private final String _id; + private final String id; private final Map> attributeFilters; private final String searchAfter; /** * Constructor for GetRuleRequest - * @param _id - Rule _id that we want to get + * @param id - Rule id that we want to get * @param attributeFilters - Attributes that we want to filter on * @param searchAfter - The sort values from the last document of the previous page, used for pagination */ - public GetRuleRequest(String _id, Map> attributeFilters, String searchAfter) { - this._id = _id; + public GetRuleRequest(String id, Map> attributeFilters, String searchAfter) { + this.id = id; this.attributeFilters = attributeFilters; this.searchAfter = searchAfter; } @@ -50,8 +51,8 @@ public GetRuleRequest(String _id, Map> attributeFilters, */ public GetRuleRequest(StreamInput in) throws IOException { super(in); - _id = in.readOptionalString(); - attributeFilters = in.readMap(Attribute::from, i -> new HashSet<>(i.readStringList())); + id = in.readOptionalString(); + attributeFilters = in.readMap(i -> Attribute.from(i, QueryGroupFeatureType.INSTANCE), i -> new HashSet<>(i.readStringList())); searchAfter = in.readOptionalString(); } @@ -63,16 +64,16 @@ public ActionRequestValidationException validate() { @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); - out.writeOptionalString(_id); + out.writeOptionalString(id); out.writeMap(attributeFilters, (output, attribute) -> attribute.writeTo(output), StreamOutput::writeStringCollection); out.writeOptionalString(searchAfter); } /** - * _id getter + * id getter */ - public String get_id() { - return _id; + public String getId() { + return id; } /** diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java index 2b4ffd6b4f431..2f45e372b8325 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java @@ -16,7 +16,6 @@ import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import java.io.IOException; import java.util.Map; @@ -43,7 +42,7 @@ * @opensearch.experimental */ public class GetRuleResponse extends ActionResponse implements ToXContent, ToXContentObject { - private final Map> rules; + private final Map rules; private final String searchAfter; private final RestStatus restStatus; @@ -53,7 +52,7 @@ public class GetRuleResponse extends ActionResponse implements ToXContent, ToXCo * @param searchAfter - The searchAfter field for the response * @param restStatus - The restStatus for the response */ - public GetRuleResponse(final Map> rules, String searchAfter, RestStatus restStatus) { + public GetRuleResponse(final Map rules, String searchAfter, RestStatus restStatus) { this.rules = rules; this.searchAfter = searchAfter; this.restStatus = restStatus; @@ -78,7 +77,7 @@ public void writeTo(StreamOutput out) throws IOException { public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); builder.startArray("rules"); - for (Map.Entry> entry : rules.entrySet()) { + for (Map.Entry entry : rules.entrySet()) { entry.getValue().toXContent(builder, new MapParams(Map.of(_ID_STRING, entry.getKey()))); } builder.endArray(); @@ -92,7 +91,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws /** * rules getter */ - public Map> getRules() { + public Map getRules() { return rules; } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java index a7c2e4cac2e04..7b9a77d78e328 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java @@ -42,6 +42,6 @@ public TransportGetRuleAction( @Override protected void doExecute(Task task, GetRuleRequest request, ActionListener listener) { - rulePersistenceService.getRule(request.get_id(), request.getAttributeFilters(), request.getSearchAfter(), listener); + rulePersistenceService.getRule(request.getId(), request.getAttributeFilters(), request.getSearchAfter(), listener); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java index 01f734e12d91e..c3a9844fc6f33 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java @@ -9,7 +9,6 @@ package org.opensearch.plugin.wlm.rule.rest; import org.opensearch.autotagging.Attribute; -import org.opensearch.client.node.NodeClient; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.plugin.wlm.rule.QueryGroupAttribute; @@ -22,6 +21,7 @@ import org.opensearch.rest.RestRequest; import org.opensearch.rest.RestResponse; import org.opensearch.rest.action.RestResponseListener; +import org.opensearch.transport.client.node.NodeClient; import java.io.IOException; import java.util.Arrays; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java index aa45dda3ba210..a903f80b2249e 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java @@ -17,7 +17,6 @@ import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.Rule; import org.opensearch.autotagging.Rule.Builder; -import org.opensearch.client.Client; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; import org.opensearch.common.util.concurrent.ThreadContext; @@ -33,9 +32,11 @@ import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; import org.opensearch.search.SearchHit; import org.opensearch.search.sort.SortOrder; +import org.opensearch.transport.client.Client; import java.io.IOException; import java.util.Arrays; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; @@ -53,12 +54,8 @@ public class RulePersistenceService { private final Client client; private final ClusterService clusterService; private static final Logger logger = LogManager.getLogger(RulePersistenceService.class); - public static final int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 100; + public static final int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 50; - /** - * Constructor for RulePersistenceService - * @param client {@link Client} - The client to be used by RulePersistenceService - */ @Inject public RulePersistenceService(final ClusterService clusterService, final Client client) { this.clusterService = clusterService; @@ -90,18 +87,14 @@ public void getRule( * @param id - The id of the rule to get * @param listener - ActionListener for GetRuleResponse */ - private void fetchRuleById(String id, ActionListener listener) { - ThreadContext.StoredContext storedContext = client.threadPool().getThreadContext().stashContext(); - client.prepareGet(RULES_INDEX, id).execute(ActionListener.wrap(getResponse -> { - try (ThreadContext.StoredContext context = storedContext) { - handleGetOneRuleResponse(id, getResponse, listener); - } - }, e -> { - try (ThreadContext.StoredContext context = storedContext) { - logger.error("Failed to fetch rule with ID {}: {}", id, e.getMessage()); - listener.onFailure(e); - } - })); + void fetchRuleById(String id, ActionListener listener) { + try (ThreadContext.StoredContext context = getContext()) { + client.prepareGet(RULES_INDEX, id) + .execute(ActionListener.wrap(getResponse -> handleGetOneRuleResponse(id, getResponse, listener), e -> { + logger.error("Failed to fetch rule with ID {}: {}", id, e.getMessage()); + listener.onFailure(e); + })); + } } /** @@ -112,7 +105,7 @@ private void fetchRuleById(String id, ActionListener listener) */ private void handleGetOneRuleResponse(String id, GetResponse getResponse, ActionListener listener) { if (getResponse.isExists()) { - try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { + try (ThreadContext.StoredContext context = getContext()) { XContentParser parser = MediaTypeRegistry.JSON.xContent() .createParser( NamedXContentRegistry.EMPTY, @@ -142,18 +135,44 @@ private void handleGetOneRuleResponse(String id, GetResponse getResponse, Action * @param listener - ActionListener for GetRuleResponse */ private void fetchAllRules(Map> attributeFilters, String searchAfter, ActionListener listener) { - ThreadContext.StoredContext storedContext = client.threadPool().getThreadContext().stashContext(); - SearchRequestBuilder searchRequest = buildGetAllRuleSearchRequest(attributeFilters, searchAfter); - searchRequest.execute(ActionListener.wrap(searchResponse -> { - try (ThreadContext.StoredContext context = storedContext) { - listener.onResponse(handleGetAllRuleResponse(searchResponse)); + try (ThreadContext.StoredContext context = getContext()) { + client.prepareSearch(RULES_INDEX) + .setSize(0) + .execute( + ActionListener.wrap(countResponse -> handleCountResponse(countResponse, attributeFilters, searchAfter, listener), e -> { + logger.error("Failed to check if index is empty: {}", e.getMessage()); + listener.onFailure(e); + }) + ); + } + } + + /** + * Processes the count response from a search query on the rules index. + * If no rules exist, it responds with an empty result. + * Otherwise, it constructs and executes a search request to retrieve all rules. + * @param countResponse The response from the count query on the rules index. + * @param attributeFilters A map of attribute filters to apply in the search query. + * @param searchAfter The searchAfter parameter for pagination. + * @param listener The action listener to handle the final response or failure. + */ + void handleCountResponse( + SearchResponse countResponse, + Map> attributeFilters, + String searchAfter, + ActionListener listener + ) { + try (ThreadContext.StoredContext context = getContext()) { + if (countResponse.getHits().getTotalHits().value() == 0) { + listener.onResponse(new GetRuleResponse(new HashMap<>(), null, RestStatus.OK)); + return; } - }, e -> { - try (ThreadContext.StoredContext context = storedContext) { + SearchRequestBuilder searchRequest = buildGetAllRuleSearchRequest(attributeFilters, searchAfter); + searchRequest.execute(ActionListener.wrap(searchResponse -> handleGetAllRuleResponse(searchResponse, listener), e -> { logger.error("Failed to fetch all rules: {}", e.getMessage()); listener.onFailure(e); - } - })); + })); + } } /** @@ -163,58 +182,72 @@ private void fetchAllRules(Map> attributeFilters, String * @param searchAfter A cursor to enable pagination, used to fetch results after a specific document. */ SearchRequestBuilder buildGetAllRuleSearchRequest(Map> attributeFilters, String searchAfter) { - BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); - for (Map.Entry> entry : attributeFilters.entrySet()) { - Attribute attribute = entry.getKey(); - Set values = entry.getValue(); - if (values != null && !values.isEmpty()) { - BoolQueryBuilder attributeQuery = QueryBuilders.boolQuery(); - for (String value : values) { - attributeQuery.should(QueryBuilders.matchQuery(attribute.getName(), value)); + try (ThreadContext.StoredContext context = getContext()) { + BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); + for (Map.Entry> entry : attributeFilters.entrySet()) { + Attribute attribute = entry.getKey(); + Set values = entry.getValue(); + if (values != null && !values.isEmpty()) { + BoolQueryBuilder attributeQuery = QueryBuilders.boolQuery(); + for (String value : values) { + attributeQuery.should(QueryBuilders.matchQuery(attribute.getName(), value)); + } + boolQuery.must(attributeQuery); } - boolQuery.must(attributeQuery); } + boolQuery.filter(QueryBuilders.existsQuery(QueryGroupFeatureType.NAME)); + SearchRequestBuilder searchRequest = client.prepareSearch(RULES_INDEX) + .setQuery(boolQuery) + .setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST) + .addSort(_ID_STRING, SortOrder.ASC); + if (searchAfter != null) { + searchRequest.searchAfter(new Object[] { searchAfter }); + } + return searchRequest; } - boolQuery.filter(QueryBuilders.existsQuery(QueryGroupFeatureType.NAME)); - SearchRequestBuilder searchRequest = client.prepareSearch(RULES_INDEX) - .setQuery(boolQuery) - .setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST) - .addSort(_ID_STRING, SortOrder.ASC); - if (searchAfter != null) { - searchRequest.searchAfter(new Object[] { searchAfter }); - } - return searchRequest; } /** * Process searchResponse from index and send a GetRuleResponse * @param searchResponse - Response received from index + * @param listener - ActionListener for GetRuleResponse */ - GetRuleResponse handleGetAllRuleResponse(SearchResponse searchResponse) { + void handleGetAllRuleResponse(SearchResponse searchResponse, ActionListener listener) { List hits = Arrays.asList(searchResponse.getHits().getHits()); - Map> ruleMap = hits.stream().map(hit -> { - String hitId = hit.getId(); - try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { - XContentParser parser = MediaTypeRegistry.JSON.xContent() - .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, hit.getSourceAsString()); - return Map.entry(hitId, Builder.fromXContent(parser, QueryGroupFeatureType.INSTANCE).build()); - } catch (IOException e) { - logger.info( - "Issue met when parsing rule from hit, the feature type for rule id {} is probably not query_group: {}", - hitId, - e.getMessage() - ); - return null; - } - }).filter(Objects::nonNull).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - String nextSearchAfter = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); - return new GetRuleResponse(ruleMap, nextSearchAfter, RestStatus.OK); + try (ThreadContext.StoredContext context = getContext()) { + Map ruleMap = hits.stream().map(hit -> { + String hitId = hit.getId(); + try { + XContentParser parser = MediaTypeRegistry.JSON.xContent() + .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, hit.getSourceAsString()); + return Map.entry(hitId, Builder.fromXContent(parser, QueryGroupFeatureType.INSTANCE).build()); + } catch (IOException e) { + logger.info( + "Issue met when parsing rule from hit, the feature type for rule id {} is probably not query_group: {}", + hitId, + e.getMessage() + ); + return null; + } + }).filter(Objects::nonNull).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + String nextSearchAfter = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); + listener.onResponse(new GetRuleResponse(ruleMap, nextSearchAfter, RestStatus.OK)); + } + } + + private ThreadContext.StoredContext getContext() { + return client.threadPool().getThreadContext().stashContext(); + } + + private boolean isExistingQueryGroup(String queryGroupId) { + return clusterService.state().metadata().queryGroups().containsKey(queryGroupId); } - /** - * client getter - */ public Client getClient() { return client; } + + public ClusterService getClusterService() { + return clusterService; + } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java index 387cdb1bae6e7..e00614055d250 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java @@ -10,7 +10,9 @@ import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.Rule; -import org.opensearch.client.Client; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; @@ -18,6 +20,7 @@ import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import org.opensearch.plugin.wlm.rule.service.RulePersistenceService; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.client.Client; import java.util.Map; import java.util.Set; @@ -30,8 +33,8 @@ public class RuleTestUtils { public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; - public static final String LABEL_ONE = "label_one"; - public static final String LABEL_TWO = "label_two"; + public static final String FEATURE_VALUE_ONE = "feature_value_one"; + public static final String FEATURE_VALUE_TWO = "feature_value_two"; public static final String PATTERN_ONE = "pattern_1"; public static final String PATTERN_TWO = "pattern_2"; public static final String DESCRIPTION_ONE = "description_1"; @@ -39,57 +42,60 @@ public class RuleTestUtils { public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; public static final String SEARCH_AFTER = "search_after_id"; - public static final Map> attributes = Map.of(QueryGroupAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE)); - public static final Rule ruleOne = Rule.builder() + public static final Map> ATTRIBUTE_MAP = Map.of(QueryGroupAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE)); + public static final Rule ruleOne = Rule.builder() .description(DESCRIPTION_ONE) .featureType(QueryGroupFeatureType.INSTANCE) - .label(LABEL_ONE) + .featureValue(FEATURE_VALUE_ONE) .attributeMap(Map.of(QueryGroupAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE))) .updatedAt(TIMESTAMP_ONE) .build(); - public static final Rule ruleTwo = Rule.builder() + public static final Rule ruleTwo = Rule.builder() .description(DESCRIPTION_TWO) .featureType(QueryGroupFeatureType.INSTANCE) - .label(LABEL_TWO) + .featureValue(FEATURE_VALUE_TWO) .attributeMap(Map.of(QueryGroupAttribute.INDEX_PATTERN, Set.of(PATTERN_TWO))) .updatedAt(TIMESTAMP_TWO) .build(); - public static Map> ruleMap() { + public static Map ruleMap() { return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); } - public static RulePersistenceService setUpRulePersistenceService() { + public static RulePersistenceService setUpRulePersistenceService(Map queryGroupMap) { Client client = mock(Client.class); ClusterService clusterService = mock(ClusterService.class); + ClusterState clusterState = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); ThreadPool threadPool = mock(ThreadPool.class); + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); when(client.threadPool()).thenReturn(threadPool); when(threadPool.getThreadContext()).thenReturn(threadContext); + when(clusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + when(metadata.queryGroups()).thenReturn(queryGroupMap); return new RulePersistenceService(clusterService, client); } - public static void assertEqualRules( - Map> mapOne, - Map> mapTwo, - boolean ruleUpdated - ) { + public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { assertEquals(mapOne.size(), mapTwo.size()); - for (Map.Entry> entry : mapOne.entrySet()) { + for (Map.Entry entry : mapOne.entrySet()) { String id = entry.getKey(); assertTrue(mapTwo.containsKey(id)); - Rule one = mapOne.get(id); - Rule two = mapTwo.get(id); + Rule one = mapOne.get(id); + Rule two = mapTwo.get(id); assertEqualRule(one, two, ruleUpdated); } } - public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { + public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { if (ruleUpdated) { assertEquals(one.getDescription(), two.getDescription()); assertEquals(one.getFeatureType(), two.getFeatureType()); - assertEquals(one.getLabel(), two.getLabel()); + assertEquals(one.getFeatureValue(), two.getFeatureValue()); + assertEquals(one.getAttributeMap(), two.getAttributeMap()); assertEquals(one.getAttributeMap(), two.getAttributeMap()); } else { assertEquals(one, two); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java index 0b8e9d186a6cb..e0811b2d85e4f 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java @@ -29,12 +29,12 @@ public class GetRuleRequestTests extends OpenSearchTestCase { */ public void testSerialization() throws IOException { GetRuleRequest request = new GetRuleRequest(_ID_ONE, Map.of(INDEX_PATTERN, Set.of(PATTERN_ONE)), null); - assertEquals(_ID_ONE, request.get_id()); + assertEquals(_ID_ONE, request.getId()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); GetRuleRequest otherRequest = new GetRuleRequest(streamInput); - assertEquals(request.get_id(), otherRequest.get_id()); + assertEquals(request.getId(), otherRequest.getId()); assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters()); } @@ -43,12 +43,12 @@ public void testSerialization() throws IOException { */ public void testSerializationWithNull() throws IOException { GetRuleRequest request = new GetRuleRequest((String) null, new HashMap<>(), SEARCH_AFTER); - assertNull(request.get_id()); + assertNull(request.getId()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); GetRuleRequest otherRequest = new GetRuleRequest(streamInput); - assertEquals(request.get_id(), otherRequest.get_id()); + assertEquals(request.getId(), otherRequest.getId()); assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters()); } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java index 9c5198704058c..8c9983a1bb770 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java @@ -15,7 +15,6 @@ import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; @@ -35,7 +34,7 @@ public class GetRuleResponseTests extends OpenSearchTestCase { * Test case to verify the serialization and deserialization of GetRuleResponse */ public void testSerializationSingleRule() throws IOException { - Map> map = new HashMap<>(); + Map map = new HashMap<>(); map.put(_ID_ONE, ruleOne); GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), null, RestStatus.OK); assertEquals(response.getRules(), map); @@ -70,7 +69,7 @@ public void testSerializationMultipleRule() throws IOException { * Test case to verify the serialization and deserialization of GetRuleResponse when the result is empty */ public void testSerializationNull() throws IOException { - Map> map = new HashMap<>(); + Map map = new HashMap<>(); GetRuleResponse response = new GetRuleResponse(map, SEARCH_AFTER, RestStatus.OK); assertEquals(response.getRules(), map); @@ -87,7 +86,7 @@ public void testSerializationNull() throws IOException { * Test case to verify the toXContent of GetRuleResponse */ public void testToXContentGetSingleRule() throws IOException { - Map> map = new HashMap<>(); + Map map = new HashMap<>(); map.put(_ID_ONE, ruleOne); GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), SEARCH_AFTER, RestStatus.OK); XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); @@ -100,7 +99,7 @@ public void testToXContentGetSingleRule() throws IOException { + " \"index_pattern\" : [\n" + " \"pattern_1\"\n" + " ],\n" - + " \"query_group\" : \"label_one\",\n" + + " \"query_group\" : \"feature_value_one\",\n" + " \"updated_at\" : \"2024-01-26T08:58:57.558Z\"\n" + " }\n" + " ],\n" diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java index da426760086b7..d995cf19a3bd7 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java @@ -11,15 +11,12 @@ import org.opensearch.ResourceNotFoundException; import org.opensearch.action.get.GetRequestBuilder; import org.opensearch.action.get.GetResponse; -import org.opensearch.action.search.SearchRequestBuilder; -import org.opensearch.client.Client; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.action.ActionListener; import org.opensearch.core.xcontent.ToXContent; -import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; -import org.opensearch.search.sort.SortOrder; import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.transport.client.Client; import java.io.IOException; import java.util.HashMap; @@ -27,30 +24,25 @@ import org.mockito.ArgumentCaptor; -import static org.opensearch.autotagging.Rule._ID_STRING; import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; import static org.opensearch.plugin.wlm.RuleTestUtils.assertEqualRules; import static org.opensearch.plugin.wlm.RuleTestUtils.ruleOne; import static org.opensearch.plugin.wlm.RuleTestUtils.setUpRulePersistenceService; -import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST; import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.RULES_INDEX; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.clearInvocations; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @SuppressWarnings("unchecked") public class RulePersistenceServiceTests extends OpenSearchTestCase { - public void testGetRuleById() throws IOException { String ruleSource = ruleOne.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString(); ActionListener listener = mock(ActionListener.class); - RulePersistenceService rulePersistenceService = setUpRulePersistenceService(); + RulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); Client client = rulePersistenceService.getClient(); GetRequestBuilder getRequestBuilder = mock(GetRequestBuilder.class); GetResponse getResponse = mock(GetResponse.class); @@ -76,7 +68,7 @@ public void testGetRuleById() throws IOException { public void testGetRuleByIdNotFound() { String nonExistentRuleId = "non-existent-rule"; - RulePersistenceService rulePersistenceService = setUpRulePersistenceService(); + RulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); Client client = rulePersistenceService.getClient(); GetRequestBuilder getRequestBuilder = mock(GetRequestBuilder.class); GetResponse getResponse = mock(GetResponse.class); @@ -99,34 +91,4 @@ public void testGetRuleByIdNotFound() { assertTrue(exception instanceof ResourceNotFoundException); clearInvocations(client, getRequestBuilder, getResponse, listener); } - - public void testBuildGetAllRuleSearchRequest() { - RulePersistenceService rulePersistenceService = setUpRulePersistenceService(); - Client client = rulePersistenceService.getClient(); - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - - when(client.prepareSearch(anyString())).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.setQuery(any(BoolQueryBuilder.class))).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST)).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.addSort(eq(_ID_STRING), eq(SortOrder.ASC))).thenReturn(searchRequestBuilder); - - rulePersistenceService.buildGetAllRuleSearchRequest(ruleOne.getAttributeMap(), _ID_ONE); - ArgumentCaptor captor = ArgumentCaptor.forClass(Object[].class); - verify(searchRequestBuilder).searchAfter(captor.capture()); - assertEquals(_ID_ONE, captor.getValue()[0]); - } - - public void testBuildGetAllRuleSearchRequest_noSearchAfter() { - RulePersistenceService rulePersistenceService = setUpRulePersistenceService(); - Client client = rulePersistenceService.getClient(); - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - - when(client.prepareSearch(anyString())).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.setQuery(any(BoolQueryBuilder.class))).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST)).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.addSort(eq(_ID_STRING), eq(SortOrder.ASC))).thenReturn(searchRequestBuilder); - - rulePersistenceService.buildGetAllRuleSearchRequest(ruleOne.getAttributeMap(), null); - verify(searchRequestBuilder, times(0)).searchAfter(any()); - } } From cdf8e194966f658240d8636975d1e1b6166d7e07 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 25 Mar 2025 14:43:51 -0700 Subject: [PATCH 04/20] modify based on comments Signed-off-by: Ruirui Zhang --- .../plugin/wlm/rule/QueryGroupAttribute.java | 11 + .../wlm/rule/QueryGroupFeatureType.java | 10 + .../plugin/wlm/rule/action/GetRuleAction.java | 9 + .../wlm/rule/rest/RestGetRuleAction.java | 3 + .../rule/service/RulePersistenceService.java | 220 ++++++------------ .../service/RulePersistenceServiceTests.java | 101 ++++---- 6 files changed, 159 insertions(+), 195 deletions(-) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java index f57a92bce5c9d..5357a344da407 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java @@ -15,8 +15,12 @@ /** * Attributes specific to the query group feature. + * @opensearch.experimental */ public enum QueryGroupAttribute implements Attribute { + /** + * Represents the index_pattern attribute in QueryGroupAttribute + */ INDEX_PATTERN("index_pattern"); private final String name; @@ -31,6 +35,10 @@ public String getName() { return name; } + /** + * Retrieves the QueryGroupAttribute from a name string + * @param name - attribute name + */ public static QueryGroupAttribute fromName(String name) { for (QueryGroupAttribute attr : QueryGroupAttribute.values()) { if (attr.getName().equals(name)) { @@ -40,6 +48,9 @@ public static QueryGroupAttribute fromName(String name) { throw new IllegalArgumentException("Unknown QueryGroupAttribute: " + name); } + /** + * Converts the QueryGroupAttribute values into a map with attribute names as keys. + */ public static Map toMap() { Map map = new HashMap<>(); for (QueryGroupAttribute attr : QueryGroupAttribute.values()) { diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java index feb762b8ad99b..42f45d3d7c5d6 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java @@ -14,8 +14,18 @@ import java.util.Map; +/** + * Represents a feature type specific to the query group feature + * @opensearch.experimental + */ public class QueryGroupFeatureType implements FeatureType { + /** + * The instance for QueryGroupFeatureType + */ public static final QueryGroupFeatureType INSTANCE = new QueryGroupFeatureType(); + /** + * Name for QueryGroupFeatureType + */ public static final String NAME = "query_group"; private static final int MAX_ATTRIBUTE_VALUES = 10; private static final int MAX_ATTRIBUTE_VALUE_LENGTH = 100; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java index 2ee030576c18e..ea53653d2fd67 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java @@ -16,10 +16,19 @@ */ public class GetRuleAction extends ActionType { + /** + * An instance of GetRuleAction + */ public static final GetRuleAction INSTANCE = new GetRuleAction(); + /** + * Name for GetRuleAction + */ public static final String NAME = "cluster:admin/opensearch/wlm/rule/_get"; + /** + * Default constructor for GetRuleAction + */ private GetRuleAction() { super(NAME, GetRuleResponse::new); } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java index c3a9844fc6f33..41dd780b464a5 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java @@ -39,6 +39,9 @@ * @opensearch.experimental */ public class RestGetRuleAction extends BaseRestHandler { + /** + * Field name used for search pagination with the search_after mechanism + */ public static final String SEARCH_AFTER_STRING = "search_after"; /** diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java index a903f80b2249e..ada07ec4705f7 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java @@ -11,7 +11,6 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.ResourceNotFoundException; -import org.opensearch.action.get.GetResponse; import org.opensearch.action.search.SearchRequestBuilder; import org.opensearch.action.search.SearchResponse; import org.opensearch.autotagging.Attribute; @@ -19,6 +18,7 @@ import org.opensearch.autotagging.Rule.Builder; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; +import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; @@ -36,7 +36,6 @@ import java.io.IOException; import java.util.Arrays; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; @@ -50,12 +49,23 @@ * @opensearch.experimental */ public class RulePersistenceService { + /** + * The system index name used for storing rules + */ public static final String RULES_INDEX = ".rules"; private final Client client; private final ClusterService clusterService; private static final Logger logger = LogManager.getLogger(RulePersistenceService.class); + /** + * The maximum number of results allowed per GET request + */ public static final int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 50; + /** + * Constructor for RulePersistenceService + * @param clusterService {@link ClusterService} - The cluster service to be used by RulePersistenceService + * @param client {@link Settings} - The client to be used by RulePersistenceService + */ @Inject public RulePersistenceService(final ClusterService clusterService, final Client client) { this.clusterService = clusterService; @@ -63,8 +73,9 @@ public RulePersistenceService(final ClusterService clusterService, final Client } /** - * Entry point for the get rule api logic in persistence service. - * @param id - The id of the rule to get. Get all matching rules when id is null + * Entry point for the get rule api logic in persistence service. If id is provided, we only get a single rule. + * Otherwise, we get all rules that satisfy the attributeFilters. + * @param id - The id of the rule to get. * @param attributeFilters - A map containing the attributes that user want to filter on * @param searchAfter - The sort values from the last document of the previous page, used for pagination * @param listener - ActionListener for GetRuleResponse @@ -75,100 +86,18 @@ public void getRule( String searchAfter, ActionListener listener ) { - if (id != null) { - fetchRuleById(id, listener); - } else { - fetchAllRules(attributeFilters, searchAfter, listener); - } - } - - /** - * Fetch a single rule from system index using id - * @param id - The id of the rule to get - * @param listener - ActionListener for GetRuleResponse - */ - void fetchRuleById(String id, ActionListener listener) { - try (ThreadContext.StoredContext context = getContext()) { - client.prepareGet(RULES_INDEX, id) - .execute(ActionListener.wrap(getResponse -> handleGetOneRuleResponse(id, getResponse, listener), e -> { - logger.error("Failed to fetch rule with ID {}: {}", id, e.getMessage()); - listener.onFailure(e); - })); - } - } - - /** - * Process getResponse from index and send a GetRuleResponse - * @param id - The id of the rule to get - * @param getResponse - Response received from index - * @param listener - ActionListener for GetRuleResponse - */ - private void handleGetOneRuleResponse(String id, GetResponse getResponse, ActionListener listener) { - if (getResponse.isExists()) { - try (ThreadContext.StoredContext context = getContext()) { - XContentParser parser = MediaTypeRegistry.JSON.xContent() - .createParser( - NamedXContentRegistry.EMPTY, - DeprecationHandler.THROW_UNSUPPORTED_OPERATION, - getResponse.getSourceAsString() - ); - listener.onResponse( - new GetRuleResponse( - Map.of(id, Builder.fromXContent(parser, QueryGroupFeatureType.INSTANCE).build()), - null, - RestStatus.OK - ) - ); - } catch (IOException e) { - logger.error("Error parsing rule with ID {}: {}", id, e.getMessage()); - listener.onFailure(e); - } - } else { - listener.onFailure(new ResourceNotFoundException("Rule with ID " + id + " not found.")); - } - } - - /** - * Fetch all rule from system index based on attributeFilters. - * @param attributeFilters - A map containing the attributes that user want to filter on - * @param searchAfter - The sort values from the last document of the previous page, used for pagination - * @param listener - ActionListener for GetRuleResponse - */ - private void fetchAllRules(Map> attributeFilters, String searchAfter, ActionListener listener) { + // Stash the current thread context when interacting with system index to perform + // operations as the system itself, bypassing authorization checks. This ensures that + // actions within this block are trusted and executed with system-level privileges. try (ThreadContext.StoredContext context = getContext()) { - client.prepareSearch(RULES_INDEX) - .setSize(0) - .execute( - ActionListener.wrap(countResponse -> handleCountResponse(countResponse, attributeFilters, searchAfter, listener), e -> { - logger.error("Failed to check if index is empty: {}", e.getMessage()); - listener.onFailure(e); - }) - ); - } - } - - /** - * Processes the count response from a search query on the rules index. - * If no rules exist, it responds with an empty result. - * Otherwise, it constructs and executes a search request to retrieve all rules. - * @param countResponse The response from the count query on the rules index. - * @param attributeFilters A map of attribute filters to apply in the search query. - * @param searchAfter The searchAfter parameter for pagination. - * @param listener The action listener to handle the final response or failure. - */ - void handleCountResponse( - SearchResponse countResponse, - Map> attributeFilters, - String searchAfter, - ActionListener listener - ) { - try (ThreadContext.StoredContext context = getContext()) { - if (countResponse.getHits().getTotalHits().value() == 0) { - listener.onResponse(new GetRuleResponse(new HashMap<>(), null, RestStatus.OK)); - return; + BoolQueryBuilder boolQuery = buildGetRuleQuery(id, attributeFilters); + SearchRequestBuilder searchRequest = client.prepareSearch(RULES_INDEX) + .setQuery(boolQuery) + .setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST); + if (searchAfter != null) { + searchRequest.addSort(_ID_STRING, SortOrder.ASC).searchAfter(new Object[] { searchAfter }); } - SearchRequestBuilder searchRequest = buildGetAllRuleSearchRequest(attributeFilters, searchAfter); - searchRequest.execute(ActionListener.wrap(searchResponse -> handleGetAllRuleResponse(searchResponse, listener), e -> { + searchRequest.execute(ActionListener.wrap(searchResponse -> handleGetRuleResponse(id, searchResponse, listener), e -> { logger.error("Failed to fetch all rules: {}", e.getMessage()); listener.onFailure(e); })); @@ -176,35 +105,29 @@ void handleCountResponse( } /** - * Builds a search request to retrieve all rules from the rules index, applying attribute-based filters - * and ensuring that the rules are associated with the query group feature type. + * Builds a bool query to retrieve rules from the rules index, applying attribute-based filters + * when needed and ensuring that the rules are associated with the query group feature type. + * @param id The ID of the rule to fetch. If not null, the search will return only this specific rule. * @param attributeFilters A map of attributes to their associated set of values used to filter the rules. - * @param searchAfter A cursor to enable pagination, used to fetch results after a specific document. */ - SearchRequestBuilder buildGetAllRuleSearchRequest(Map> attributeFilters, String searchAfter) { - try (ThreadContext.StoredContext context = getContext()) { - BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); - for (Map.Entry> entry : attributeFilters.entrySet()) { - Attribute attribute = entry.getKey(); - Set values = entry.getValue(); - if (values != null && !values.isEmpty()) { - BoolQueryBuilder attributeQuery = QueryBuilders.boolQuery(); - for (String value : values) { - attributeQuery.should(QueryBuilders.matchQuery(attribute.getName(), value)); - } - boolQuery.must(attributeQuery); + BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeFilters) { + BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); + if (id != null) { + return boolQuery.must(QueryBuilders.termQuery(_ID_STRING, id)); + } + for (Map.Entry> entry : attributeFilters.entrySet()) { + Attribute attribute = entry.getKey(); + Set values = entry.getValue(); + if (values != null && !values.isEmpty()) { + BoolQueryBuilder attributeQuery = QueryBuilders.boolQuery(); + for (String value : values) { + attributeQuery.should(QueryBuilders.matchQuery(attribute.getName(), value)); } + boolQuery.must(attributeQuery); } - boolQuery.filter(QueryBuilders.existsQuery(QueryGroupFeatureType.NAME)); - SearchRequestBuilder searchRequest = client.prepareSearch(RULES_INDEX) - .setQuery(boolQuery) - .setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST) - .addSort(_ID_STRING, SortOrder.ASC); - if (searchAfter != null) { - searchRequest.searchAfter(new Object[] { searchAfter }); - } - return searchRequest; } + boolQuery.filter(QueryBuilders.existsQuery(QueryGroupFeatureType.NAME)); + return boolQuery; } /** @@ -212,41 +135,52 @@ SearchRequestBuilder buildGetAllRuleSearchRequest(Map> at * @param searchResponse - Response received from index * @param listener - ActionListener for GetRuleResponse */ - void handleGetAllRuleResponse(SearchResponse searchResponse, ActionListener listener) { + void handleGetRuleResponse(String id, SearchResponse searchResponse, ActionListener listener) { List hits = Arrays.asList(searchResponse.getHits().getHits()); - try (ThreadContext.StoredContext context = getContext()) { - Map ruleMap = hits.stream().map(hit -> { - String hitId = hit.getId(); - try { - XContentParser parser = MediaTypeRegistry.JSON.xContent() - .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, hit.getSourceAsString()); - return Map.entry(hitId, Builder.fromXContent(parser, QueryGroupFeatureType.INSTANCE).build()); - } catch (IOException e) { - logger.info( - "Issue met when parsing rule from hit, the feature type for rule id {} is probably not query_group: {}", - hitId, - e.getMessage() - ); - return null; - } - }).filter(Objects::nonNull).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - String nextSearchAfter = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); - listener.onResponse(new GetRuleResponse(ruleMap, nextSearchAfter, RestStatus.OK)); + if (id != null && hits.isEmpty()) { + logger.error("Rule with ID " + id + " not found."); + listener.onFailure(new ResourceNotFoundException("Rule with ID " + id + " doesn't exist in the .rules index.")); + return; } + Map ruleMap = hits.stream() + .map(hit -> parseRule(hit.getId(), hit.getSourceAsString())) + .filter(Objects::nonNull) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + String nextSearchAfter = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); + listener.onResponse(new GetRuleResponse(ruleMap, nextSearchAfter, RestStatus.OK)); } - private ThreadContext.StoredContext getContext() { - return client.threadPool().getThreadContext().stashContext(); + /** + * Parses a source string into a Rule object + * @param id - document id for the Rule object + * @param source - The raw source string representing the rule to be parsed + */ + private Map.Entry parseRule(String id, String source) { + try ( + XContentParser parser = MediaTypeRegistry.JSON.xContent() + .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, source) + ) { + return Map.entry(id, Builder.fromXContent(parser, QueryGroupFeatureType.INSTANCE).build()); + } catch (IOException e) { + logger.info("Issue met when parsing rule for ID {}: {}", id, e.getMessage()); + return null; + } } - private boolean isExistingQueryGroup(String queryGroupId) { - return clusterService.state().metadata().queryGroups().containsKey(queryGroupId); + private ThreadContext.StoredContext getContext() { + return client.threadPool().getThreadContext().stashContext(); } + /** + * client getter + */ public Client getClient() { return client; } + /** + * clusterService getter + */ public ClusterService getClusterService() { return clusterService; } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java index d995cf19a3bd7..6bfa8b96cf16e 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java @@ -8,30 +8,27 @@ package org.opensearch.plugin.wlm.rule.service; -import org.opensearch.ResourceNotFoundException; -import org.opensearch.action.get.GetRequestBuilder; -import org.opensearch.action.get.GetResponse; -import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.action.search.SearchRequestBuilder; +import org.opensearch.action.search.SearchResponse; import org.opensearch.core.action.ActionListener; -import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; +import org.opensearch.search.sort.SortOrder; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.transport.client.Client; -import java.io.IOException; import java.util.HashMap; -import java.util.Map; - -import org.mockito.ArgumentCaptor; +import static org.opensearch.autotagging.Rule._ID_STRING; +import static org.opensearch.plugin.wlm.RuleTestUtils.ATTRIBUTE_MAP; import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; -import static org.opensearch.plugin.wlm.RuleTestUtils.assertEqualRules; -import static org.opensearch.plugin.wlm.RuleTestUtils.ruleOne; +import static org.opensearch.plugin.wlm.RuleTestUtils._ID_TWO; import static org.opensearch.plugin.wlm.RuleTestUtils.setUpRulePersistenceService; -import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.RULES_INDEX; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.clearInvocations; +import static org.mockito.Mockito.anyInt; +import static org.mockito.Mockito.anyString; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -39,56 +36,56 @@ @SuppressWarnings("unchecked") public class RulePersistenceServiceTests extends OpenSearchTestCase { - public void testGetRuleById() throws IOException { - String ruleSource = ruleOne.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString(); - ActionListener listener = mock(ActionListener.class); + public void testBuildGetRuleQuery_WithId() { RulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); - Client client = rulePersistenceService.getClient(); - GetRequestBuilder getRequestBuilder = mock(GetRequestBuilder.class); - GetResponse getResponse = mock(GetResponse.class); + BoolQueryBuilder query = rulePersistenceService.buildGetRuleQuery(_ID_ONE, new HashMap<>()); + assertTrue(query.hasClauses()); + assertEquals(QueryBuilders.termQuery(_ID_STRING, _ID_ONE).toString(), query.must().get(0).toString()); + } - when(getResponse.isExists()).thenReturn(true); - when(getResponse.getSourceAsString()).thenReturn(ruleSource); - when(client.prepareGet(eq(RULES_INDEX), eq(_ID_ONE))).thenReturn(getRequestBuilder); - doAnswer(invocation -> { - ActionListener actionListener = invocation.getArgument(0); - actionListener.onResponse(getResponse); - return null; - }).when(getRequestBuilder).execute(any(ActionListener.class)); + public void testBuildGetRuleQuery_WithFilters() { + RulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); + BoolQueryBuilder query = rulePersistenceService.buildGetRuleQuery(null, ATTRIBUTE_MAP); + assertTrue(query.hasClauses()); + assertEquals(1, query.must().size()); + assertTrue(query.filter().contains(QueryBuilders.existsQuery(QueryGroupFeatureType.NAME))); + } - rulePersistenceService.getRule(_ID_ONE, new HashMap<>(), null, listener); + public void testGetRule_WithId() { + RulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); + Client client = rulePersistenceService.getClient(); + ActionListener listener = mock(ActionListener.class); + SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + SetupMocksForGetRule(client, searchRequestBuilder); - ArgumentCaptor captor = ArgumentCaptor.forClass(GetRuleResponse.class); - verify(listener).onResponse(captor.capture()); - GetRuleResponse response = captor.getValue(); - assertNotNull(response); - assertEqualRules(Map.of(_ID_ONE, ruleOne), response.getRules(), false); - clearInvocations(client, getRequestBuilder, getResponse, listener); + rulePersistenceService.getRule(_ID_ONE, new HashMap<>(), null, listener); + verify(client).prepareSearch(RulePersistenceService.RULES_INDEX); + verify(searchRequestBuilder).setQuery(any()); + verify(searchRequestBuilder).execute(any()); } - public void testGetRuleByIdNotFound() { - String nonExistentRuleId = "non-existent-rule"; + public void testGetRule_WithSearchAfter() { RulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); Client client = rulePersistenceService.getClient(); - GetRequestBuilder getRequestBuilder = mock(GetRequestBuilder.class); - GetResponse getResponse = mock(GetResponse.class); ActionListener listener = mock(ActionListener.class); + SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + SetupMocksForGetRule(client, searchRequestBuilder); + when(searchRequestBuilder.addSort(anyString(), any(SortOrder.class))).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.searchAfter(any())).thenReturn(searchRequestBuilder); - when(client.prepareGet(RULES_INDEX, nonExistentRuleId)).thenReturn(getRequestBuilder); - when(getResponse.isExists()).thenReturn(false); + rulePersistenceService.getRule(null, new HashMap<>(), _ID_TWO, listener); + verify(searchRequestBuilder).addSort(_ID_STRING, SortOrder.ASC); + verify(searchRequestBuilder).searchAfter(new Object[] { _ID_TWO }); + } + public void SetupMocksForGetRule(Client client, SearchRequestBuilder searchRequestBuilder) { + when(client.prepareSearch(anyString())).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.setQuery(any())).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.setSize(anyInt())).thenReturn(searchRequestBuilder); doAnswer(invocation -> { - ActionListener actionListener = invocation.getArgument(0); - actionListener.onResponse(getResponse); + ActionListener searchListener = invocation.getArgument(0); + searchListener.onResponse(mock(SearchResponse.class)); return null; - }).when(getRequestBuilder).execute(any(ActionListener.class)); - - rulePersistenceService.getRule(nonExistentRuleId, new HashMap<>(), null, listener); - - ArgumentCaptor captor = ArgumentCaptor.forClass(Exception.class); - verify(listener).onFailure(captor.capture()); - Exception exception = captor.getValue(); - assertTrue(exception instanceof ResourceNotFoundException); - clearInvocations(client, getRequestBuilder, getResponse, listener); + }).when(searchRequestBuilder).execute(any()); } } From 28ca8b1a35974a026c94f403118acb2589483550 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Sun, 30 Mar 2025 23:30:53 -0700 Subject: [PATCH 05/20] extract common logics to libs Signed-off-by: Ruirui Zhang --- .../rule/action/GetRuleRequest.java | 31 +++------ .../rule/action/GetRuleResponse.java | 26 +------- .../opensearch/rule/action/package-info.java | 12 ++++ .../rule/rest/RestGetRuleAction.java | 44 ++++++------- .../opensearch/rule/rest/package-info.java | 12 ++++ .../rule/service/RulePersistenceService.java | 63 ++++++++++-------- .../opensearch/rule/service/package-info.java | 12 ++++ ...tRuleAction.java => GetWlmRuleAction.java} | 14 ++-- .../wlm/rule/action/GetWlmRuleRequest.java | 52 +++++++++++++++ .../wlm/rule/action/GetWlmRuleResponse.java | 55 ++++++++++++++++ ...on.java => TransportGetWlmRuleAction.java} | 20 +++--- .../wlm/rule/rest/RestGetWlmRuleAction.java | 64 +++++++++++++++++++ .../service/WlmRulePersistenceService.java | 50 +++++++++++++++ .../plugin/wlm/WorkloadGroupTestUtils.java | 2 +- .../RestDeleteWorkloadGroupActionTests.java | 5 -- .../WlmRuleTestUtils.java} | 12 ++-- ...Tests.java => GetWlmRuleRequestTests.java} | 16 ++--- ...ests.java => GetWlmRuleResponseTests.java} | 28 ++++---- ...ts.java => RestGetWlmRuleActionTests.java} | 5 +- ...va => WlmRulePersistenceServiceTests.java} | 26 ++++---- 20 files changed, 386 insertions(+), 163 deletions(-) rename {plugins/workload-management/src/main/java/org/opensearch/plugin/wlm => libs/autotagging-commons/src/main/java/org/opensearch}/rule/action/GetRuleRequest.java (71%) rename {plugins/workload-management/src/main/java/org/opensearch/plugin/wlm => libs/autotagging-commons/src/main/java/org/opensearch}/rule/action/GetRuleResponse.java (77%) create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/action/package-info.java rename {plugins/workload-management/src/main/java/org/opensearch/plugin/wlm => libs/autotagging-commons/src/main/java/org/opensearch}/rule/rest/RestGetRuleAction.java (65%) create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/package-info.java rename {plugins/workload-management/src/main/java/org/opensearch/plugin/wlm => libs/autotagging-commons/src/main/java/org/opensearch}/rule/service/RulePersistenceService.java (78%) create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/package-info.java rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/{GetRuleAction.java => GetWlmRuleAction.java} (60%) create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponse.java rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/{TransportGetRuleAction.java => TransportGetWlmRuleAction.java} (58%) create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{RuleTestUtils.java => rule/WlmRuleTestUtils.java} (90%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/{GetRuleRequestTests.java => GetWlmRuleRequestTests.java} (71%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/{GetRuleResponseTests.java => GetWlmRuleResponseTests.java} (78%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/{RestGetRuleActionTests.java => RestGetWlmRuleActionTests.java} (88%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/{RulePersistenceServiceTests.java => WlmRulePersistenceServiceTests.java} (76%) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java similarity index 71% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java rename to libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java index 181bd9583ce34..6cae3e860c3ad 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequest.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java @@ -6,14 +6,14 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.rule.action; +package org.opensearch.rule.action; import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.FeatureType; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import java.io.IOException; import java.util.HashSet; @@ -23,36 +23,27 @@ /** * A request for get Rule * Example Request: + * The endpoint "localhost:9200/_wlm/rule" is specific to the Workload Management feature to manage rules * curl -X GET "localhost:9200/_wlm/rule" - get all rules * curl -X GET "localhost:9200/_wlm/rule/{_id}" - get single rule by id - * curl -X GET "localhost:9200/_wlm/rule?index_pattern=a,b" - get all rules containing index_pattern as a or b + * curl -X GET "localhost:9200/_wlm/rule?index_pattern=a,b" - get all rules containing attribute index_pattern as a or b * @opensearch.experimental */ -public class GetRuleRequest extends ActionRequest { +public abstract class GetRuleRequest extends ActionRequest { private final String id; private final Map> attributeFilters; private final String searchAfter; - /** - * Constructor for GetRuleRequest - * @param id - Rule id that we want to get - * @param attributeFilters - Attributes that we want to filter on - * @param searchAfter - The sort values from the last document of the previous page, used for pagination - */ public GetRuleRequest(String id, Map> attributeFilters, String searchAfter) { this.id = id; this.attributeFilters = attributeFilters; this.searchAfter = searchAfter; } - /** - * Constructor for GetRuleRequest - * @param in - A {@link StreamInput} object - */ public GetRuleRequest(StreamInput in) throws IOException { super(in); id = in.readOptionalString(); - attributeFilters = in.readMap(i -> Attribute.from(i, QueryGroupFeatureType.INSTANCE), i -> new HashSet<>(i.readStringList())); + attributeFilters = in.readMap(i -> Attribute.from(i, retrieveFeatureTypeInstance()), i -> new HashSet<>(i.readStringList())); searchAfter = in.readOptionalString(); } @@ -70,22 +61,18 @@ public void writeTo(StreamOutput out) throws IOException { } /** - * id getter + * Abstract method for subclasses to provide specific FeatureType Instance */ + protected abstract FeatureType retrieveFeatureTypeInstance(); + public String getId() { return id; } - /** - * attributeFilters getter - */ public Map> getAttributeFilters() { return attributeFilters; } - /** - * searchAfter getter - */ public String getSearchAfter() { return searchAfter; } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java similarity index 77% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java rename to libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java index 2f45e372b8325..a92d3c4d79df8 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponse.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.rule.action; +package org.opensearch.rule.action; import org.opensearch.autotagging.Rule; import org.opensearch.core.action.ActionResponse; @@ -21,7 +21,6 @@ import java.util.Map; import static org.opensearch.autotagging.Rule._ID_STRING; -import static org.opensearch.plugin.wlm.rule.rest.RestGetRuleAction.SEARCH_AFTER_STRING; /** * Response for the get API for Rule. @@ -41,27 +40,17 @@ * } * @opensearch.experimental */ -public class GetRuleResponse extends ActionResponse implements ToXContent, ToXContentObject { +public abstract class GetRuleResponse extends ActionResponse implements ToXContent, ToXContentObject { private final Map rules; private final String searchAfter; private final RestStatus restStatus; - /** - * Constructor for GetRuleResponse - * @param rules - The Map of Rules to be included in the response - * @param searchAfter - The searchAfter field for the response - * @param restStatus - The restStatus for the response - */ public GetRuleResponse(final Map rules, String searchAfter, RestStatus restStatus) { this.rules = rules; this.searchAfter = searchAfter; this.restStatus = restStatus; } - /** - * Constructor for GetRuleResponse - * @param in - A {@link StreamInput} object - */ public GetRuleResponse(StreamInput in) throws IOException { this(in.readMap(StreamInput::readString, Rule::new), in.readOptionalString(), RestStatus.readFrom(in)); } @@ -82,29 +71,20 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws } builder.endArray(); if (searchAfter != null && !searchAfter.isEmpty()) { - builder.field(SEARCH_AFTER_STRING, new Object[] { searchAfter }); + builder.field("search_after", new Object[] { searchAfter }); } builder.endObject(); return builder; } - /** - * rules getter - */ public Map getRules() { return rules; } - /** - * restStatus getter - */ public RestStatus getRestStatus() { return restStatus; } - /** - * searchAfter getter - */ public String getSearchAfter() { return searchAfter; } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/package-info.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/package-info.java new file mode 100644 index 0000000000000..91913aff23eac --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/package-info.java @@ -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. + */ + +/** + * This package contains abstract action classes for rules + */ +package org.opensearch.rule.action; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java similarity index 65% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java rename to libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java index 41dd780b464a5..686b4cbf2406e 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleAction.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java @@ -6,21 +6,20 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.rule.rest; +package org.opensearch.rule.rest; +import org.opensearch.action.ActionType; import org.opensearch.autotagging.Attribute; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; -import org.opensearch.plugin.wlm.rule.QueryGroupAttribute; -import org.opensearch.plugin.wlm.rule.action.GetRuleAction; -import org.opensearch.plugin.wlm.rule.action.GetRuleRequest; -import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; 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 org.opensearch.rule.action.GetRuleRequest; +import org.opensearch.rule.action.GetRuleResponse; import org.opensearch.transport.client.node.NodeClient; import java.io.IOException; @@ -32,35 +31,21 @@ import java.util.Set; import static org.opensearch.autotagging.Rule._ID_STRING; -import static org.opensearch.rest.RestRequest.Method.GET; /** * Rest action to get a Rule * @opensearch.experimental */ -public class RestGetRuleAction extends BaseRestHandler { - /** - * Field name used for search pagination with the search_after mechanism - */ +public abstract class RestGetRuleAction extends BaseRestHandler { public static final String SEARCH_AFTER_STRING = "search_after"; - /** - * Constructor for RestGetRuleAction - */ public RestGetRuleAction() {} @Override - public String getName() { - return "get_rule"; - } + public abstract String getName(); - /** - * The list of {@link Route}s that this RestHandler is responsible for handling. - */ @Override - public List routes() { - return List.of(new Route(GET, "_wlm/rule/"), new Route(GET, "_wlm/rule/{_id}")); - } + public abstract List routes(); @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { @@ -70,14 +55,14 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli continue; } String[] valuesArray = request.param(attributeName).split(","); - attributeFilters.put(QueryGroupAttribute.fromName(attributeName), new HashSet<>(Arrays.asList(valuesArray))); + attributeFilters.put(getAttributeFromName(attributeName), new HashSet<>(Arrays.asList(valuesArray))); } - final GetRuleRequest getRuleRequest = new GetRuleRequest( + final GetRuleRequest getRuleRequest = buildGetRuleRequest( request.param(_ID_STRING), attributeFilters, request.param(SEARCH_AFTER_STRING) ); - return channel -> client.execute(GetRuleAction.INSTANCE, getRuleRequest, getRuleResponse(channel)); + return channel -> client.execute(retrieveGetRuleActionInstance(), getRuleRequest, getRuleResponse(channel)); } private RestResponseListener getRuleResponse(final RestChannel channel) { @@ -88,4 +73,13 @@ public RestResponse buildResponse(final GetRuleResponse response) throws Excepti } }; } + + protected abstract Attribute getAttributeFromName(String name); + + /** + * Abstract method for subclasses to provide specific ActionType Instance + */ + protected abstract > T retrieveGetRuleActionInstance(); + + protected abstract GetRuleRequest buildGetRuleRequest(String id, Map> attributeFilters, String searchAfter); } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/package-info.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/package-info.java new file mode 100644 index 0000000000000..c1000b90b1856 --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/package-info.java @@ -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. + */ + +/** + * This package contains abstract rest classes for rules + */ +package org.opensearch.rule.rest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java similarity index 78% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java rename to libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java index ada07ec4705f7..69c957c94bf3d 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceService.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java @@ -6,40 +6,46 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.rule.service; +package org.opensearch.rule.service; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.ResourceAlreadyExistsException; import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.admin.indices.create.CreateIndexRequest; +import org.opensearch.action.admin.indices.create.CreateIndexResponse; +import org.opensearch.action.index.IndexRequest; import org.opensearch.action.search.SearchRequestBuilder; import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.update.UpdateRequest; import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.FeatureType; import org.opensearch.autotagging.Rule; import org.opensearch.autotagging.Rule.Builder; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.action.ActionListener; +import org.opensearch.core.action.ActionResponse; import org.opensearch.core.rest.RestStatus; -import org.opensearch.core.xcontent.DeprecationHandler; -import org.opensearch.core.xcontent.MediaTypeRegistry; -import org.opensearch.core.xcontent.NamedXContentRegistry; -import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.core.xcontent.*; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilders; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; +import org.opensearch.rule.action.GetRuleResponse; import org.opensearch.search.SearchHit; import org.opensearch.search.sort.SortOrder; import org.opensearch.transport.client.Client; import java.io.IOException; +import java.time.Instant; import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; +import java.util.HashMap; import java.util.stream.Collectors; import static org.opensearch.autotagging.Rule._ID_STRING; @@ -48,7 +54,7 @@ * This class encapsulates the logic to manage the lifecycle of rules at index level * @opensearch.experimental */ -public class RulePersistenceService { +public abstract class RulePersistenceService { /** * The system index name used for storing rules */ @@ -56,16 +62,12 @@ public class RulePersistenceService { private final Client client; private final ClusterService clusterService; private static final Logger logger = LogManager.getLogger(RulePersistenceService.class); + /** * The maximum number of results allowed per GET request */ public static final int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 50; - /** - * Constructor for RulePersistenceService - * @param clusterService {@link ClusterService} - The cluster service to be used by RulePersistenceService - * @param client {@link Settings} - The client to be used by RulePersistenceService - */ @Inject public RulePersistenceService(final ClusterService clusterService, final Client client) { this.clusterService = clusterService; @@ -84,7 +86,7 @@ public void getRule( String id, Map> attributeFilters, String searchAfter, - ActionListener listener + ActionListener listener ) { // Stash the current thread context when interacting with system index to perform // operations as the system itself, bypassing authorization checks. This ensures that @@ -93,7 +95,7 @@ public void getRule( BoolQueryBuilder boolQuery = buildGetRuleQuery(id, attributeFilters); SearchRequestBuilder searchRequest = client.prepareSearch(RULES_INDEX) .setQuery(boolQuery) - .setSize(MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST); + .setSize(getMaxReturnSizeAllowedPerGetRequest()); if (searchAfter != null) { searchRequest.addSort(_ID_STRING, SortOrder.ASC).searchAfter(new Object[] { searchAfter }); } @@ -110,7 +112,7 @@ public void getRule( * @param id The ID of the rule to fetch. If not null, the search will return only this specific rule. * @param attributeFilters A map of attributes to their associated set of values used to filter the rules. */ - BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeFilters) { + public BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeFilters) { BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); if (id != null) { return boolQuery.must(QueryBuilders.termQuery(_ID_STRING, id)); @@ -126,7 +128,7 @@ BoolQueryBuilder buildGetRuleQuery(String id, Map> attrib boolQuery.must(attributeQuery); } } - boolQuery.filter(QueryBuilders.existsQuery(QueryGroupFeatureType.NAME)); + boolQuery.filter(QueryBuilders.existsQuery(retrieveFeatureTypeInstance().getName())); return boolQuery; } @@ -135,7 +137,7 @@ BoolQueryBuilder buildGetRuleQuery(String id, Map> attrib * @param searchResponse - Response received from index * @param listener - ActionListener for GetRuleResponse */ - void handleGetRuleResponse(String id, SearchResponse searchResponse, ActionListener listener) { + void handleGetRuleResponse(String id, SearchResponse searchResponse, ActionListener listener) { List hits = Arrays.asList(searchResponse.getHits().getHits()); if (id != null && hits.isEmpty()) { logger.error("Rule with ID " + id + " not found."); @@ -147,7 +149,7 @@ void handleGetRuleResponse(String id, SearchResponse searchResponse, ActionListe .filter(Objects::nonNull) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); String nextSearchAfter = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); - listener.onResponse(new GetRuleResponse(ruleMap, nextSearchAfter, RestStatus.OK)); + listener.onResponse(buildGetRuleResponse(ruleMap, nextSearchAfter, RestStatus.OK)); } /** @@ -160,27 +162,36 @@ private Map.Entry parseRule(String id, String source) { XContentParser parser = MediaTypeRegistry.JSON.xContent() .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, source) ) { - return Map.entry(id, Builder.fromXContent(parser, QueryGroupFeatureType.INSTANCE).build()); + return Map.entry(id, Builder.fromXContent(parser, retrieveFeatureTypeInstance()).build()); } catch (IOException e) { logger.info("Issue met when parsing rule for ID {}: {}", id, e.getMessage()); return null; } } - private ThreadContext.StoredContext getContext() { - return client.threadPool().getThreadContext().stashContext(); + private boolean isExistingQueryGroup(String queryGroupId) { + return clusterService.state().metadata().queryGroups().containsKey(queryGroupId); } /** - * client getter + * Abstract method for subclasses to provide specific FeatureType Instance */ + protected abstract FeatureType retrieveFeatureTypeInstance(); + + protected abstract T buildGetRuleResponse(Map ruleMap, String nextSearchAfter, RestStatus restStatus); + + private ThreadContext.StoredContext getContext() { + return client.threadPool().getThreadContext().stashContext(); + } + + public int getMaxReturnSizeAllowedPerGetRequest() { + return MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST; + } + public Client getClient() { return client; } - /** - * clusterService getter - */ public ClusterService getClusterService() { return clusterService; } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/package-info.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/package-info.java new file mode 100644 index 0000000000000..f51b281884dc6 --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/package-info.java @@ -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. + */ + +/** + * This package contains abstract service classes for rules + */ +package org.opensearch.rule.service; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java similarity index 60% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java index ea53653d2fd67..95ecbe6d6a4ec 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java @@ -14,22 +14,22 @@ * Action type for getting Rules in workload management * @opensearch.experimental */ -public class GetRuleAction extends ActionType { +public class GetWlmRuleAction extends ActionType { /** - * An instance of GetRuleAction + * An instance of GetWlmRuleAction */ - public static final GetRuleAction INSTANCE = new GetRuleAction(); + public static final GetWlmRuleAction INSTANCE = new GetWlmRuleAction(); /** - * Name for GetRuleAction + * Name for GetWlmRuleAction */ public static final String NAME = "cluster:admin/opensearch/wlm/rule/_get"; /** - * Default constructor for GetRuleAction + * Default constructor for GetWlmRuleAction */ - private GetRuleAction() { - super(NAME, GetRuleResponse::new); + private GetWlmRuleAction() { + super(NAME, GetWlmRuleResponse::new); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java new file mode 100644 index 0000000000000..d2b3e6f9c962f --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java @@ -0,0 +1,52 @@ +/* + * 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.rule.action; + +import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.FeatureType; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; +import org.opensearch.rule.action.GetRuleRequest; + +import java.io.IOException; +import java.util.Map; +import java.util.Set; + +/** + * A request to get workload management Rules in workload management + * Example Request: + * curl -X GET "localhost:9200/_wlm/rule" - get all rules + * curl -X GET "localhost:9200/_wlm/rule/{_id}" - get single rule by id + * curl -X GET "localhost:9200/_wlm/rule?index_pattern=a,b" - get all rules containing attribute index_pattern as a or b + * @opensearch.experimental + */ +public class GetWlmRuleRequest extends GetRuleRequest { + /** + * Constructor for GetWlmRuleRequest + * @param id - Rule id to get + * @param attributeFilters - A map containing the attributes to filter on + * @param searchAfter - A string used for pagination + */ + public GetWlmRuleRequest(String id, Map> attributeFilters, String searchAfter) { + super(id, attributeFilters, searchAfter); + } + + /** + * Constructs a new GetWlmRuleRequest instance from StreamInput + * @param in The {@link StreamInput} from which to read the request data. + */ + public GetWlmRuleRequest(StreamInput in) throws IOException { + super(in); + } + + @Override + protected FeatureType retrieveFeatureTypeInstance() { + return QueryGroupFeatureType.INSTANCE; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponse.java new file mode 100644 index 0000000000000..c3ab2a80ecf1a --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponse.java @@ -0,0 +1,55 @@ +/* + * 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.rule.action; + +import org.opensearch.autotagging.Rule; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.rule.action.GetRuleResponse; + +import java.io.IOException; +import java.util.Map; + +/** + * A response to get workload management Rules in workload management + * Example response: + * { + * "rules": [ + * { + * "_id": "z1MJApUB0zgMcDmz-UQq", + * "description": "Rule for tagging query_group_id to index123" + * "index_pattern": ["index123"], + * "query_group": "query_group_id", + * "updated_at": "2025-02-14T01:19:22.589Z" + * }, + * ... + * ], + * "search_after": ["z1MJApUB0zgMcDmz-UQq"] + * } + * @opensearch.experimental + */ +public class GetWlmRuleResponse extends GetRuleResponse { + /** + * Constructor for GetWlmRuleResponse + * @param rules - A map of rule IDs to objects representing the retrieved rules + * @param searchAfter - A string used for pagination + * @param restStatus - The {@link RestStatus} indicating the status of the request. + */ + public GetWlmRuleResponse(Map rules, String searchAfter, RestStatus restStatus) { + super(rules, searchAfter, restStatus); + } + + /** + * Constructs a new GetWlmRuleResponse instance from StreamInput + * @param in The {@link StreamInput} from which to read the response data. + */ + public GetWlmRuleResponse(StreamInput in) throws IOException { + super(in); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java similarity index 58% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java index 7b9a77d78e328..ca20ea4d1815a 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java @@ -12,36 +12,36 @@ import org.opensearch.action.support.HandledTransportAction; import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; -import org.opensearch.plugin.wlm.rule.service.RulePersistenceService; +import org.opensearch.plugin.wlm.rule.service.WlmRulePersistenceService; import org.opensearch.tasks.Task; import org.opensearch.transport.TransportService; /** - * Transport action to get Rule + * Transport action to get workload management Rules * @opensearch.experimental */ -public class TransportGetRuleAction extends HandledTransportAction { +public class TransportGetWlmRuleAction extends HandledTransportAction { - private final RulePersistenceService rulePersistenceService; + private final WlmRulePersistenceService rulePersistenceService; /** - * Constructor for TransportGetRuleAction + * Constructor for TransportGetWlmRuleAction * @param transportService - a {@link TransportService} object * @param actionFilters - a {@link ActionFilters} object - * @param rulePersistenceService - a {@link RulePersistenceService} object + * @param rulePersistenceService - a {@link WlmRulePersistenceService} object */ @Inject - public TransportGetRuleAction( + public TransportGetWlmRuleAction( TransportService transportService, ActionFilters actionFilters, - RulePersistenceService rulePersistenceService + WlmRulePersistenceService rulePersistenceService ) { - super(GetRuleAction.NAME, transportService, actionFilters, GetRuleRequest::new); + super(GetWlmRuleAction.NAME, transportService, actionFilters, GetWlmRuleRequest::new); this.rulePersistenceService = rulePersistenceService; } @Override - protected void doExecute(Task task, GetRuleRequest request, ActionListener listener) { + protected void doExecute(Task task, GetWlmRuleRequest request, ActionListener listener) { rulePersistenceService.getRule(request.getId(), request.getAttributeFilters(), request.getSearchAfter(), listener); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java new file mode 100644 index 0000000000000..11f7776c77c21 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.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.rule.rest; + +import org.opensearch.action.ActionType; +import org.opensearch.autotagging.Attribute; +import org.opensearch.plugin.wlm.rule.QueryGroupAttribute; +import org.opensearch.plugin.wlm.rule.action.GetWlmRuleAction; +import org.opensearch.plugin.wlm.rule.action.GetWlmRuleRequest; +import org.opensearch.rule.action.GetRuleRequest; +import org.opensearch.rule.action.GetRuleResponse; +import org.opensearch.rule.rest.RestGetRuleAction; + +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.opensearch.rest.RestRequest.Method.GET; + +/** + * Rest action to get workload management Rules + * @opensearch.experimental + */ +public class RestGetWlmRuleAction extends RestGetRuleAction { + + /** + * Constructor for RestGetWlmRuleAction + */ + public RestGetWlmRuleAction() { + super(); + } + + @Override + public String getName() { + return "get_rule"; + } + + @Override + public List routes() { + return List.of(new Route(GET, "_wlm/rule/"), new Route(GET, "_wlm/rule/{_id}")); + } + + @Override + protected Attribute getAttributeFromName(String name) { + return QueryGroupAttribute.fromName(name); + } + + @Override + @SuppressWarnings("unchecked") + protected > T retrieveGetRuleActionInstance() { + return (T) GetWlmRuleAction.INSTANCE; + } + + @Override + protected GetRuleRequest buildGetRuleRequest(String id, Map> attributeFilters, String searchAfter) { + return new GetWlmRuleRequest(id, attributeFilters, searchAfter); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java new file mode 100644 index 0000000000000..5ee6a1d65dbb0 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java @@ -0,0 +1,50 @@ +/* + * 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.rule.service; + +import org.opensearch.autotagging.FeatureType; +import org.opensearch.autotagging.Rule; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; +import org.opensearch.plugin.wlm.rule.action.GetWlmRuleResponse; +import org.opensearch.rule.action.GetRuleResponse; +import org.opensearch.rule.service.RulePersistenceService; +import org.opensearch.transport.client.Client; + +import java.util.Map; + +/** + * This class encapsulates the logic to manage the lifecycle of workload management rules at index level + * @opensearch.experimental + */ +@SuppressWarnings("unchecked") +public class WlmRulePersistenceService extends RulePersistenceService { + /** + * Constructor for WlmRulePersistenceService + * @param clusterService {@link ClusterService} - The cluster service to be used by RulePersistenceService + * @param client {@link Settings} - The client to be used by RulePersistenceService + */ + @Inject + public WlmRulePersistenceService(ClusterService clusterService, Client client) { + super(clusterService, client); + } + + @Override + protected FeatureType retrieveFeatureTypeInstance() { + return QueryGroupFeatureType.INSTANCE; + } + + @Override + protected T buildGetRuleResponse(Map ruleMap, String nextSearchAfter, RestStatus restStatus) { + return (T) new GetWlmRuleResponse(ruleMap, nextSearchAfter, restStatus); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadGroupTestUtils.java index bac644a172c1e..dac26c3c0f929 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadGroupTestUtils.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm; +package org.opensearch.plugin.wlm.querygroup; import org.opensearch.cluster.ClusterName; import org.opensearch.cluster.ClusterState; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java index 4269c6d64ab97..df59ba3103db5 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java @@ -11,13 +11,8 @@ import org.opensearch.action.support.clustermanager.AcknowledgedResponse; import org.opensearch.common.CheckedConsumer; import org.opensearch.common.unit.TimeValue; -<<<<<<<< HEAD:plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupAction; import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupRequest; -======== -import org.opensearch.plugin.wlm.querygroup.action.DeleteQueryGroupAction; -import org.opensearch.plugin.wlm.querygroup.action.DeleteQueryGroupRequest; ->>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteQueryGroupActionTests.java import org.opensearch.rest.RestChannel; import org.opensearch.rest.RestHandler; import org.opensearch.rest.RestRequest; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java similarity index 90% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java index e00614055d250..f374ae8f5b703 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/RuleTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm; +package org.opensearch.plugin.wlm.rule; import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.Rule; @@ -16,9 +16,7 @@ import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.plugin.wlm.rule.QueryGroupAttribute; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.plugin.wlm.rule.service.RulePersistenceService; +import org.opensearch.plugin.wlm.rule.service.WlmRulePersistenceService; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.Client; @@ -30,7 +28,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -public class RuleTestUtils { +public class WlmRuleTestUtils { public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; public static final String FEATURE_VALUE_ONE = "feature_value_one"; @@ -63,7 +61,7 @@ public static Map ruleMap() { return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); } - public static RulePersistenceService setUpRulePersistenceService(Map queryGroupMap) { + public static WlmRulePersistenceService setUpRulePersistenceService(Map queryGroupMap) { Client client = mock(Client.class); ClusterService clusterService = mock(ClusterService.class); ClusterState clusterState = mock(ClusterState.class); @@ -76,7 +74,7 @@ public static RulePersistenceService setUpRulePersistenceService(Map mapOne, Map mapTwo, boolean ruleUpdated) { diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java similarity index 71% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java index e0811b2d85e4f..d3eb6fba53821 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java @@ -17,23 +17,23 @@ import java.util.Map; import java.util.Set; -import static org.opensearch.plugin.wlm.RuleTestUtils.PATTERN_ONE; -import static org.opensearch.plugin.wlm.RuleTestUtils.SEARCH_AFTER; -import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; import static org.opensearch.plugin.wlm.rule.QueryGroupAttribute.INDEX_PATTERN; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.PATTERN_ONE; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.SEARCH_AFTER; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils._ID_ONE; -public class GetRuleRequestTests extends OpenSearchTestCase { +public class GetWlmRuleRequestTests extends OpenSearchTestCase { /** * Test case to verify the serialization and deserialization of GetRuleRequest */ public void testSerialization() throws IOException { - GetRuleRequest request = new GetRuleRequest(_ID_ONE, Map.of(INDEX_PATTERN, Set.of(PATTERN_ONE)), null); + GetWlmRuleRequest request = new GetWlmRuleRequest(_ID_ONE, Map.of(INDEX_PATTERN, Set.of(PATTERN_ONE)), null); assertEquals(_ID_ONE, request.getId()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); - GetRuleRequest otherRequest = new GetRuleRequest(streamInput); + GetWlmRuleRequest otherRequest = new GetWlmRuleRequest(streamInput); assertEquals(request.getId(), otherRequest.getId()); assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters()); } @@ -42,12 +42,12 @@ public void testSerialization() throws IOException { * Test case to verify the serialization and deserialization of GetRuleRequest when name is null */ public void testSerializationWithNull() throws IOException { - GetRuleRequest request = new GetRuleRequest((String) null, new HashMap<>(), SEARCH_AFTER); + GetWlmRuleRequest request = new GetWlmRuleRequest((String) null, new HashMap<>(), SEARCH_AFTER); assertNull(request.getId()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); - GetRuleRequest otherRequest = new GetRuleRequest(streamInput); + GetWlmRuleRequest otherRequest = new GetWlmRuleRequest(streamInput); assertEquals(request.getId(), otherRequest.getId()); assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters()); } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponseTests.java similarity index 78% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponseTests.java index 8c9983a1bb770..76f1f7a0a2534 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetRuleResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponseTests.java @@ -21,14 +21,14 @@ import java.util.HashMap; import java.util.Map; -import static org.opensearch.plugin.wlm.RuleTestUtils.SEARCH_AFTER; -import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; -import static org.opensearch.plugin.wlm.RuleTestUtils.assertEqualRules; -import static org.opensearch.plugin.wlm.RuleTestUtils.ruleMap; -import static org.opensearch.plugin.wlm.RuleTestUtils.ruleOne; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.SEARCH_AFTER; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils._ID_ONE; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.assertEqualRules; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.ruleMap; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.ruleOne; import static org.mockito.Mockito.mock; -public class GetRuleResponseTests extends OpenSearchTestCase { +public class GetWlmRuleResponseTests extends OpenSearchTestCase { /** * Test case to verify the serialization and deserialization of GetRuleResponse @@ -36,14 +36,14 @@ public class GetRuleResponseTests extends OpenSearchTestCase { public void testSerializationSingleRule() throws IOException { Map map = new HashMap<>(); map.put(_ID_ONE, ruleOne); - GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), null, RestStatus.OK); + GetWlmRuleResponse response = new GetWlmRuleResponse(Map.of(_ID_ONE, ruleOne), null, RestStatus.OK); assertEquals(response.getRules(), map); BytesStreamOutput out = new BytesStreamOutput(); response.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); - GetRuleResponse otherResponse = new GetRuleResponse(streamInput); + GetWlmRuleResponse otherResponse = new GetWlmRuleResponse(streamInput); assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); assertEqualRules(response.getRules(), otherResponse.getRules(), false); } @@ -52,14 +52,14 @@ public void testSerializationSingleRule() throws IOException { * Test case to verify the serialization and deserialization of GetRuleResponse when the result contains multiple rules */ public void testSerializationMultipleRule() throws IOException { - GetRuleResponse response = new GetRuleResponse(ruleMap(), SEARCH_AFTER, RestStatus.OK); + GetWlmRuleResponse response = new GetWlmRuleResponse(ruleMap(), SEARCH_AFTER, RestStatus.OK); assertEquals(response.getRules(), ruleMap()); BytesStreamOutput out = new BytesStreamOutput(); response.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); - GetRuleResponse otherResponse = new GetRuleResponse(streamInput); + GetWlmRuleResponse otherResponse = new GetWlmRuleResponse(streamInput); assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); assertEquals(2, otherResponse.getRules().size()); assertEqualRules(response.getRules(), otherResponse.getRules(), false); @@ -70,14 +70,14 @@ public void testSerializationMultipleRule() throws IOException { */ public void testSerializationNull() throws IOException { Map map = new HashMap<>(); - GetRuleResponse response = new GetRuleResponse(map, SEARCH_AFTER, RestStatus.OK); + GetWlmRuleResponse response = new GetWlmRuleResponse(map, SEARCH_AFTER, RestStatus.OK); assertEquals(response.getRules(), map); BytesStreamOutput out = new BytesStreamOutput(); response.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); - GetRuleResponse otherResponse = new GetRuleResponse(streamInput); + GetWlmRuleResponse otherResponse = new GetWlmRuleResponse(streamInput); assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); assertEquals(0, otherResponse.getRules().size()); } @@ -88,7 +88,7 @@ public void testSerializationNull() throws IOException { public void testToXContentGetSingleRule() throws IOException { Map map = new HashMap<>(); map.put(_ID_ONE, ruleOne); - GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), SEARCH_AFTER, RestStatus.OK); + GetWlmRuleResponse response = new GetWlmRuleResponse(Map.of(_ID_ONE, ruleOne), SEARCH_AFTER, RestStatus.OK); XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); String expected = "{\n" @@ -115,7 +115,7 @@ public void testToXContentGetSingleRule() throws IOException { */ public void testToXContentGetZeroRule() throws IOException { XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); - GetRuleResponse otherResponse = new GetRuleResponse(new HashMap<>(), null, RestStatus.OK); + GetWlmRuleResponse otherResponse = new GetWlmRuleResponse(new HashMap<>(), null, RestStatus.OK); String actual = otherResponse.toXContent(builder, mock(ToXContent.Params.class)).toString(); String expected = "{\n" + " \"rules\" : [ ]\n" + "}"; assertEquals(expected, actual); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleActionTests.java similarity index 88% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleActionTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleActionTests.java index c09600be6fc6e..a0f61cbaf3c8b 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetRuleActionTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleActionTests.java @@ -1,3 +1,4 @@ + /* * SPDX-License-Identifier: Apache-2.0 * @@ -15,12 +16,12 @@ import static org.opensearch.rest.RestRequest.Method.GET; -public class RestGetRuleActionTests extends OpenSearchTestCase { +public class RestGetWlmRuleActionTests extends OpenSearchTestCase { /** * Test case to validate the construction for RestGetRuleAction */ public void testConstruction() { - RestGetRuleAction action = new RestGetRuleAction(); + RestGetWlmRuleAction action = new RestGetWlmRuleAction(); assertNotNull(action); assertEquals("get_rule", action.getName()); List routes = action.routes(); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java similarity index 76% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java index 6bfa8b96cf16e..4a1cc1c122ecd 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/RulePersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java @@ -14,7 +14,7 @@ import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.plugin.wlm.rule.action.GetRuleResponse; +import org.opensearch.plugin.wlm.rule.action.GetWlmRuleResponse; import org.opensearch.search.sort.SortOrder; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.transport.client.Client; @@ -22,10 +22,10 @@ import java.util.HashMap; import static org.opensearch.autotagging.Rule._ID_STRING; -import static org.opensearch.plugin.wlm.RuleTestUtils.ATTRIBUTE_MAP; -import static org.opensearch.plugin.wlm.RuleTestUtils._ID_ONE; -import static org.opensearch.plugin.wlm.RuleTestUtils._ID_TWO; -import static org.opensearch.plugin.wlm.RuleTestUtils.setUpRulePersistenceService; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.ATTRIBUTE_MAP; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils._ID_ONE; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils._ID_TWO; +import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.setUpRulePersistenceService; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.anyInt; import static org.mockito.Mockito.anyString; @@ -35,16 +35,16 @@ import static org.mockito.Mockito.when; @SuppressWarnings("unchecked") -public class RulePersistenceServiceTests extends OpenSearchTestCase { +public class WlmRulePersistenceServiceTests extends OpenSearchTestCase { public void testBuildGetRuleQuery_WithId() { - RulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); + WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); BoolQueryBuilder query = rulePersistenceService.buildGetRuleQuery(_ID_ONE, new HashMap<>()); assertTrue(query.hasClauses()); assertEquals(QueryBuilders.termQuery(_ID_STRING, _ID_ONE).toString(), query.must().get(0).toString()); } public void testBuildGetRuleQuery_WithFilters() { - RulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); + WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); BoolQueryBuilder query = rulePersistenceService.buildGetRuleQuery(null, ATTRIBUTE_MAP); assertTrue(query.hasClauses()); assertEquals(1, query.must().size()); @@ -52,22 +52,22 @@ public void testBuildGetRuleQuery_WithFilters() { } public void testGetRule_WithId() { - RulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); + WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); Client client = rulePersistenceService.getClient(); - ActionListener listener = mock(ActionListener.class); + ActionListener listener = mock(ActionListener.class); SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); SetupMocksForGetRule(client, searchRequestBuilder); rulePersistenceService.getRule(_ID_ONE, new HashMap<>(), null, listener); - verify(client).prepareSearch(RulePersistenceService.RULES_INDEX); + verify(client).prepareSearch(WlmRulePersistenceService.RULES_INDEX); verify(searchRequestBuilder).setQuery(any()); verify(searchRequestBuilder).execute(any()); } public void testGetRule_WithSearchAfter() { - RulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); + WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); Client client = rulePersistenceService.getClient(); - ActionListener listener = mock(ActionListener.class); + ActionListener listener = mock(ActionListener.class); SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); SetupMocksForGetRule(client, searchRequestBuilder); when(searchRequestBuilder.addSort(anyString(), any(SortOrder.class))).thenReturn(searchRequestBuilder); From 7de0cd2b93bd1f817ad5a83ffb59abae95d2e37d Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Mon, 31 Mar 2025 16:22:32 -0700 Subject: [PATCH 06/20] Add javadocs for libs Signed-off-by: Ruirui Zhang --- .../rule/action/GetRuleRequest.java | 19 ++++++++ .../rule/action/GetRuleResponse.java | 20 +++++++-- .../rule/rest/RestGetRuleAction.java | 19 ++++++++ .../rule/service/RulePersistenceService.java | 45 +++++++++++++------ 4 files changed, 85 insertions(+), 18 deletions(-) diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java index 6cae3e860c3ad..f3b3c61ee17f6 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java @@ -34,12 +34,22 @@ public abstract class GetRuleRequest extends ActionRequest { private final Map> attributeFilters; private final String searchAfter; + /** + * Constructor for GetRuleRequest + * @param id - Rule id to get + * @param attributeFilters - Rules will be filtered based on the attribute-value mappings. + * @param searchAfter - The sort value used for pagination. + */ public GetRuleRequest(String id, Map> attributeFilters, String searchAfter) { this.id = id; this.attributeFilters = attributeFilters; this.searchAfter = searchAfter; } + /** + * Constructs a GetRuleRequest from a StreamInput for deserialization + * @param in - The {@link StreamInput} instance to read from. + */ public GetRuleRequest(StreamInput in) throws IOException { super(in); id = in.readOptionalString(); @@ -65,14 +75,23 @@ public void writeTo(StreamOutput out) throws IOException { */ protected abstract FeatureType retrieveFeatureTypeInstance(); + /** + * id getter + */ public String getId() { return id; } + /** + * attributeFilters getter + */ public Map> getAttributeFilters() { return attributeFilters; } + /** + * searchAfter getter + */ public String getSearchAfter() { return searchAfter; } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java index a92d3c4d79df8..954c39cfdfc8e 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java @@ -45,12 +45,22 @@ public abstract class GetRuleResponse extends ActionResponse implements ToXConte private final String searchAfter; private final RestStatus restStatus; + /** + * Constructor for GetRuleResponse + * @param rules - Rules get from the request + * @param searchAfter - The sort value used for pagination. + * @param restStatus - Status of the GetRuleResponse + */ public GetRuleResponse(final Map rules, String searchAfter, RestStatus restStatus) { this.rules = rules; this.searchAfter = searchAfter; this.restStatus = restStatus; } + /** + * Constructs a GetRuleResponse from a StreamInput for deserialization + * @param in - The {@link StreamInput} instance to read from. + */ public GetRuleResponse(StreamInput in) throws IOException { this(in.readMap(StreamInput::readString, Rule::new), in.readOptionalString(), RestStatus.readFrom(in)); } @@ -77,15 +87,17 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws return builder; } + /** + * rules getter + */ public Map getRules() { return rules; } + /** + * restStatus getter + */ public RestStatus getRestStatus() { return restStatus; } - - public String getSearchAfter() { - return searchAfter; - } } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java index 686b4cbf2406e..f27bdfbbb4e59 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java @@ -37,8 +37,14 @@ * @opensearch.experimental */ public abstract class RestGetRuleAction extends BaseRestHandler { + /** + * field name used for pagination + */ public static final String SEARCH_AFTER_STRING = "search_after"; + /** + * Constructor for RestGetRuleAction + */ public RestGetRuleAction() {} @Override @@ -74,6 +80,11 @@ public RestResponse buildResponse(final GetRuleResponse response) throws Excepti }; } + /** + * Abstract method for subclasses to retrieve the Attribute corresponding + * to the attribute name. + * @param name - The name of the attribute to retrieve. + */ protected abstract Attribute getAttributeFromName(String name); /** @@ -81,5 +92,13 @@ public RestResponse buildResponse(final GetRuleResponse response) throws Excepti */ protected abstract > T retrieveGetRuleActionInstance(); + /** + * Abstract method for subclasses to construct a {@link GetRuleRequest}. This method allows subclasses + * to define their own request-building logic depending on their specific needs. + * + * @param id - The ID of the rule to retrieve. + * @param attributeFilters - A map of {@link Attribute} keys to sets of string values for filtering. + * @param searchAfter - The pagination value to fetch the next set of results. + */ protected abstract GetRuleRequest buildGetRuleRequest(String id, Map> attributeFilters, String searchAfter); } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java index 69c957c94bf3d..cb5eebb94a54c 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java @@ -10,27 +10,22 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.opensearch.ResourceAlreadyExistsException; import org.opensearch.ResourceNotFoundException; -import org.opensearch.action.admin.indices.create.CreateIndexRequest; -import org.opensearch.action.admin.indices.create.CreateIndexResponse; -import org.opensearch.action.index.IndexRequest; import org.opensearch.action.search.SearchRequestBuilder; import org.opensearch.action.search.SearchResponse; -import org.opensearch.action.update.UpdateRequest; import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.FeatureType; import org.opensearch.autotagging.Rule; import org.opensearch.autotagging.Rule.Builder; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; -import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.action.ActionListener; -import org.opensearch.core.action.ActionResponse; import org.opensearch.core.rest.RestStatus; -import org.opensearch.core.xcontent.*; +import org.opensearch.core.xcontent.DeprecationHandler; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.XContentParser; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.rule.action.GetRuleResponse; @@ -39,13 +34,11 @@ import org.opensearch.transport.client.Client; import java.io.IOException; -import java.time.Instant; import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; -import java.util.HashMap; import java.util.stream.Collectors; import static org.opensearch.autotagging.Rule._ID_STRING; @@ -68,6 +61,11 @@ public abstract class RulePersistenceService { */ public static final int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 50; + /** + * Constructor for RulePersistenceService + * @param clusterService - The {@link ClusterService} to be used by RulePersistenceService + * @param client - The {@link Client} to be used by RulePersistenceService + */ @Inject public RulePersistenceService(final ClusterService clusterService, final Client client) { this.clusterService = clusterService; @@ -109,8 +107,8 @@ public void getRule( /** * Builds a bool query to retrieve rules from the rules index, applying attribute-based filters * when needed and ensuring that the rules are associated with the query group feature type. - * @param id The ID of the rule to fetch. If not null, the search will return only this specific rule. - * @param attributeFilters A map of attributes to their associated set of values used to filter the rules. + * @param id - The ID of the rule to fetch. If not null, the search will return only this specific rule. + * @param attributeFilters - A map of attributes to their associated set of values used to filter the rules. */ public BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeFilters) { BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); @@ -178,20 +176,39 @@ private boolean isExistingQueryGroup(String queryGroupId) { */ protected abstract FeatureType retrieveFeatureTypeInstance(); - protected abstract T buildGetRuleResponse(Map ruleMap, String nextSearchAfter, RestStatus restStatus); + /** + * Builds a GetRuleResponse. Subclasses can override based on their own needs. + * @param ruleMap - The map of rules the response contains + * @param nextSearchAfter - The search after string containing a rule id used for pagination + * @param restStatus - Status of the response + */ + protected abstract T buildGetRuleResponse( + Map ruleMap, + String nextSearchAfter, + RestStatus restStatus + ); private ThreadContext.StoredContext getContext() { return client.threadPool().getThreadContext().stashContext(); } + /** + * Get the max return size allowed for page for get request when pagination is needed + */ public int getMaxReturnSizeAllowedPerGetRequest() { return MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST; } + /** + * client getter + */ public Client getClient() { return client; } + /** + * clusterService getter + */ public ClusterService getClusterService() { return clusterService; } From dc1a8b85c93e0d99a2d4e3a81f29b769ab70e00f Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 2 Apr 2025 00:08:06 -0700 Subject: [PATCH 07/20] modify based on comments Signed-off-by: Ruirui Zhang --- .../rule/action/GetRuleRequest.java | 16 +- .../rule/service/RulePersistenceService.java | 144 ++++-------------- .../rule/service/RuleProcessingService.java | 63 ++++++++ .../rule/service/RuleResponseBuilder.java | 29 ++++ .../opensearch/rule/service/RuleService.java | 74 +++++++++ .../org/opensearch/rule/autotagging/Rule.java | 6 + .../wlm/WorkloadManagementPluginModule.java | 13 +- .../wlm/rule/action/GetWlmRuleRequest.java | 15 +- .../action/TransportGetWlmRuleAction.java | 16 +- .../wlm/rule/rest/RestGetWlmRuleAction.java | 3 +- .../service/WlmRulePersistenceService.java | 33 ++-- .../service/WlmRuleProcessingService.java | 20 +++ .../rule/service/WlmRuleResponseBuilder.java | 25 +++ .../rule/action/GetWlmRuleRequestTests.java | 10 +- .../WlmRulePersistenceServiceTests.java | 5 +- 15 files changed, 310 insertions(+), 162 deletions(-) create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleProcessingService.java create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleResponseBuilder.java create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleService.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleProcessingService.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleResponseBuilder.java diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java index f3b3c61ee17f6..052ceb90bbf77 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java @@ -33,27 +33,31 @@ public abstract class GetRuleRequest extends ActionRequest { private final String id; private final Map> attributeFilters; private final String searchAfter; + private final FeatureType featureType; /** * Constructor for GetRuleRequest * @param id - Rule id to get * @param attributeFilters - Rules will be filtered based on the attribute-value mappings. * @param searchAfter - The sort value used for pagination. + * @param featureType - The feature type related to rule. */ - public GetRuleRequest(String id, Map> attributeFilters, String searchAfter) { + public GetRuleRequest(String id, Map> attributeFilters, String searchAfter, FeatureType featureType) { this.id = id; this.attributeFilters = attributeFilters; this.searchAfter = searchAfter; + this.featureType = featureType; } /** - * Constructs a GetRuleRequest from a StreamInput for deserialization + * Constructs a GetRuleRequest from a StreamInput for deserialization. * @param in - The {@link StreamInput} instance to read from. */ public GetRuleRequest(StreamInput in) throws IOException { super(in); id = in.readOptionalString(); - attributeFilters = in.readMap(i -> Attribute.from(i, retrieveFeatureTypeInstance()), i -> new HashSet<>(i.readStringList())); + featureType = FeatureType.from(in); + attributeFilters = in.readMap(i -> Attribute.from(i, featureType), i -> new HashSet<>(i.readStringList())); searchAfter = in.readOptionalString(); } @@ -66,15 +70,11 @@ public ActionRequestValidationException validate() { public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeOptionalString(id); + featureType.writeTo(out); out.writeMap(attributeFilters, (output, attribute) -> attribute.writeTo(output), StreamOutput::writeStringCollection); out.writeOptionalString(searchAfter); } - /** - * Abstract method for subclasses to provide specific FeatureType Instance - */ - protected abstract FeatureType retrieveFeatureTypeInstance(); - /** * id getter */ diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java index cb5eebb94a54c..08856b38cdb08 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java @@ -10,36 +10,20 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.opensearch.ResourceNotFoundException; import org.opensearch.action.search.SearchRequestBuilder; import org.opensearch.action.search.SearchResponse; import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule; -import org.opensearch.autotagging.Rule.Builder; import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.inject.Inject; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.core.xcontent.DeprecationHandler; -import org.opensearch.core.xcontent.MediaTypeRegistry; -import org.opensearch.core.xcontent.NamedXContentRegistry; -import org.opensearch.core.xcontent.XContentParser; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilders; -import org.opensearch.rule.action.GetRuleResponse; -import org.opensearch.search.SearchHit; import org.opensearch.search.sort.SortOrder; import org.opensearch.transport.client.Client; -import java.io.IOException; -import java.util.Arrays; -import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.Set; -import java.util.stream.Collectors; import static org.opensearch.autotagging.Rule._ID_STRING; @@ -47,31 +31,39 @@ * This class encapsulates the logic to manage the lifecycle of rules at index level * @opensearch.experimental */ -public abstract class RulePersistenceService { +public interface RulePersistenceService { + /** + * The default maximum number of results allowed per GET request + */ + int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 50; + /** * The system index name used for storing rules */ - public static final String RULES_INDEX = ".rules"; - private final Client client; - private final ClusterService clusterService; - private static final Logger logger = LogManager.getLogger(RulePersistenceService.class); + String getIndexName(); + + /** + * client getter + */ + Client getClient(); /** - * The maximum number of results allowed per GET request + * clusterService getter */ - public static final int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 50; + ClusterService getClusterService(); /** - * Constructor for RulePersistenceService - * @param clusterService - The {@link ClusterService} to be used by RulePersistenceService - * @param client - The {@link Client} to be used by RulePersistenceService + * Get the max return size allowed for page for get request when pagination is needed */ - @Inject - public RulePersistenceService(final ClusterService clusterService, final Client client) { - this.clusterService = clusterService; - this.client = client; + default int getMaxReturnSizeAllowedPerGetRequest() { + return MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST; } + /** + * logger for RulePersistenceService + */ + Logger logger = LogManager.getLogger(RulePersistenceService.class); + /** * Entry point for the get rule api logic in persistence service. If id is provided, we only get a single rule. * Otherwise, we get all rules that satisfy the attributeFilters. @@ -80,24 +72,24 @@ public RulePersistenceService(final ClusterService clusterService, final Client * @param searchAfter - The sort values from the last document of the previous page, used for pagination * @param listener - ActionListener for GetRuleResponse */ - public void getRule( + default void getRule( String id, Map> attributeFilters, String searchAfter, - ActionListener listener + ActionListener listener ) { // Stash the current thread context when interacting with system index to perform // operations as the system itself, bypassing authorization checks. This ensures that // actions within this block are trusted and executed with system-level privileges. try (ThreadContext.StoredContext context = getContext()) { BoolQueryBuilder boolQuery = buildGetRuleQuery(id, attributeFilters); - SearchRequestBuilder searchRequest = client.prepareSearch(RULES_INDEX) + SearchRequestBuilder searchRequest = getClient().prepareSearch(getIndexName()) .setQuery(boolQuery) .setSize(getMaxReturnSizeAllowedPerGetRequest()); if (searchAfter != null) { searchRequest.addSort(_ID_STRING, SortOrder.ASC).searchAfter(new Object[] { searchAfter }); } - searchRequest.execute(ActionListener.wrap(searchResponse -> handleGetRuleResponse(id, searchResponse, listener), e -> { + searchRequest.execute(ActionListener.wrap(listener::onResponse, e -> { logger.error("Failed to fetch all rules: {}", e.getMessage()); listener.onFailure(e); })); @@ -107,10 +99,10 @@ public void getRule( /** * Builds a bool query to retrieve rules from the rules index, applying attribute-based filters * when needed and ensuring that the rules are associated with the query group feature type. - * @param id - The ID of the rule to fetch. If not null, the search will return only this specific rule. - * @param attributeFilters - A map of attributes to their associated set of values used to filter the rules. + * @param id The ID of the rule to fetch. If not null, the search will return only this specific rule. + * @param attributeFilters A map of attributes to their associated set of values used to filter the rules. */ - public BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeFilters) { + default BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeFilters) { BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); if (id != null) { return boolQuery.must(QueryBuilders.termQuery(_ID_STRING, id)); @@ -130,86 +122,12 @@ public BoolQueryBuilder buildGetRuleQuery(String id, Map> return boolQuery; } - /** - * Process searchResponse from index and send a GetRuleResponse - * @param searchResponse - Response received from index - * @param listener - ActionListener for GetRuleResponse - */ - void handleGetRuleResponse(String id, SearchResponse searchResponse, ActionListener listener) { - List hits = Arrays.asList(searchResponse.getHits().getHits()); - if (id != null && hits.isEmpty()) { - logger.error("Rule with ID " + id + " not found."); - listener.onFailure(new ResourceNotFoundException("Rule with ID " + id + " doesn't exist in the .rules index.")); - return; - } - Map ruleMap = hits.stream() - .map(hit -> parseRule(hit.getId(), hit.getSourceAsString())) - .filter(Objects::nonNull) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - String nextSearchAfter = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); - listener.onResponse(buildGetRuleResponse(ruleMap, nextSearchAfter, RestStatus.OK)); - } - - /** - * Parses a source string into a Rule object - * @param id - document id for the Rule object - * @param source - The raw source string representing the rule to be parsed - */ - private Map.Entry parseRule(String id, String source) { - try ( - XContentParser parser = MediaTypeRegistry.JSON.xContent() - .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, source) - ) { - return Map.entry(id, Builder.fromXContent(parser, retrieveFeatureTypeInstance()).build()); - } catch (IOException e) { - logger.info("Issue met when parsing rule for ID {}: {}", id, e.getMessage()); - return null; - } - } - - private boolean isExistingQueryGroup(String queryGroupId) { - return clusterService.state().metadata().queryGroups().containsKey(queryGroupId); - } - /** * Abstract method for subclasses to provide specific FeatureType Instance */ - protected abstract FeatureType retrieveFeatureTypeInstance(); - - /** - * Builds a GetRuleResponse. Subclasses can override based on their own needs. - * @param ruleMap - The map of rules the response contains - * @param nextSearchAfter - The search after string containing a rule id used for pagination - * @param restStatus - Status of the response - */ - protected abstract T buildGetRuleResponse( - Map ruleMap, - String nextSearchAfter, - RestStatus restStatus - ); + FeatureType retrieveFeatureTypeInstance(); private ThreadContext.StoredContext getContext() { - return client.threadPool().getThreadContext().stashContext(); - } - - /** - * Get the max return size allowed for page for get request when pagination is needed - */ - public int getMaxReturnSizeAllowedPerGetRequest() { - return MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST; - } - - /** - * client getter - */ - public Client getClient() { - return client; - } - - /** - * clusterService getter - */ - public ClusterService getClusterService() { - return clusterService; + return getClient().threadPool().getThreadContext().stashContext(); } } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleProcessingService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleProcessingService.java new file mode 100644 index 0000000000000..95353a25d2343 --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleProcessingService.java @@ -0,0 +1,63 @@ +/* + * 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.rule.service; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.autotagging.FeatureType; +import org.opensearch.autotagging.Rule; +import org.opensearch.common.collect.Tuple; +import org.opensearch.search.SearchHit; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.opensearch.autotagging.Rule.parseRule; + +/** + * Interface defining the contract for processing and parsing rule-related logic. + * @opensearch.experimental + */ +public interface RuleProcessingService { + /** + * logger for RuleProcessingService + */ + Logger logger = LogManager.getLogger(RuleProcessingService.class); + + /** + * function to retrieve the feature type corresponding to the rule + */ + FeatureType retrieveFeatureTypeInstance(); + + /** + * Process searchResponse from index and send a GetRuleResponse + * @param id - Rule id fetched + * @param searchResponse - Response received from index + */ + default Tuple, String> parseGetRuleResponse(String id, SearchResponse searchResponse) { + List hits = Arrays.asList(searchResponse.getHits().getHits()); + if (id != null && hits.isEmpty()) { + logger.error("Rule with ID " + id + " not found."); + return new Tuple<>(new HashMap<>(), null); + } + Map ruleMap = hits.stream() + .collect( + Collectors.toMap( + SearchHit::getId, + hit -> parseRule(hit.getId(), hit.getSourceAsString(), retrieveFeatureTypeInstance()).getValue() + ) + ); + String lastValidRuleId = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); + return new Tuple<>(ruleMap, lastValidRuleId); + } +} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleResponseBuilder.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleResponseBuilder.java new file mode 100644 index 0000000000000..7ef1046effa01 --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleResponseBuilder.java @@ -0,0 +1,29 @@ +/* + * 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.rule.service; + +import org.opensearch.autotagging.Rule; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.rule.action.GetRuleResponse; + +import java.util.Map; + +/** + * A builder interface for rule lifecycle response objects. + * @opensearch.experimental + */ +public interface RuleResponseBuilder { + /** + * Builds a GetRuleResponse object containing rule data. + * @param ruleMap A map of rule IDs to their corresponding {@link Rule} objects. + * @param nextSearchAfter The ID used for pagination to fetch the next set of results, if applicable. + * @param restStatus The response rest status. + */ + T buildGetRuleResponse(Map ruleMap, String nextSearchAfter, RestStatus restStatus); +} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleService.java new file mode 100644 index 0000000000000..57fbcb4a01ff5 --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleService.java @@ -0,0 +1,74 @@ +/* + * 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.rule.service; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.autotagging.Rule; +import org.opensearch.common.collect.Tuple; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.rule.action.GetRuleRequest; +import org.opensearch.rule.action.GetRuleResponse; + +import java.util.Map; + +/** + * Service class responsible for processing the business logic for rule lifecycle operations + * @opensearch.experimental + */ +public class RuleService { + private final RulePersistenceService rulePersistenceService; + private final RuleProcessingService ruleProcessingService; + private final RuleResponseBuilder ruleResponseBuilder; + private final Logger logger = LogManager.getLogger(RuleService.class); + + /** + * Constructs a {@link RuleService} with the specified dependencies. + * @param rulePersistenceService The persistence service for retrieving and storing rules. + * @param ruleProcessingService The service responsible for processing rule data. + * @param ruleResponseBuilder The builder used to construct rule response objects. + */ + @Inject + public RuleService( + RulePersistenceService rulePersistenceService, + RuleProcessingService ruleProcessingService, + RuleResponseBuilder ruleResponseBuilder + ) { + this.rulePersistenceService = rulePersistenceService; + this.ruleProcessingService = ruleProcessingService; + this.ruleResponseBuilder = ruleResponseBuilder; + } + + /** + * Processes a request to retrieve rules based on specified criteria. + * @param request The {@link GetRuleRequest} containing rule retrieval parameters. + * @param listener The {@link ActionListener} that handles the asynchronous response. + */ + public void processGetRuleRequest(GetRuleRequest request, ActionListener listener) { + rulePersistenceService.getRule(request.getId(), request.getAttributeFilters(), request.getSearchAfter(), new ActionListener<>() { + @Override + public void onResponse(SearchResponse searchResponse) { + Tuple, String> responseParams = ruleProcessingService.parseGetRuleResponse( + request.getId(), + searchResponse + ); + listener.onResponse(ruleResponseBuilder.buildGetRuleResponse(responseParams.v1(), responseParams.v2(), RestStatus.OK)); + } + + @Override + public void onFailure(Exception e) { + logger.error("Failed to fetch rules: {}", e.getMessage()); + listener.onFailure(e); + } + }); + } +} diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/Rule.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/Rule.java index e1b003f0085c6..f4f3d524ff642 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/Rule.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/Rule.java @@ -8,9 +8,14 @@ package org.opensearch.rule.autotagging; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; 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 org.opensearch.core.xcontent.DeprecationHandler; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParseException; @@ -56,6 +61,7 @@ public class Rule implements Writeable, ToXContentObject { * updated_at field */ public static final String UPDATED_AT_STRING = "updated_at"; + public static final Logger logger = LogManager.getLogger(Rule.class); /** * Main constructor diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java index 0ed3781aae409..e65c67b148660 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java @@ -10,7 +10,14 @@ import org.opensearch.common.inject.AbstractModule; import org.opensearch.common.inject.Singleton; -import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; +import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService;; +import org.opensearch.plugin.wlm.rule.service.WlmRulePersistenceService; +import org.opensearch.plugin.wlm.rule.service.WlmRuleProcessingService; +import org.opensearch.plugin.wlm.rule.service.WlmRuleResponseBuilder; +import org.opensearch.rule.service.RulePersistenceService; +import org.opensearch.rule.service.RuleProcessingService; +import org.opensearch.rule.service.RuleResponseBuilder; +import org.opensearch.rule.service.RuleService; /** * Guice Module to manage WorkloadManagement related objects @@ -26,5 +33,9 @@ protected void configure() { // Bind WorkloadGroupPersistenceService as a singleton to ensure a single instance is used, // preventing multiple throttling key registrations in the constructor. bind(WorkloadGroupPersistenceService.class).in(Singleton.class); + bind(RulePersistenceService.class).to(WlmRulePersistenceService.class).in(Singleton.class); + bind(RuleProcessingService.class).to(WlmRuleProcessingService.class).in(Singleton.class); + bind(RuleResponseBuilder.class).to(WlmRuleResponseBuilder.class).in(Singleton.class); + bind(RuleService.class).in(Singleton.class); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java index d2b3e6f9c962f..97120e9425ed1 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java @@ -11,7 +11,6 @@ import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.FeatureType; import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import org.opensearch.rule.action.GetRuleRequest; import java.io.IOException; @@ -20,10 +19,6 @@ /** * A request to get workload management Rules in workload management - * Example Request: - * curl -X GET "localhost:9200/_wlm/rule" - get all rules - * curl -X GET "localhost:9200/_wlm/rule/{_id}" - get single rule by id - * curl -X GET "localhost:9200/_wlm/rule?index_pattern=a,b" - get all rules containing attribute index_pattern as a or b * @opensearch.experimental */ public class GetWlmRuleRequest extends GetRuleRequest { @@ -32,9 +27,10 @@ public class GetWlmRuleRequest extends GetRuleRequest { * @param id - Rule id to get * @param attributeFilters - A map containing the attributes to filter on * @param searchAfter - A string used for pagination + * @param featureType - The featureType instance for wlm */ - public GetWlmRuleRequest(String id, Map> attributeFilters, String searchAfter) { - super(id, attributeFilters, searchAfter); + public GetWlmRuleRequest(String id, Map> attributeFilters, String searchAfter, FeatureType featureType) { + super(id, attributeFilters, searchAfter, featureType); } /** @@ -44,9 +40,4 @@ public GetWlmRuleRequest(String id, Map> attributeFilters public GetWlmRuleRequest(StreamInput in) throws IOException { super(in); } - - @Override - protected FeatureType retrieveFeatureTypeInstance() { - return QueryGroupFeatureType.INSTANCE; - } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java index ca20ea4d1815a..85bc110b61a03 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java @@ -12,7 +12,7 @@ import org.opensearch.action.support.HandledTransportAction; import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; -import org.opensearch.plugin.wlm.rule.service.WlmRulePersistenceService; +import org.opensearch.rule.service.RuleService; import org.opensearch.tasks.Task; import org.opensearch.transport.TransportService; @@ -22,26 +22,22 @@ */ public class TransportGetWlmRuleAction extends HandledTransportAction { - private final WlmRulePersistenceService rulePersistenceService; + private final RuleService ruleService; /** * Constructor for TransportGetWlmRuleAction * @param transportService - a {@link TransportService} object * @param actionFilters - a {@link ActionFilters} object - * @param rulePersistenceService - a {@link WlmRulePersistenceService} object + * @param ruleService - a {@link RuleService} object */ @Inject - public TransportGetWlmRuleAction( - TransportService transportService, - ActionFilters actionFilters, - WlmRulePersistenceService rulePersistenceService - ) { + public TransportGetWlmRuleAction(TransportService transportService, ActionFilters actionFilters, RuleService ruleService) { super(GetWlmRuleAction.NAME, transportService, actionFilters, GetWlmRuleRequest::new); - this.rulePersistenceService = rulePersistenceService; + this.ruleService = ruleService; } @Override protected void doExecute(Task task, GetWlmRuleRequest request, ActionListener listener) { - rulePersistenceService.getRule(request.getId(), request.getAttributeFilters(), request.getSearchAfter(), listener); + ruleService.processGetRuleRequest(request, listener); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java index 11f7776c77c21..55a07082928dd 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java @@ -11,6 +11,7 @@ import org.opensearch.action.ActionType; import org.opensearch.autotagging.Attribute; import org.opensearch.plugin.wlm.rule.QueryGroupAttribute; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import org.opensearch.plugin.wlm.rule.action.GetWlmRuleAction; import org.opensearch.plugin.wlm.rule.action.GetWlmRuleRequest; import org.opensearch.rule.action.GetRuleRequest; @@ -59,6 +60,6 @@ protected > T retrieveGetRuleAct @Override protected GetRuleRequest buildGetRuleRequest(String id, Map> attributeFilters, String searchAfter) { - return new GetWlmRuleRequest(id, attributeFilters, searchAfter); + return new GetWlmRuleRequest(id, attributeFilters, searchAfter, QueryGroupFeatureType.INSTANCE); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java index 5ee6a1d65dbb0..b4d77f3ca0a90 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java @@ -9,25 +9,28 @@ package org.opensearch.plugin.wlm.rule.service; import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; import org.opensearch.common.settings.Settings; -import org.opensearch.core.rest.RestStatus; import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.plugin.wlm.rule.action.GetWlmRuleResponse; -import org.opensearch.rule.action.GetRuleResponse; import org.opensearch.rule.service.RulePersistenceService; import org.opensearch.transport.client.Client; -import java.util.Map; - /** * This class encapsulates the logic to manage the lifecycle of workload management rules at index level * @opensearch.experimental */ @SuppressWarnings("unchecked") -public class WlmRulePersistenceService extends RulePersistenceService { +public class WlmRulePersistenceService implements RulePersistenceService { + public static final String RULES_INDEX = ".rules"; + private final ClusterService clusterService; + private final Client client; + + @Override + public String getIndexName() { + return RULES_INDEX; + } + /** * Constructor for WlmRulePersistenceService * @param clusterService {@link ClusterService} - The cluster service to be used by RulePersistenceService @@ -35,16 +38,22 @@ public class WlmRulePersistenceService extends RulePersistenceService { */ @Inject public WlmRulePersistenceService(ClusterService clusterService, Client client) { - super(clusterService, client); + this.client = client; + this.clusterService = clusterService; } @Override - protected FeatureType retrieveFeatureTypeInstance() { - return QueryGroupFeatureType.INSTANCE; + public Client getClient() { + return client; } @Override - protected T buildGetRuleResponse(Map ruleMap, String nextSearchAfter, RestStatus restStatus) { - return (T) new GetWlmRuleResponse(ruleMap, nextSearchAfter, restStatus); + public ClusterService getClusterService() { + return clusterService; + } + + @Override + public FeatureType retrieveFeatureTypeInstance() { + return QueryGroupFeatureType.INSTANCE; } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleProcessingService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleProcessingService.java new file mode 100644 index 0000000000000..3e50ad02ad5d2 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleProcessingService.java @@ -0,0 +1,20 @@ +/* + * 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.rule.service; + +import org.opensearch.autotagging.FeatureType; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; +import org.opensearch.rule.service.RuleProcessingService; + +public class WlmRuleProcessingService implements RuleProcessingService { + @Override + public FeatureType retrieveFeatureTypeInstance() { + return QueryGroupFeatureType.INSTANCE; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleResponseBuilder.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleResponseBuilder.java new file mode 100644 index 0000000000000..819d5542bab5e --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleResponseBuilder.java @@ -0,0 +1,25 @@ +/* + * 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.rule.service; + +import org.opensearch.autotagging.Rule; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.plugin.wlm.rule.action.GetWlmRuleResponse; +import org.opensearch.rule.action.GetRuleResponse; +import org.opensearch.rule.service.RuleResponseBuilder; + +import java.util.Map; + +@SuppressWarnings("unchecked") +public class WlmRuleResponseBuilder implements RuleResponseBuilder { + @Override + public T buildGetRuleResponse(Map ruleMap, String nextSearchAfter, RestStatus restStatus) { + return (T) new GetWlmRuleResponse(ruleMap, nextSearchAfter, restStatus); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java index d3eb6fba53821..ccf65d250814f 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java @@ -10,6 +10,7 @@ import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; @@ -28,7 +29,12 @@ public class GetWlmRuleRequestTests extends OpenSearchTestCase { * Test case to verify the serialization and deserialization of GetRuleRequest */ public void testSerialization() throws IOException { - GetWlmRuleRequest request = new GetWlmRuleRequest(_ID_ONE, Map.of(INDEX_PATTERN, Set.of(PATTERN_ONE)), null); + GetWlmRuleRequest request = new GetWlmRuleRequest( + _ID_ONE, + Map.of(INDEX_PATTERN, Set.of(PATTERN_ONE)), + null, + QueryGroupFeatureType.INSTANCE + ); assertEquals(_ID_ONE, request.getId()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); @@ -42,7 +48,7 @@ public void testSerialization() throws IOException { * Test case to verify the serialization and deserialization of GetRuleRequest when name is null */ public void testSerializationWithNull() throws IOException { - GetWlmRuleRequest request = new GetWlmRuleRequest((String) null, new HashMap<>(), SEARCH_AFTER); + GetWlmRuleRequest request = new GetWlmRuleRequest((String) null, new HashMap<>(), SEARCH_AFTER, QueryGroupFeatureType.INSTANCE); assertNull(request.getId()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java index 4a1cc1c122ecd..30d3359078140 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java @@ -14,7 +14,6 @@ import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.plugin.wlm.rule.action.GetWlmRuleResponse; import org.opensearch.search.sort.SortOrder; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.transport.client.Client; @@ -54,7 +53,7 @@ public void testBuildGetRuleQuery_WithFilters() { public void testGetRule_WithId() { WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); Client client = rulePersistenceService.getClient(); - ActionListener listener = mock(ActionListener.class); + ActionListener listener = mock(ActionListener.class); SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); SetupMocksForGetRule(client, searchRequestBuilder); @@ -67,7 +66,7 @@ public void testGetRule_WithId() { public void testGetRule_WithSearchAfter() { WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); Client client = rulePersistenceService.getClient(); - ActionListener listener = mock(ActionListener.class); + ActionListener listener = mock(ActionListener.class); SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); SetupMocksForGetRule(client, searchRequestBuilder); when(searchRequestBuilder.addSort(anyString(), any(SortOrder.class))).thenReturn(searchRequestBuilder); From 35abfd3cf9af67f1184309395be1e939c3c2ee7c Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 2 Apr 2025 20:51:25 -0700 Subject: [PATCH 08/20] rebase modify based on comments Signed-off-by: Ruirui Zhang --- .../rule/Utils/IndexStoredRuleParser.java | 40 ++++++ .../rule/Utils/IndexStoredRuleUtils.java | 41 ++++++ .../opensearch/rule/Utils}/package-info.java | 2 +- .../rule/action/GetRuleRequest.java | 2 +- .../rule/action/GetRuleResponse.java | 2 +- .../IndexStoredRulePersistenceService.java | 126 +++++++++++++++++ .../rule/service/RulePersistenceService.java | 122 +---------------- .../rule/service/RuleProcessingService.java | 63 --------- .../rule/service/RuleResponseBuilder.java | 29 ---- .../opensearch/rule/service/RuleService.java | 74 ---------- .../wlm/WorkloadManagementPluginModule.java | 5 +- .../wlm/rule/action/GetWlmRuleAction.java | 5 +- .../wlm/rule/action/GetWlmRuleRequest.java | 43 ------ .../wlm/rule/action/GetWlmRuleResponse.java | 55 -------- .../action/TransportGetWlmRuleAction.java | 25 ++-- .../wlm/rule/rest/RestGetWlmRuleAction.java | 3 +- .../service/WlmRulePersistenceService.java | 59 -------- .../service/WlmRuleProcessingService.java | 20 --- .../rule/service/WlmRuleResponseBuilder.java | 25 ---- .../plugin/wlm/rule/WlmRuleTestUtils.java | 41 +++--- .../rule/action/GetWlmRuleRequestTests.java | 60 --------- .../rule/action/GetWlmRuleResponseTests.java | 123 ----------------- .../WlmRulePersistenceServiceTests.java | 127 +++++++----------- 23 files changed, 302 insertions(+), 790 deletions(-) create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleParser.java create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleUtils.java rename {plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service => libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils}/package-info.java (86%) create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleProcessingService.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleResponseBuilder.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleService.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponse.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleProcessingService.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleResponseBuilder.java delete mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java delete mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponseTests.java diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleParser.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleParser.java new file mode 100644 index 0000000000000..0e921b7619c60 --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleParser.java @@ -0,0 +1,40 @@ +/* + * 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.rule.Utils; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.autotagging.FeatureType; +import org.opensearch.autotagging.Rule; +import org.opensearch.core.xcontent.DeprecationHandler; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.XContentParser; + +import java.io.IOException; + +public class IndexStoredRuleParser { + private static final Logger logger = LogManager.getLogger(IndexStoredRuleParser.class); + + /** + * Parses a source string into a Rule object + * @param source - The raw source string representing the rule to be parsed + */ + public static Rule parseRule(String source, FeatureType featureType) { + try ( + XContentParser parser = MediaTypeRegistry.JSON.xContent() + .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, source) + ) { + return Rule.Builder.fromXContent(parser, featureType).build(); + } catch (IOException e) { + logger.info("Issue met when parsing rule {}: {}", source, e.getMessage()); + throw new RuntimeException("Cannot parse rule from index: " + source); + } + } +} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleUtils.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleUtils.java new file mode 100644 index 0000000000000..e9f29fbe3c7c0 --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleUtils.java @@ -0,0 +1,41 @@ +/* + * 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.rule.Utils; + +import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.FeatureType; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.QueryBuilders; + +import java.util.Map; +import java.util.Set; + +import static org.opensearch.autotagging.Rule._ID_STRING; + +public class IndexStoredRuleUtils { + public static BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeFilters, FeatureType featureType) { + BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); + if (id != null) { + return boolQuery.must(QueryBuilders.termQuery(_ID_STRING, id)); + } + for (Map.Entry> entry : attributeFilters.entrySet()) { + Attribute attribute = entry.getKey(); + Set values = entry.getValue(); + if (values != null && !values.isEmpty()) { + BoolQueryBuilder attributeQuery = QueryBuilders.boolQuery(); + for (String value : values) { + attributeQuery.should(QueryBuilders.matchQuery(attribute.getName(), value)); + } + boolQuery.must(attributeQuery); + } + } + boolQuery.filter(QueryBuilders.existsQuery(featureType.getName())); + return boolQuery; + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/package-info.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/package-info.java similarity index 86% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/package-info.java rename to libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/package-info.java index 603145e617cf9..12990034ed50d 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/package-info.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/package-info.java @@ -9,4 +9,4 @@ /** * Package for the service classes related to query groups in WorkloadManagementPlugin */ -package org.opensearch.plugin.wlm.rule.service; +package org.opensearch.rule.Utils; diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java index 052ceb90bbf77..03d1899a2e6a8 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java @@ -29,7 +29,7 @@ * curl -X GET "localhost:9200/_wlm/rule?index_pattern=a,b" - get all rules containing attribute index_pattern as a or b * @opensearch.experimental */ -public abstract class GetRuleRequest extends ActionRequest { +public class GetRuleRequest extends ActionRequest { private final String id; private final Map> attributeFilters; private final String searchAfter; diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java index 954c39cfdfc8e..110522ff86af6 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java @@ -40,7 +40,7 @@ * } * @opensearch.experimental */ -public abstract class GetRuleResponse extends ActionResponse implements ToXContent, ToXContentObject { +public class GetRuleResponse extends ActionResponse implements ToXContent, ToXContentObject { private final Map rules; private final String searchAfter; private final RestStatus restStatus; diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java new file mode 100644 index 0000000000000..9ca715e4cecc0 --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -0,0 +1,126 @@ +/* + * 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.rule.service; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.search.SearchRequestBuilder; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.FeatureType; +import org.opensearch.autotagging.Rule; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.rule.Utils.IndexStoredRuleParser; +import org.opensearch.rule.Utils.IndexStoredRuleUtils; +import org.opensearch.rule.action.GetRuleRequest; +import org.opensearch.rule.action.GetRuleResponse; +import org.opensearch.search.SearchHit; +import org.opensearch.search.sort.SortOrder; +import org.opensearch.transport.client.Client; + +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.opensearch.autotagging.Rule._ID_STRING; + +/** + * This class encapsulates the logic to manage the lifecycle of rules at index level + * @opensearch.experimental + */ +public class IndexStoredRulePersistenceService implements RulePersistenceService { + /** + * The system index name used for storing rules + */ + private final String indexName; + private final Client client; + private final FeatureType featureType; + private final int maxReturnSizeAllowedPerGetRequest; + private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class); + + public IndexStoredRulePersistenceService( + String indexName, + Client client, + FeatureType featureType, + int maxReturnSizeAllowedPerGetRequest + ) { + this.indexName = indexName; + this.client = client; + this.featureType = featureType; + this.maxReturnSizeAllowedPerGetRequest = maxReturnSizeAllowedPerGetRequest; + } + + public void getRule(GetRuleRequest getRuleRequest, ActionListener listener) { + getRuleFromIndex(getRuleRequest.getId(), getRuleRequest.getAttributeFilters(), getRuleRequest.getSearchAfter(), listener); + } + + /** + * Entry point for the get rule api logic in persistence service. If id is provided, we only get a single rule. + * Otherwise, we get all rules that satisfy the attributeFilters. + * @param id - The id of the rule to get. + * @param attributeFilters - A map containing the attributes that user want to filter on + * @param searchAfter - The sort values from the last document of the previous page, used for pagination + * @param listener - ActionListener for GetRuleResponse + */ + public void getRuleFromIndex( + String id, + Map> attributeFilters, + String searchAfter, + ActionListener listener + ) { + // Stash the current thread context when interacting with system index to perform + // operations as the system itself, bypassing authorization checks. This ensures that + // actions within this block are trusted and executed with system-level privileges. + try (ThreadContext.StoredContext context = getContext()) { + BoolQueryBuilder boolQuery = IndexStoredRuleUtils.buildGetRuleQuery(id, attributeFilters, featureType); + SearchRequestBuilder searchRequest = client.prepareSearch(indexName) + .setQuery(boolQuery) + .setSize(maxReturnSizeAllowedPerGetRequest); + if (searchAfter != null) { + searchRequest.addSort(_ID_STRING, SortOrder.ASC).searchAfter(new Object[] { searchAfter }); + } + searchRequest.execute(ActionListener.wrap(searchResponse -> handleGetRuleResponse(id, searchResponse, listener), e -> { + logger.error("Failed to fetch all rules: {}", e.getMessage()); + listener.onFailure(e); + })); + } + } + + /** + * Process searchResponse from index and send a GetRuleResponse + * @param searchResponse - Response received from index + * @param listener - ActionListener for GetRuleResponse + */ + void handleGetRuleResponse(String id, SearchResponse searchResponse, ActionListener listener) { + List hits = Arrays.asList(searchResponse.getHits().getHits()); + if (id != null && hits.isEmpty()) { + logger.error("Rule with ID " + id + " not found."); + listener.onFailure(new ResourceNotFoundException("Rule with ID " + id + " doesn't exist in the .rules index.")); + return; + } + Map ruleMap = hits.stream() + .collect(Collectors.toMap(SearchHit::getId, hit -> IndexStoredRuleParser.parseRule(hit.getSourceAsString(), featureType))); + String nextSearchAfter = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); + listener.onResponse(new GetRuleResponse(ruleMap, nextSearchAfter, RestStatus.OK)); + } + + private ThreadContext.StoredContext getContext() { + return client.threadPool().getThreadContext().stashContext(); + } + + public Client getClient() { + return client; + } +} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java index 08856b38cdb08..751607bbc9c92 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java @@ -8,126 +8,10 @@ package org.opensearch.rule.service; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.opensearch.action.search.SearchRequestBuilder; -import org.opensearch.action.search.SearchResponse; -import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; -import org.opensearch.index.query.BoolQueryBuilder; -import org.opensearch.index.query.QueryBuilders; -import org.opensearch.search.sort.SortOrder; -import org.opensearch.transport.client.Client; +import org.opensearch.rule.action.GetRuleRequest; +import org.opensearch.rule.action.GetRuleResponse; -import java.util.Map; -import java.util.Set; - -import static org.opensearch.autotagging.Rule._ID_STRING; - -/** - * This class encapsulates the logic to manage the lifecycle of rules at index level - * @opensearch.experimental - */ public interface RulePersistenceService { - /** - * The default maximum number of results allowed per GET request - */ - int MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST = 50; - - /** - * The system index name used for storing rules - */ - String getIndexName(); - - /** - * client getter - */ - Client getClient(); - - /** - * clusterService getter - */ - ClusterService getClusterService(); - - /** - * Get the max return size allowed for page for get request when pagination is needed - */ - default int getMaxReturnSizeAllowedPerGetRequest() { - return MAX_RETURN_SIZE_ALLOWED_PER_GET_REQUEST; - } - - /** - * logger for RulePersistenceService - */ - Logger logger = LogManager.getLogger(RulePersistenceService.class); - - /** - * Entry point for the get rule api logic in persistence service. If id is provided, we only get a single rule. - * Otherwise, we get all rules that satisfy the attributeFilters. - * @param id - The id of the rule to get. - * @param attributeFilters - A map containing the attributes that user want to filter on - * @param searchAfter - The sort values from the last document of the previous page, used for pagination - * @param listener - ActionListener for GetRuleResponse - */ - default void getRule( - String id, - Map> attributeFilters, - String searchAfter, - ActionListener listener - ) { - // Stash the current thread context when interacting with system index to perform - // operations as the system itself, bypassing authorization checks. This ensures that - // actions within this block are trusted and executed with system-level privileges. - try (ThreadContext.StoredContext context = getContext()) { - BoolQueryBuilder boolQuery = buildGetRuleQuery(id, attributeFilters); - SearchRequestBuilder searchRequest = getClient().prepareSearch(getIndexName()) - .setQuery(boolQuery) - .setSize(getMaxReturnSizeAllowedPerGetRequest()); - if (searchAfter != null) { - searchRequest.addSort(_ID_STRING, SortOrder.ASC).searchAfter(new Object[] { searchAfter }); - } - searchRequest.execute(ActionListener.wrap(listener::onResponse, e -> { - logger.error("Failed to fetch all rules: {}", e.getMessage()); - listener.onFailure(e); - })); - } - } - - /** - * Builds a bool query to retrieve rules from the rules index, applying attribute-based filters - * when needed and ensuring that the rules are associated with the query group feature type. - * @param id The ID of the rule to fetch. If not null, the search will return only this specific rule. - * @param attributeFilters A map of attributes to their associated set of values used to filter the rules. - */ - default BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeFilters) { - BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); - if (id != null) { - return boolQuery.must(QueryBuilders.termQuery(_ID_STRING, id)); - } - for (Map.Entry> entry : attributeFilters.entrySet()) { - Attribute attribute = entry.getKey(); - Set values = entry.getValue(); - if (values != null && !values.isEmpty()) { - BoolQueryBuilder attributeQuery = QueryBuilders.boolQuery(); - for (String value : values) { - attributeQuery.should(QueryBuilders.matchQuery(attribute.getName(), value)); - } - boolQuery.must(attributeQuery); - } - } - boolQuery.filter(QueryBuilders.existsQuery(retrieveFeatureTypeInstance().getName())); - return boolQuery; - } - - /** - * Abstract method for subclasses to provide specific FeatureType Instance - */ - FeatureType retrieveFeatureTypeInstance(); - - private ThreadContext.StoredContext getContext() { - return getClient().threadPool().getThreadContext().stashContext(); - } + void getRule(GetRuleRequest request, ActionListener listener); } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleProcessingService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleProcessingService.java deleted file mode 100644 index 95353a25d2343..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleProcessingService.java +++ /dev/null @@ -1,63 +0,0 @@ -/* - * 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.rule.service; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.opensearch.action.search.SearchResponse; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule; -import org.opensearch.common.collect.Tuple; -import org.opensearch.search.SearchHit; - -import java.util.Arrays; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; - -import static org.opensearch.autotagging.Rule.parseRule; - -/** - * Interface defining the contract for processing and parsing rule-related logic. - * @opensearch.experimental - */ -public interface RuleProcessingService { - /** - * logger for RuleProcessingService - */ - Logger logger = LogManager.getLogger(RuleProcessingService.class); - - /** - * function to retrieve the feature type corresponding to the rule - */ - FeatureType retrieveFeatureTypeInstance(); - - /** - * Process searchResponse from index and send a GetRuleResponse - * @param id - Rule id fetched - * @param searchResponse - Response received from index - */ - default Tuple, String> parseGetRuleResponse(String id, SearchResponse searchResponse) { - List hits = Arrays.asList(searchResponse.getHits().getHits()); - if (id != null && hits.isEmpty()) { - logger.error("Rule with ID " + id + " not found."); - return new Tuple<>(new HashMap<>(), null); - } - Map ruleMap = hits.stream() - .collect( - Collectors.toMap( - SearchHit::getId, - hit -> parseRule(hit.getId(), hit.getSourceAsString(), retrieveFeatureTypeInstance()).getValue() - ) - ); - String lastValidRuleId = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); - return new Tuple<>(ruleMap, lastValidRuleId); - } -} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleResponseBuilder.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleResponseBuilder.java deleted file mode 100644 index 7ef1046effa01..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleResponseBuilder.java +++ /dev/null @@ -1,29 +0,0 @@ -/* - * 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.rule.service; - -import org.opensearch.autotagging.Rule; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.rule.action.GetRuleResponse; - -import java.util.Map; - -/** - * A builder interface for rule lifecycle response objects. - * @opensearch.experimental - */ -public interface RuleResponseBuilder { - /** - * Builds a GetRuleResponse object containing rule data. - * @param ruleMap A map of rule IDs to their corresponding {@link Rule} objects. - * @param nextSearchAfter The ID used for pagination to fetch the next set of results, if applicable. - * @param restStatus The response rest status. - */ - T buildGetRuleResponse(Map ruleMap, String nextSearchAfter, RestStatus restStatus); -} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleService.java deleted file mode 100644 index 57fbcb4a01ff5..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RuleService.java +++ /dev/null @@ -1,74 +0,0 @@ -/* - * 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.rule.service; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.opensearch.action.search.SearchResponse; -import org.opensearch.autotagging.Rule; -import org.opensearch.common.collect.Tuple; -import org.opensearch.common.inject.Inject; -import org.opensearch.core.action.ActionListener; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.rule.action.GetRuleRequest; -import org.opensearch.rule.action.GetRuleResponse; - -import java.util.Map; - -/** - * Service class responsible for processing the business logic for rule lifecycle operations - * @opensearch.experimental - */ -public class RuleService { - private final RulePersistenceService rulePersistenceService; - private final RuleProcessingService ruleProcessingService; - private final RuleResponseBuilder ruleResponseBuilder; - private final Logger logger = LogManager.getLogger(RuleService.class); - - /** - * Constructs a {@link RuleService} with the specified dependencies. - * @param rulePersistenceService The persistence service for retrieving and storing rules. - * @param ruleProcessingService The service responsible for processing rule data. - * @param ruleResponseBuilder The builder used to construct rule response objects. - */ - @Inject - public RuleService( - RulePersistenceService rulePersistenceService, - RuleProcessingService ruleProcessingService, - RuleResponseBuilder ruleResponseBuilder - ) { - this.rulePersistenceService = rulePersistenceService; - this.ruleProcessingService = ruleProcessingService; - this.ruleResponseBuilder = ruleResponseBuilder; - } - - /** - * Processes a request to retrieve rules based on specified criteria. - * @param request The {@link GetRuleRequest} containing rule retrieval parameters. - * @param listener The {@link ActionListener} that handles the asynchronous response. - */ - public void processGetRuleRequest(GetRuleRequest request, ActionListener listener) { - rulePersistenceService.getRule(request.getId(), request.getAttributeFilters(), request.getSearchAfter(), new ActionListener<>() { - @Override - public void onResponse(SearchResponse searchResponse) { - Tuple, String> responseParams = ruleProcessingService.parseGetRuleResponse( - request.getId(), - searchResponse - ); - listener.onResponse(ruleResponseBuilder.buildGetRuleResponse(responseParams.v1(), responseParams.v2(), RestStatus.OK)); - } - - @Override - public void onFailure(Exception e) { - logger.error("Failed to fetch rules: {}", e.getMessage()); - listener.onFailure(e); - } - }); - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java index e65c67b148660..77d938a85fdaa 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java @@ -14,10 +14,8 @@ import org.opensearch.plugin.wlm.rule.service.WlmRulePersistenceService; import org.opensearch.plugin.wlm.rule.service.WlmRuleProcessingService; import org.opensearch.plugin.wlm.rule.service.WlmRuleResponseBuilder; +import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.rule.service.RulePersistenceService; -import org.opensearch.rule.service.RuleProcessingService; -import org.opensearch.rule.service.RuleResponseBuilder; -import org.opensearch.rule.service.RuleService; /** * Guice Module to manage WorkloadManagement related objects @@ -37,5 +35,6 @@ protected void configure() { bind(RuleProcessingService.class).to(WlmRuleProcessingService.class).in(Singleton.class); bind(RuleResponseBuilder.class).to(WlmRuleResponseBuilder.class).in(Singleton.class); bind(RuleService.class).in(Singleton.class); + bind(RulePersistenceService.class).to(IndexStoredRulePersistenceService.class); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java index 95ecbe6d6a4ec..080d417f1319d 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java @@ -9,12 +9,13 @@ package org.opensearch.plugin.wlm.rule.action; import org.opensearch.action.ActionType; +import org.opensearch.rule.action.GetRuleResponse; /** * Action type for getting Rules in workload management * @opensearch.experimental */ -public class GetWlmRuleAction extends ActionType { +public class GetWlmRuleAction extends ActionType { /** * An instance of GetWlmRuleAction @@ -30,6 +31,6 @@ public class GetWlmRuleAction extends ActionType { * Default constructor for GetWlmRuleAction */ private GetWlmRuleAction() { - super(NAME, GetWlmRuleResponse::new); + super(NAME, GetRuleResponse::new); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java deleted file mode 100644 index 97120e9425ed1..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequest.java +++ /dev/null @@ -1,43 +0,0 @@ -/* - * 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.rule.action; - -import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.rule.action.GetRuleRequest; - -import java.io.IOException; -import java.util.Map; -import java.util.Set; - -/** - * A request to get workload management Rules in workload management - * @opensearch.experimental - */ -public class GetWlmRuleRequest extends GetRuleRequest { - /** - * Constructor for GetWlmRuleRequest - * @param id - Rule id to get - * @param attributeFilters - A map containing the attributes to filter on - * @param searchAfter - A string used for pagination - * @param featureType - The featureType instance for wlm - */ - public GetWlmRuleRequest(String id, Map> attributeFilters, String searchAfter, FeatureType featureType) { - super(id, attributeFilters, searchAfter, featureType); - } - - /** - * Constructs a new GetWlmRuleRequest instance from StreamInput - * @param in The {@link StreamInput} from which to read the request data. - */ - public GetWlmRuleRequest(StreamInput in) throws IOException { - super(in); - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponse.java deleted file mode 100644 index c3ab2a80ecf1a..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponse.java +++ /dev/null @@ -1,55 +0,0 @@ -/* - * 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.rule.action; - -import org.opensearch.autotagging.Rule; -import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.rule.action.GetRuleResponse; - -import java.io.IOException; -import java.util.Map; - -/** - * A response to get workload management Rules in workload management - * Example response: - * { - * "rules": [ - * { - * "_id": "z1MJApUB0zgMcDmz-UQq", - * "description": "Rule for tagging query_group_id to index123" - * "index_pattern": ["index123"], - * "query_group": "query_group_id", - * "updated_at": "2025-02-14T01:19:22.589Z" - * }, - * ... - * ], - * "search_after": ["z1MJApUB0zgMcDmz-UQq"] - * } - * @opensearch.experimental - */ -public class GetWlmRuleResponse extends GetRuleResponse { - /** - * Constructor for GetWlmRuleResponse - * @param rules - A map of rule IDs to objects representing the retrieved rules - * @param searchAfter - A string used for pagination - * @param restStatus - The {@link RestStatus} indicating the status of the request. - */ - public GetWlmRuleResponse(Map rules, String searchAfter, RestStatus restStatus) { - super(rules, searchAfter, restStatus); - } - - /** - * Constructs a new GetWlmRuleResponse instance from StreamInput - * @param in The {@link StreamInput} from which to read the response data. - */ - public GetWlmRuleResponse(StreamInput in) throws IOException { - super(in); - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java index 85bc110b61a03..e4bb603f869a7 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java @@ -12,7 +12,10 @@ import org.opensearch.action.support.HandledTransportAction; import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; -import org.opensearch.rule.service.RuleService; +import org.opensearch.rule.action.GetRuleRequest; +import org.opensearch.rule.action.GetRuleResponse; +import org.opensearch.rule.service.IndexStoredRulePersistenceService; +import org.opensearch.rule.service.RulePersistenceService; import org.opensearch.tasks.Task; import org.opensearch.transport.TransportService; @@ -20,24 +23,28 @@ * Transport action to get workload management Rules * @opensearch.experimental */ -public class TransportGetWlmRuleAction extends HandledTransportAction { +public class TransportGetWlmRuleAction extends HandledTransportAction { - private final RuleService ruleService; + private final RulePersistenceService rulePersistenceService; /** * Constructor for TransportGetWlmRuleAction * @param transportService - a {@link TransportService} object * @param actionFilters - a {@link ActionFilters} object - * @param ruleService - a {@link RuleService} object + * @param rulePersistenceService - a {@link RulePersistenceService} object */ @Inject - public TransportGetWlmRuleAction(TransportService transportService, ActionFilters actionFilters, RuleService ruleService) { - super(GetWlmRuleAction.NAME, transportService, actionFilters, GetWlmRuleRequest::new); - this.ruleService = ruleService; + public TransportGetWlmRuleAction( + TransportService transportService, + ActionFilters actionFilters, + RulePersistenceService rulePersistenceService + ) { + super(GetWlmRuleAction.NAME, transportService, actionFilters, GetRuleRequest::new); + this.rulePersistenceService = rulePersistenceService; } @Override - protected void doExecute(Task task, GetWlmRuleRequest request, ActionListener listener) { - ruleService.processGetRuleRequest(request, listener); + protected void doExecute(Task task, GetRuleRequest request, ActionListener listener) { + rulePersistenceService.getRule(request, listener); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java index 55a07082928dd..437dd8b246091 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java @@ -13,7 +13,6 @@ import org.opensearch.plugin.wlm.rule.QueryGroupAttribute; import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; import org.opensearch.plugin.wlm.rule.action.GetWlmRuleAction; -import org.opensearch.plugin.wlm.rule.action.GetWlmRuleRequest; import org.opensearch.rule.action.GetRuleRequest; import org.opensearch.rule.action.GetRuleResponse; import org.opensearch.rule.rest.RestGetRuleAction; @@ -60,6 +59,6 @@ protected > T retrieveGetRuleAct @Override protected GetRuleRequest buildGetRuleRequest(String id, Map> attributeFilters, String searchAfter) { - return new GetWlmRuleRequest(id, attributeFilters, searchAfter, QueryGroupFeatureType.INSTANCE); + return new GetRuleRequest(id, attributeFilters, searchAfter, QueryGroupFeatureType.INSTANCE); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java deleted file mode 100644 index b4d77f3ca0a90..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceService.java +++ /dev/null @@ -1,59 +0,0 @@ -/* - * 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.rule.service; - -import org.opensearch.autotagging.FeatureType; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.inject.Inject; -import org.opensearch.common.settings.Settings; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.rule.service.RulePersistenceService; -import org.opensearch.transport.client.Client; - -/** - * This class encapsulates the logic to manage the lifecycle of workload management rules at index level - * @opensearch.experimental - */ -@SuppressWarnings("unchecked") -public class WlmRulePersistenceService implements RulePersistenceService { - public static final String RULES_INDEX = ".rules"; - private final ClusterService clusterService; - private final Client client; - - @Override - public String getIndexName() { - return RULES_INDEX; - } - - /** - * Constructor for WlmRulePersistenceService - * @param clusterService {@link ClusterService} - The cluster service to be used by RulePersistenceService - * @param client {@link Settings} - The client to be used by RulePersistenceService - */ - @Inject - public WlmRulePersistenceService(ClusterService clusterService, Client client) { - this.client = client; - this.clusterService = clusterService; - } - - @Override - public Client getClient() { - return client; - } - - @Override - public ClusterService getClusterService() { - return clusterService; - } - - @Override - public FeatureType retrieveFeatureTypeInstance() { - return QueryGroupFeatureType.INSTANCE; - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleProcessingService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleProcessingService.java deleted file mode 100644 index 3e50ad02ad5d2..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleProcessingService.java +++ /dev/null @@ -1,20 +0,0 @@ -/* - * 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.rule.service; - -import org.opensearch.autotagging.FeatureType; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.rule.service.RuleProcessingService; - -public class WlmRuleProcessingService implements RuleProcessingService { - @Override - public FeatureType retrieveFeatureTypeInstance() { - return QueryGroupFeatureType.INSTANCE; - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleResponseBuilder.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleResponseBuilder.java deleted file mode 100644 index 819d5542bab5e..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/service/WlmRuleResponseBuilder.java +++ /dev/null @@ -1,25 +0,0 @@ -/* - * 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.rule.service; - -import org.opensearch.autotagging.Rule; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.plugin.wlm.rule.action.GetWlmRuleResponse; -import org.opensearch.rule.action.GetRuleResponse; -import org.opensearch.rule.service.RuleResponseBuilder; - -import java.util.Map; - -@SuppressWarnings("unchecked") -public class WlmRuleResponseBuilder implements RuleResponseBuilder { - @Override - public T buildGetRuleResponse(Map ruleMap, String nextSearchAfter, RestStatus restStatus) { - return (T) new GetWlmRuleResponse(ruleMap, nextSearchAfter, restStatus); - } -} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java index f374ae8f5b703..6a96918633e9b 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java @@ -10,23 +10,12 @@ import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.Rule; -import org.opensearch.cluster.ClusterState; -import org.opensearch.cluster.metadata.Metadata; -import org.opensearch.cluster.metadata.QueryGroup; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.settings.Settings; -import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.plugin.wlm.rule.service.WlmRulePersistenceService; -import org.opensearch.threadpool.ThreadPool; -import org.opensearch.transport.client.Client; import java.util.Map; import java.util.Set; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; public class WlmRuleTestUtils { public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; @@ -61,21 +50,21 @@ public static Map ruleMap() { return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); } - public static WlmRulePersistenceService setUpRulePersistenceService(Map queryGroupMap) { - Client client = mock(Client.class); - ClusterService clusterService = mock(ClusterService.class); - ClusterState clusterState = mock(ClusterState.class); - Metadata metadata = mock(Metadata.class); - ThreadPool threadPool = mock(ThreadPool.class); - - ThreadContext threadContext = new ThreadContext(Settings.EMPTY); - when(client.threadPool()).thenReturn(threadPool); - when(threadPool.getThreadContext()).thenReturn(threadContext); - when(clusterService.state()).thenReturn(clusterState); - when(clusterState.metadata()).thenReturn(metadata); - when(metadata.queryGroups()).thenReturn(queryGroupMap); - return new WlmRulePersistenceService(clusterService, client); - } + // public static WlmRulePersistenceService setUpRulePersistenceService(Map queryGroupMap) { + // Client client = mock(Client.class); + // ClusterService clusterService = mock(ClusterService.class); + // ClusterState clusterState = mock(ClusterState.class); + // Metadata metadata = mock(Metadata.class); + // ThreadPool threadPool = mock(ThreadPool.class); + // + // ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + // when(client.threadPool()).thenReturn(threadPool); + // when(threadPool.getThreadContext()).thenReturn(threadContext); + // when(clusterService.state()).thenReturn(clusterState); + // when(clusterState.metadata()).thenReturn(metadata); + // when(metadata.queryGroups()).thenReturn(queryGroupMap); + // return new WlmRulePersistenceService(clusterService, client); + // } public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { assertEquals(mapOne.size(), mapTwo.size()); diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java deleted file mode 100644 index ccf65d250814f..0000000000000 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleRequestTests.java +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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.rule.action; - -import org.opensearch.common.io.stream.BytesStreamOutput; -import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.test.OpenSearchTestCase; - -import java.io.IOException; -import java.util.HashMap; -import java.util.Map; -import java.util.Set; - -import static org.opensearch.plugin.wlm.rule.QueryGroupAttribute.INDEX_PATTERN; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.PATTERN_ONE; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.SEARCH_AFTER; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils._ID_ONE; - -public class GetWlmRuleRequestTests extends OpenSearchTestCase { - - /** - * Test case to verify the serialization and deserialization of GetRuleRequest - */ - public void testSerialization() throws IOException { - GetWlmRuleRequest request = new GetWlmRuleRequest( - _ID_ONE, - Map.of(INDEX_PATTERN, Set.of(PATTERN_ONE)), - null, - QueryGroupFeatureType.INSTANCE - ); - assertEquals(_ID_ONE, request.getId()); - BytesStreamOutput out = new BytesStreamOutput(); - request.writeTo(out); - StreamInput streamInput = out.bytes().streamInput(); - GetWlmRuleRequest otherRequest = new GetWlmRuleRequest(streamInput); - assertEquals(request.getId(), otherRequest.getId()); - assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters()); - } - - /** - * Test case to verify the serialization and deserialization of GetRuleRequest when name is null - */ - public void testSerializationWithNull() throws IOException { - GetWlmRuleRequest request = new GetWlmRuleRequest((String) null, new HashMap<>(), SEARCH_AFTER, QueryGroupFeatureType.INSTANCE); - assertNull(request.getId()); - BytesStreamOutput out = new BytesStreamOutput(); - request.writeTo(out); - StreamInput streamInput = out.bytes().streamInput(); - GetWlmRuleRequest otherRequest = new GetWlmRuleRequest(streamInput); - assertEquals(request.getId(), otherRequest.getId()); - assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters()); - } -} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponseTests.java deleted file mode 100644 index 76f1f7a0a2534..0000000000000 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleResponseTests.java +++ /dev/null @@ -1,123 +0,0 @@ -/* - * 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.rule.action; - -import org.opensearch.autotagging.Rule; -import org.opensearch.common.io.stream.BytesStreamOutput; -import org.opensearch.common.xcontent.json.JsonXContent; -import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.core.xcontent.ToXContent; -import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.test.OpenSearchTestCase; - -import java.io.IOException; -import java.util.HashMap; -import java.util.Map; - -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.SEARCH_AFTER; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils._ID_ONE; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.assertEqualRules; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.ruleMap; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.ruleOne; -import static org.mockito.Mockito.mock; - -public class GetWlmRuleResponseTests extends OpenSearchTestCase { - - /** - * Test case to verify the serialization and deserialization of GetRuleResponse - */ - public void testSerializationSingleRule() throws IOException { - Map map = new HashMap<>(); - map.put(_ID_ONE, ruleOne); - GetWlmRuleResponse response = new GetWlmRuleResponse(Map.of(_ID_ONE, ruleOne), null, RestStatus.OK); - assertEquals(response.getRules(), map); - - BytesStreamOutput out = new BytesStreamOutput(); - response.writeTo(out); - StreamInput streamInput = out.bytes().streamInput(); - - GetWlmRuleResponse otherResponse = new GetWlmRuleResponse(streamInput); - assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); - assertEqualRules(response.getRules(), otherResponse.getRules(), false); - } - - /** - * Test case to verify the serialization and deserialization of GetRuleResponse when the result contains multiple rules - */ - public void testSerializationMultipleRule() throws IOException { - GetWlmRuleResponse response = new GetWlmRuleResponse(ruleMap(), SEARCH_AFTER, RestStatus.OK); - assertEquals(response.getRules(), ruleMap()); - - BytesStreamOutput out = new BytesStreamOutput(); - response.writeTo(out); - StreamInput streamInput = out.bytes().streamInput(); - - GetWlmRuleResponse otherResponse = new GetWlmRuleResponse(streamInput); - assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); - assertEquals(2, otherResponse.getRules().size()); - assertEqualRules(response.getRules(), otherResponse.getRules(), false); - } - - /** - * Test case to verify the serialization and deserialization of GetRuleResponse when the result is empty - */ - public void testSerializationNull() throws IOException { - Map map = new HashMap<>(); - GetWlmRuleResponse response = new GetWlmRuleResponse(map, SEARCH_AFTER, RestStatus.OK); - assertEquals(response.getRules(), map); - - BytesStreamOutput out = new BytesStreamOutput(); - response.writeTo(out); - StreamInput streamInput = out.bytes().streamInput(); - - GetWlmRuleResponse otherResponse = new GetWlmRuleResponse(streamInput); - assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); - assertEquals(0, otherResponse.getRules().size()); - } - - /** - * Test case to verify the toXContent of GetRuleResponse - */ - public void testToXContentGetSingleRule() throws IOException { - Map map = new HashMap<>(); - map.put(_ID_ONE, ruleOne); - GetWlmRuleResponse response = new GetWlmRuleResponse(Map.of(_ID_ONE, ruleOne), SEARCH_AFTER, RestStatus.OK); - XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); - String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); - String expected = "{\n" - + " \"rules\" : [\n" - + " {\n" - + " \"_id\" : \"AgfUO5Ja9yfvhdONlYi3TQ==\",\n" - + " \"description\" : \"description_1\",\n" - + " \"index_pattern\" : [\n" - + " \"pattern_1\"\n" - + " ],\n" - + " \"query_group\" : \"feature_value_one\",\n" - + " \"updated_at\" : \"2024-01-26T08:58:57.558Z\"\n" - + " }\n" - + " ],\n" - + " \"search_after\" : [\n" - + " \"search_after_id\"\n" - + " ]\n" - + "}"; - assertEquals(expected, actual); - } - - /** - * Test case to verify toXContent of GetRuleResponse when the result contains zero Rule - */ - public void testToXContentGetZeroRule() throws IOException { - XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); - GetWlmRuleResponse otherResponse = new GetWlmRuleResponse(new HashMap<>(), null, RestStatus.OK); - String actual = otherResponse.toXContent(builder, mock(ToXContent.Params.class)).toString(); - String expected = "{\n" + " \"rules\" : [ ]\n" + "}"; - assertEquals(expected, actual); - } -} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java index 30d3359078140..6df56b15ed875 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java @@ -8,83 +8,60 @@ package org.opensearch.plugin.wlm.rule.service; -import org.opensearch.action.search.SearchRequestBuilder; -import org.opensearch.action.search.SearchResponse; -import org.opensearch.core.action.ActionListener; -import org.opensearch.index.query.BoolQueryBuilder; -import org.opensearch.index.query.QueryBuilders; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.search.sort.SortOrder; import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.transport.client.Client; - -import java.util.HashMap; - -import static org.opensearch.autotagging.Rule._ID_STRING; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.ATTRIBUTE_MAP; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils._ID_ONE; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils._ID_TWO; -import static org.opensearch.plugin.wlm.rule.WlmRuleTestUtils.setUpRulePersistenceService; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.anyInt; -import static org.mockito.Mockito.anyString; -import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; @SuppressWarnings("unchecked") public class WlmRulePersistenceServiceTests extends OpenSearchTestCase { - public void testBuildGetRuleQuery_WithId() { - WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); - BoolQueryBuilder query = rulePersistenceService.buildGetRuleQuery(_ID_ONE, new HashMap<>()); - assertTrue(query.hasClauses()); - assertEquals(QueryBuilders.termQuery(_ID_STRING, _ID_ONE).toString(), query.must().get(0).toString()); - } - - public void testBuildGetRuleQuery_WithFilters() { - WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); - BoolQueryBuilder query = rulePersistenceService.buildGetRuleQuery(null, ATTRIBUTE_MAP); - assertTrue(query.hasClauses()); - assertEquals(1, query.must().size()); - assertTrue(query.filter().contains(QueryBuilders.existsQuery(QueryGroupFeatureType.NAME))); - } - - public void testGetRule_WithId() { - WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); - Client client = rulePersistenceService.getClient(); - ActionListener listener = mock(ActionListener.class); - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - SetupMocksForGetRule(client, searchRequestBuilder); - - rulePersistenceService.getRule(_ID_ONE, new HashMap<>(), null, listener); - verify(client).prepareSearch(WlmRulePersistenceService.RULES_INDEX); - verify(searchRequestBuilder).setQuery(any()); - verify(searchRequestBuilder).execute(any()); - } - - public void testGetRule_WithSearchAfter() { - WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); - Client client = rulePersistenceService.getClient(); - ActionListener listener = mock(ActionListener.class); - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - SetupMocksForGetRule(client, searchRequestBuilder); - when(searchRequestBuilder.addSort(anyString(), any(SortOrder.class))).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.searchAfter(any())).thenReturn(searchRequestBuilder); - - rulePersistenceService.getRule(null, new HashMap<>(), _ID_TWO, listener); - verify(searchRequestBuilder).addSort(_ID_STRING, SortOrder.ASC); - verify(searchRequestBuilder).searchAfter(new Object[] { _ID_TWO }); - } - - public void SetupMocksForGetRule(Client client, SearchRequestBuilder searchRequestBuilder) { - when(client.prepareSearch(anyString())).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.setQuery(any())).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.setSize(anyInt())).thenReturn(searchRequestBuilder); - doAnswer(invocation -> { - ActionListener searchListener = invocation.getArgument(0); - searchListener.onResponse(mock(SearchResponse.class)); - return null; - }).when(searchRequestBuilder).execute(any()); - } + // public void testBuildGetRuleQuery_WithId() { + // WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); + // BoolQueryBuilder query = rulePersistenceService.buildGetRuleQuery(_ID_ONE, new HashMap<>()); + // assertTrue(query.hasClauses()); + // assertEquals(QueryBuilders.termQuery(_ID_STRING, _ID_ONE).toString(), query.must().get(0).toString()); + // } + // + // public void testBuildGetRuleQuery_WithFilters() { + // WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); + // BoolQueryBuilder query = rulePersistenceService.buildGetRuleQuery(null, ATTRIBUTE_MAP); + // assertTrue(query.hasClauses()); + // assertEquals(1, query.must().size()); + // assertTrue(query.filter().contains(QueryBuilders.existsQuery(QueryGroupFeatureType.NAME))); + // } + // + // public void testGetRule_WithId() { + // WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); + // Client client = rulePersistenceService.getClient(); + // ActionListener listener = mock(ActionListener.class); + // SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + // SetupMocksForGetRule(client, searchRequestBuilder); + // + // rulePersistenceService.getRule(_ID_ONE, new HashMap<>(), null, listener); + // verify(client).prepareSearch(WlmRulePersistenceService.RULES_INDEX); + // verify(searchRequestBuilder).setQuery(any()); + // verify(searchRequestBuilder).execute(any()); + // } + // + // public void testGetRule_WithSearchAfter() { + // WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); + // Client client = rulePersistenceService.getClient(); + // ActionListener listener = mock(ActionListener.class); + // SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + // SetupMocksForGetRule(client, searchRequestBuilder); + // when(searchRequestBuilder.addSort(anyString(), any(SortOrder.class))).thenReturn(searchRequestBuilder); + // when(searchRequestBuilder.searchAfter(any())).thenReturn(searchRequestBuilder); + // + // rulePersistenceService.getRule(null, new HashMap<>(), _ID_TWO, listener); + // verify(searchRequestBuilder).addSort(_ID_STRING, SortOrder.ASC); + // verify(searchRequestBuilder).searchAfter(new Object[] { _ID_TWO }); + // } + // + // public void SetupMocksForGetRule(Client client, SearchRequestBuilder searchRequestBuilder) { + // when(client.prepareSearch(anyString())).thenReturn(searchRequestBuilder); + // when(searchRequestBuilder.setQuery(any())).thenReturn(searchRequestBuilder); + // when(searchRequestBuilder.setSize(anyInt())).thenReturn(searchRequestBuilder); + // doAnswer(invocation -> { + // ActionListener searchListener = invocation.getArgument(0); + // searchListener.onResponse(mock(SearchResponse.class)); + // return null; + // }).when(searchRequestBuilder).execute(any()); + // } } From 5b48e3caa33264bab3ee246bc5116ed33ee34ff5 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Thu, 3 Apr 2025 13:17:23 -0700 Subject: [PATCH 09/20] modify based on comments Signed-off-by: Ruirui Zhang --- .../IndexStoredRulePersistenceService.java | 37 +++-- .../rule/service/RulePersistenceService.java | 10 ++ .../IndexStoredRuleParser.java | 13 +- .../IndexStoredRuleUtils.java | 18 ++- .../rule/{Utils => utils}/package-info.java | 2 +- .../org/opensearch/rule/autotagging/Rule.java | 6 - .../org/opensearch/rule/RuleTestUtils.java | 153 ++++++++++++++++++ .../rule/action/GetRuleRequestTests.java | 89 +--------- .../rule/action/GetRuleResponseTests.java | 51 +++--- ...ndexStoredRulePersistenceServiceTests.java | 109 +++---------- .../utils/IndexStoredRuleParserTests.java | 54 +++++++ .../rule/utils/IndexStoredRuleUtilsTests.java | 42 +++++ .../opensearch/rule/utils/RuleTestUtils.java | 79 --------- .../plugin/wlm/WorkloadManagementPlugin.java | 8 +- .../action/TransportGetWlmRuleAction.java | 1 - .../plugin/wlm/rule/WlmRuleTestUtils.java | 91 ----------- .../WlmRulePersistenceServiceTests.java | 67 -------- 17 files changed, 365 insertions(+), 465 deletions(-) rename libs/autotagging-commons/src/main/java/org/opensearch/rule/{Utils => utils}/IndexStoredRuleParser.java (82%) rename libs/autotagging-commons/src/main/java/org/opensearch/rule/{Utils => utils}/IndexStoredRuleUtils.java (69%) rename libs/autotagging-commons/src/main/java/org/opensearch/rule/{Utils => utils}/package-info.java (89%) create mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java create mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java create mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java delete mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java delete mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java delete mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java index 9ca715e4cecc0..52cd6e97a75b7 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -20,10 +20,10 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; import org.opensearch.index.query.BoolQueryBuilder; -import org.opensearch.rule.Utils.IndexStoredRuleParser; -import org.opensearch.rule.Utils.IndexStoredRuleUtils; import org.opensearch.rule.action.GetRuleRequest; import org.opensearch.rule.action.GetRuleResponse; +import org.opensearch.rule.utils.IndexStoredRuleParser; +import org.opensearch.rule.utils.IndexStoredRuleUtils; import org.opensearch.search.SearchHit; import org.opensearch.search.sort.SortOrder; import org.opensearch.transport.client.Client; @@ -47,27 +47,35 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService private final String indexName; private final Client client; private final FeatureType featureType; - private final int maxReturnSizeAllowedPerGetRequest; + private final int maxRulesPerGetRequest; private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class); - public IndexStoredRulePersistenceService( - String indexName, - Client client, - FeatureType featureType, - int maxReturnSizeAllowedPerGetRequest - ) { + /** + * Constructs an instance of {@link IndexStoredRulePersistenceService} with the specified parameters. + * This service handles persistence and retrieval of stored rules within an OpenSearch index. + * @param indexName - The name of the OpenSearch index where the rules are stored. + * @param client - The OpenSearch client used to interact with the OpenSearch cluster. + * @param featureType - The feature type associated with the stored rules. + * @param maxRulesPerGetRequest - The maximum number of rules that can be returned in a single get request. + */ + public IndexStoredRulePersistenceService(String indexName, Client client, FeatureType featureType, int maxRulesPerGetRequest) { this.indexName = indexName; this.client = client; this.featureType = featureType; - this.maxReturnSizeAllowedPerGetRequest = maxReturnSizeAllowedPerGetRequest; + this.maxRulesPerGetRequest = maxRulesPerGetRequest; } + /** + * Entry point for the get rule api logic in persistence service. + * @param getRuleRequest the getRuleRequest to process. + * @param listener the listener for GetRuleResponse. + */ public void getRule(GetRuleRequest getRuleRequest, ActionListener listener) { getRuleFromIndex(getRuleRequest.getId(), getRuleRequest.getAttributeFilters(), getRuleRequest.getSearchAfter(), listener); } /** - * Entry point for the get rule api logic in persistence service. If id is provided, we only get a single rule. + * Get rules from index. If id is provided, we only get a single rule. * Otherwise, we get all rules that satisfy the attributeFilters. * @param id - The id of the rule to get. * @param attributeFilters - A map containing the attributes that user want to filter on @@ -85,9 +93,7 @@ public void getRuleFromIndex( // actions within this block are trusted and executed with system-level privileges. try (ThreadContext.StoredContext context = getContext()) { BoolQueryBuilder boolQuery = IndexStoredRuleUtils.buildGetRuleQuery(id, attributeFilters, featureType); - SearchRequestBuilder searchRequest = client.prepareSearch(indexName) - .setQuery(boolQuery) - .setSize(maxReturnSizeAllowedPerGetRequest); + SearchRequestBuilder searchRequest = client.prepareSearch(indexName).setQuery(boolQuery).setSize(maxRulesPerGetRequest); if (searchAfter != null) { searchRequest.addSort(_ID_STRING, SortOrder.ASC).searchAfter(new Object[] { searchAfter }); } @@ -120,6 +126,9 @@ private ThreadContext.StoredContext getContext() { return client.threadPool().getThreadContext().stashContext(); } + /** + * client getter + */ public Client getClient() { return client; } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java index 751607bbc9c92..adcccfc429210 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java @@ -12,6 +12,16 @@ import org.opensearch.rule.action.GetRuleRequest; import org.opensearch.rule.action.GetRuleResponse; +/** + * Interface for a service that handles rule persistence CRUD operations. + * @opensearch.experimental + */ public interface RulePersistenceService { + + /** + * Get rules based on the provided request. + * @param request The request containing the details for retrieving the rule. + * @param listener The listener that will handle the response or failure. + */ void getRule(GetRuleRequest request, ActionListener listener); } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleParser.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java similarity index 82% rename from libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleParser.java rename to libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java index 0e921b7619c60..4d2e808a22c2d 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleParser.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.rule.Utils; +package org.opensearch.rule.utils; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -19,12 +19,23 @@ import java.io.IOException; +/** + * Utility class for parsing index stored rules into Rule objects. + * @opensearch.experimental + */ public class IndexStoredRuleParser { + + /** + * constructor for IndexStoredRuleParser + */ + public IndexStoredRuleParser() {} + private static final Logger logger = LogManager.getLogger(IndexStoredRuleParser.class); /** * Parses a source string into a Rule object * @param source - The raw source string representing the rule to be parsed + * @param featureType - The feature type to associate with the parsed rule */ public static Rule parseRule(String source, FeatureType featureType) { try ( diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleUtils.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java similarity index 69% rename from libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleUtils.java rename to libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java index e9f29fbe3c7c0..d43557a81b4f5 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/IndexStoredRuleUtils.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.rule.Utils; +package org.opensearch.rule.utils; import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.FeatureType; @@ -18,7 +18,23 @@ import static org.opensearch.autotagging.Rule._ID_STRING; +/** + * Utility class that provides methods for the lifecycle of rules. + * @opensearch.experimental + */ public class IndexStoredRuleUtils { + + /** + * constructor for IndexStoredRuleUtils + */ + public IndexStoredRuleUtils() {} + + /** + * Builds a Boolean query to retrieve a rule by its ID or attribute filters. + * @param id The ID of the rule to search for. If null, no ID-based filtering is applied. + * @param attributeFilters A map of attributes and their corresponding filter values. This allows filtering by specific attribute values. + * @param featureType The feature type that is required in the query. + */ public static BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeFilters, FeatureType featureType) { BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); if (id != null) { diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/package-info.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/package-info.java similarity index 89% rename from libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/package-info.java rename to libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/package-info.java index 12990034ed50d..5b2961fb3fb2e 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/Utils/package-info.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/package-info.java @@ -9,4 +9,4 @@ /** * Package for the service classes related to query groups in WorkloadManagementPlugin */ -package org.opensearch.rule.Utils; +package org.opensearch.rule.utils; diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/Rule.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/Rule.java index f4f3d524ff642..e1b003f0085c6 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/Rule.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/Rule.java @@ -8,14 +8,9 @@ package org.opensearch.rule.autotagging; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; 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 org.opensearch.core.xcontent.DeprecationHandler; -import org.opensearch.core.xcontent.MediaTypeRegistry; -import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParseException; @@ -61,7 +56,6 @@ public class Rule implements Writeable, ToXContentObject { * updated_at field */ public static final String UPDATED_AT_STRING = "updated_at"; - public static final Logger logger = LogManager.getLogger(Rule.class); /** * Main constructor diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java new file mode 100644 index 0000000000000..14388d65a5e20 --- /dev/null +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java @@ -0,0 +1,153 @@ +/* + * 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.rule; + +import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.AutoTaggingRegistry; +import org.opensearch.autotagging.FeatureType; +import org.opensearch.autotagging.Rule; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.rule.service.IndexStoredRulePersistenceService; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.client.Client; + +import java.util.Map; +import java.util.Set; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class RuleTestUtils { + public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; + public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; + public static final String FEATURE_VALUE_ONE = "feature_value_one"; + public static final String FEATURE_VALUE_TWO = "feature_value_two"; + public static final String PATTERN_ONE = "pattern_1"; + public static final String PATTERN_TWO = "pattern_2"; + public static final String DESCRIPTION_ONE = "description_1"; + public static final String DESCRIPTION_TWO = "description_2"; + public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; + public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; + public static final String SEARCH_AFTER = "search_after_id"; + public static final String TEST_INDEX_NAME = ".test_index_for_rule"; + public static final Map> ATTRIBUTE_MAP = Map.of(MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, Set.of(PATTERN_ONE)); + public static final Rule ruleOne = Rule.builder() + .description(DESCRIPTION_ONE) + .featureType(MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_ONE) + .attributeMap(ATTRIBUTE_MAP) + .updatedAt(TIMESTAMP_ONE) + .build(); + + public static final Rule ruleTwo = Rule.builder() + .description(DESCRIPTION_TWO) + .featureType(MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_TWO) + .attributeMap(Map.of(MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(PATTERN_TWO))) + .updatedAt(TIMESTAMP_TWO) + .build(); + + public static Map ruleMap() { + return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); + } + + public static IndexStoredRulePersistenceService setUpIndexStoredRulePersistenceService(Map queryGroupMap) { + Client client = mock(Client.class); + ClusterService clusterService = mock(ClusterService.class); + ClusterState clusterState = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + ThreadPool threadPool = mock(ThreadPool.class); + + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + when(client.threadPool()).thenReturn(threadPool); + when(threadPool.getThreadContext()).thenReturn(threadContext); + when(clusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + when(metadata.queryGroups()).thenReturn(queryGroupMap); + return new IndexStoredRulePersistenceService(TEST_INDEX_NAME, client, MockRuleFeatureType.INSTANCE, 50); + } + + public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { + assertEquals(mapOne.size(), mapTwo.size()); + for (Map.Entry entry : mapOne.entrySet()) { + String id = entry.getKey(); + assertTrue(mapTwo.containsKey(id)); + Rule one = mapOne.get(id); + Rule two = mapTwo.get(id); + assertEqualRule(one, two, ruleUpdated); + } + } + + public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { + if (ruleUpdated) { + assertEquals(one.getDescription(), two.getDescription()); + assertEquals(one.getFeatureType(), two.getFeatureType()); + assertEquals(one.getFeatureValue(), two.getFeatureValue()); + assertEquals(one.getAttributeMap(), two.getAttributeMap()); + assertEquals(one.getAttributeMap(), two.getAttributeMap()); + } else { + assertEquals(one, two); + } + } + + public static class MockRuleFeatureType implements FeatureType { + + public static final MockRuleFeatureType INSTANCE = new MockRuleFeatureType(); + + private MockRuleFeatureType() {} + + static { + INSTANCE.registerFeatureType(); + } + + @Override + public String getName() { + return "mock_feature_type"; + } + + @Override + public Map getAllowedAttributesRegistry() { + return Map.of( + "mock_attribute_one", + MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, + "mock_attribute_two", + MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO + ); + } + + @Override + public void registerFeatureType() { + AutoTaggingRegistry.registerFeatureType(INSTANCE); + } + } + + public enum MockRuleAttributes implements Attribute { + MOCK_RULE_ATTRIBUTE_ONE("mock_attribute_one"), + MOCK_RULE_ATTRIBUTE_TWO("mock_attribute_two"); + ; + + private final String name; + + MockRuleAttributes(String name) { + this.name = name; + } + + @Override + public String getName() { + return name; + } + } +} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java index a451a58356606..61f452e26f624 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java @@ -10,27 +10,25 @@ import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.rule.GetRuleRequest; -import org.opensearch.rule.autotagging.Attribute; -import org.opensearch.rule.autotagging.Rule; -import org.opensearch.rule.utils.RuleTestUtils; +import org.opensearch.rule.RuleTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; import java.util.HashMap; -import java.util.Map; -import java.util.Set; + +import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_MAP; +import static org.opensearch.rule.RuleTestUtils.MockRuleFeatureType; +import static org.opensearch.rule.RuleTestUtils.SEARCH_AFTER; +import static org.opensearch.rule.RuleTestUtils._ID_ONE; public class GetRuleRequestTests extends OpenSearchTestCase { + /** * Test case to verify the serialization and deserialization of GetRuleRequest */ public void testSerialization() throws IOException { GetRuleRequest request = new GetRuleRequest(_ID_ONE, ATTRIBUTE_MAP, null, RuleTestUtils.MockRuleFeatureType.INSTANCE); assertEquals(_ID_ONE, request.getId()); - assertNull(request.validate()); - assertNull(request.getSearchAfter()); - assertEquals(RuleTestUtils.MockRuleFeatureType.INSTANCE, request.getFeatureType()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); @@ -43,12 +41,7 @@ public void testSerialization() throws IOException { * Test case to verify the serialization and deserialization of GetRuleRequest when name is null */ public void testSerializationWithNull() throws IOException { - GetRuleRequest request = new GetRuleRequest( - (String) null, - new HashMap<>(), - SEARCH_AFTER, - RuleTestUtils.MockRuleFeatureType.INSTANCE - ); + GetRuleRequest request = new GetRuleRequest((String) null, new HashMap<>(), SEARCH_AFTER, MockRuleFeatureType.INSTANCE); assertNull(request.getId()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); @@ -57,70 +50,4 @@ public void testSerializationWithNull() throws IOException { assertEquals(request.getId(), otherRequest.getId()); assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters()); } - - public void testValidate() { - GetRuleRequest request = new GetRuleRequest("", ATTRIBUTE_MAP, null, RuleTestUtils.MockRuleFeatureType.INSTANCE); - assertThrows(IllegalArgumentException.class, request::validate); - request = new GetRuleRequest(_ID_ONE, ATTRIBUTE_MAP, "", RuleTestUtils.MockRuleFeatureType.INSTANCE); - assertThrows(IllegalArgumentException.class, request::validate); - } - - public static final String _ID_ONE = "id_1"; - public static final String SEARCH_AFTER = "search_after"; - public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; - public static final String FEATURE_VALUE_ONE = "feature_value_one"; - public static final String FEATURE_VALUE_TWO = "feature_value_two"; - public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; - public static final String ATTRIBUTE_VALUE_TWO = "mock_attribute_two"; - public static final String DESCRIPTION_ONE = "description_1"; - public static final String DESCRIPTION_TWO = "description_2"; - public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; - public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; - public static final Map> ATTRIBUTE_MAP = Map.of( - RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, - Set.of(ATTRIBUTE_VALUE_ONE) - ); - - public static final Rule ruleOne = Rule.builder() - .description(DESCRIPTION_ONE) - .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_ONE) - .attributeMap(ATTRIBUTE_MAP) - .updatedAt(TIMESTAMP_ONE) - .build(); - - public static final Rule ruleTwo = Rule.builder() - .description(DESCRIPTION_TWO) - .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_TWO) - .attributeMap(Map.of(RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(ATTRIBUTE_VALUE_TWO))) - .updatedAt(TIMESTAMP_TWO) - .build(); - - public static Map ruleMap() { - return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); - } - - public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { - assertEquals(mapOne.size(), mapTwo.size()); - for (Map.Entry entry : mapOne.entrySet()) { - String id = entry.getKey(); - assertTrue(mapTwo.containsKey(id)); - Rule one = mapOne.get(id); - Rule two = mapTwo.get(id); - assertEqualRule(one, two, ruleUpdated); - } - } - - public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { - if (ruleUpdated) { - assertEquals(one.getDescription(), two.getDescription()); - assertEquals(one.getFeatureType(), two.getFeatureType()); - assertEquals(one.getFeatureValue(), two.getFeatureValue()); - assertEquals(one.getAttributeMap(), two.getAttributeMap()); - assertEquals(one.getAttributeMap(), two.getAttributeMap()); - } else { - assertEquals(one, two); - } - } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java index f01bb94fab276..5619bf4ee8567 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java @@ -8,45 +8,27 @@ package org.opensearch.rule.action; +import org.opensearch.autotagging.Rule; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.rule.GetRuleResponse; -import org.opensearch.rule.autotagging.Attribute; -import org.opensearch.rule.autotagging.Rule; -import org.opensearch.rule.utils.RuleTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; import java.util.HashMap; import java.util.Map; -import java.util.Set; -import static org.opensearch.rule.action.GetRuleRequestTests.SEARCH_AFTER; -import static org.opensearch.rule.action.GetRuleRequestTests._ID_ONE; -import static org.opensearch.rule.action.GetRuleRequestTests.assertEqualRules; -import static org.opensearch.rule.action.GetRuleRequestTests.ruleMap; +import static org.opensearch.rule.RuleTestUtils.SEARCH_AFTER; +import static org.opensearch.rule.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.RuleTestUtils.assertEqualRules; +import static org.opensearch.rule.RuleTestUtils.ruleMap; +import static org.opensearch.rule.RuleTestUtils.ruleOne; import static org.mockito.Mockito.mock; public class GetRuleResponseTests extends OpenSearchTestCase { - public static final String FEATURE_VALUE_ONE = "feature_value_one"; - public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; - public static final String DESCRIPTION_ONE = "description_1"; - public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; - static final Map> ATTRIBUTE_MAP = Map.of( - RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, - Set.of(ATTRIBUTE_VALUE_ONE) - ); - - public static final Rule ruleOne = Rule.builder() - .description(DESCRIPTION_ONE) - .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_ONE) - .attributeMap(ATTRIBUTE_MAP) - .updatedAt(TIMESTAMP_ONE) - .build(); /** * Test case to verify the serialization and deserialization of GetRuleResponse @@ -54,7 +36,7 @@ public class GetRuleResponseTests extends OpenSearchTestCase { public void testSerializationSingleRule() throws IOException { Map map = new HashMap<>(); map.put(_ID_ONE, ruleOne); - GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), null); + GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), null, RestStatus.OK); assertEquals(response.getRules(), map); BytesStreamOutput out = new BytesStreamOutput(); @@ -62,6 +44,7 @@ public void testSerializationSingleRule() throws IOException { StreamInput streamInput = out.bytes().streamInput(); GetRuleResponse otherResponse = new GetRuleResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); assertEqualRules(response.getRules(), otherResponse.getRules(), false); } @@ -69,7 +52,7 @@ public void testSerializationSingleRule() throws IOException { * Test case to verify the serialization and deserialization of GetRuleResponse when the result contains multiple rules */ public void testSerializationMultipleRule() throws IOException { - GetRuleResponse response = new GetRuleResponse(ruleMap(), SEARCH_AFTER); + GetRuleResponse response = new GetRuleResponse(ruleMap(), SEARCH_AFTER, RestStatus.OK); assertEquals(response.getRules(), ruleMap()); BytesStreamOutput out = new BytesStreamOutput(); @@ -77,6 +60,7 @@ public void testSerializationMultipleRule() throws IOException { StreamInput streamInput = out.bytes().streamInput(); GetRuleResponse otherResponse = new GetRuleResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); assertEquals(2, otherResponse.getRules().size()); assertEqualRules(response.getRules(), otherResponse.getRules(), false); } @@ -86,7 +70,7 @@ public void testSerializationMultipleRule() throws IOException { */ public void testSerializationNull() throws IOException { Map map = new HashMap<>(); - GetRuleResponse response = new GetRuleResponse(map, SEARCH_AFTER); + GetRuleResponse response = new GetRuleResponse(map, SEARCH_AFTER, RestStatus.OK); assertEquals(response.getRules(), map); BytesStreamOutput out = new BytesStreamOutput(); @@ -94,6 +78,7 @@ public void testSerializationNull() throws IOException { StreamInput streamInput = out.bytes().streamInput(); GetRuleResponse otherResponse = new GetRuleResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); assertEquals(0, otherResponse.getRules().size()); } @@ -103,23 +88,23 @@ public void testSerializationNull() throws IOException { public void testToXContentGetSingleRule() throws IOException { Map map = new HashMap<>(); map.put(_ID_ONE, ruleOne); - GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), SEARCH_AFTER); + GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), SEARCH_AFTER, RestStatus.OK); XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); String expected = "{\n" + " \"rules\" : [\n" + " {\n" - + " \"_id\" : \"id_1\",\n" + + " \"_id\" : \"AgfUO5Ja9yfvhdONlYi3TQ==\",\n" + " \"description\" : \"description_1\",\n" + " \"mock_attribute_one\" : [\n" - + " \"mock_attribute_one\"\n" + + " \"pattern_1\"\n" + " ],\n" + " \"mock_feature_type\" : \"feature_value_one\",\n" + " \"updated_at\" : \"2024-01-26T08:58:57.558Z\"\n" + " }\n" + " ],\n" + " \"search_after\" : [\n" - + " \"search_after\"\n" + + " \"search_after_id\"\n" + " ]\n" + "}"; assertEquals(expected, actual); @@ -130,7 +115,7 @@ public void testToXContentGetSingleRule() throws IOException { */ public void testToXContentGetZeroRule() throws IOException { XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); - GetRuleResponse otherResponse = new GetRuleResponse(new HashMap<>(), null); + GetRuleResponse otherResponse = new GetRuleResponse(new HashMap<>(), null, RestStatus.OK); String actual = otherResponse.toXContent(builder, mock(ToXContent.Params.class)).toString(); String expected = "{\n" + " \"rules\" : [ ]\n" + "}"; assertEquals(expected, actual); diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java index 09a693dfd2b4e..73a3b1cd962af 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java @@ -33,10 +33,14 @@ import org.opensearch.rule.RuleQueryMapper; import org.opensearch.rule.autotagging.Rule; import org.opensearch.rule.utils.RuleTestUtils; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.bytes.BytesArray; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.rule.action.GetRuleResponse; import org.opensearch.search.SearchHit; import org.opensearch.search.SearchHits; import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.Client; import java.util.HashMap; @@ -44,14 +48,18 @@ import org.mockito.ArgumentCaptor; import static org.opensearch.rule.XContentRuleParserTests.VALID_JSON; -import static org.opensearch.rule.utils.RuleTestUtils.TEST_INDEX_NAME; -import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.RuleTestUtils.TEST_INDEX_NAME; +import static org.opensearch.rule.RuleTestUtils._ID_ONE; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.argThat; +import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_MAP; +import static org.opensearch.rule.RuleTestUtils.TEST_INDEX_NAME; +import static org.opensearch.rule.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.RuleTestUtils.setUpIndexStoredRulePersistenceService; +import static org.opensearch.rule.utils.IndexStoredRuleParserTests.VALID_JSON; import static org.mockito.Mockito.any; import static org.mockito.Mockito.anyInt; -import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -59,32 +67,8 @@ @SuppressWarnings("unchecked") public class IndexStoredRulePersistenceServiceTests extends OpenSearchTestCase { - public static final int MAX_VALUES_PER_PAGE = 50; - public void testGetRuleByIdSuccess() { - GetRuleRequest getRuleRequest = mock(GetRuleRequest.class); - when(getRuleRequest.getId()).thenReturn(_ID_ONE); - when(getRuleRequest.getAttributeFilters()).thenReturn(new HashMap<>()); - QueryBuilder queryBuilder = mock(QueryBuilder.class); - RuleQueryMapper mockRuleQueryMapper = mock(RuleQueryMapper.class); - RuleEntityParser mockRuleEntityParser = mock(RuleEntityParser.class); - Rule mockRule = mock(Rule.class); - - when(mockRuleEntityParser.parse(anyString())).thenReturn(mockRule); - when(mockRuleQueryMapper.from(getRuleRequest)).thenReturn(queryBuilder); - when(queryBuilder.filter(any())).thenReturn(queryBuilder); - - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - Client client = setUpMockClient(searchRequestBuilder); - - RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( - TEST_INDEX_NAME, - client, - MAX_VALUES_PER_PAGE, - mockRuleEntityParser, - mockRuleQueryMapper - ); - + IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); SearchResponse searchResponse = mock(SearchResponse.class); SearchHits searchHits = new SearchHits(new SearchHit[] { new SearchHit(1) }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1.0f); when(searchResponse.getHits()).thenReturn(searchHits); @@ -92,58 +76,22 @@ public void testGetRuleByIdSuccess() { hit.sourceRef(new BytesArray(VALID_JSON)); ActionListener listener = mock(ActionListener.class); - - doAnswer((invocation) -> { - ActionListener actionListener = invocation.getArgument(0); - actionListener.onResponse(searchResponse); - return null; - }).when(searchRequestBuilder).execute(any(ActionListener.class)); - - when(getRuleRequest.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); - rulePersistenceService.getRule(getRuleRequest, listener); + rulePersistenceService.handleGetRuleResponse(_ID_ONE, searchResponse, listener); ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(GetRuleResponse.class); verify(listener).onResponse(responseCaptor.capture()); GetRuleResponse response = responseCaptor.getValue(); assertEquals(response.getRules().size(), 1); + assertEquals(RestStatus.OK, response.getRestStatus()); } public void testGetRuleByIdNotFound() { - GetRuleRequest getRuleRequest = mock(GetRuleRequest.class); - when(getRuleRequest.getId()).thenReturn(_ID_ONE); - QueryBuilder queryBuilder = mock(QueryBuilder.class); - RuleQueryMapper mockRuleQueryMapper = mock(RuleQueryMapper.class); - RuleEntityParser mockRuleEntityParser = mock(RuleEntityParser.class); - Rule mockRule = mock(Rule.class); - - when(mockRuleEntityParser.parse(anyString())).thenReturn(mockRule); - when(mockRuleQueryMapper.from(getRuleRequest)).thenReturn(queryBuilder); - when(queryBuilder.filter(any())).thenReturn(queryBuilder); - - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - Client client = setUpMockClient(searchRequestBuilder); - - RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( - TEST_INDEX_NAME, - client, - MAX_VALUES_PER_PAGE, - mockRuleEntityParser, - mockRuleQueryMapper - ); - + IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); SearchResponse searchResponse = mock(SearchResponse.class); when(searchResponse.getHits()).thenReturn(new SearchHits(new SearchHit[] {}, new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1.0f)); ActionListener listener = mock(ActionListener.class); - - doAnswer(invocationOnMock -> { - ActionListener actionListener = invocationOnMock.getArgument(0); - actionListener.onResponse(searchResponse); - return null; - }).when(searchRequestBuilder).execute(any(ActionListener.class)); - - when(getRuleRequest.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); - rulePersistenceService.getRule(getRuleRequest, listener); + rulePersistenceService.handleGetRuleResponse(_ID_ONE, searchResponse, listener); ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); verify(listener).onFailure(exceptionCaptor.capture()); @@ -151,24 +99,17 @@ public void testGetRuleByIdNotFound() { assertTrue(exception instanceof ResourceNotFoundException); } - private Client setUpMockClient(SearchRequestBuilder searchRequestBuilder) { - Client client = mock(Client.class); - ClusterService clusterService = mock(ClusterService.class); - ClusterState clusterState = mock(ClusterState.class); - Metadata metadata = mock(Metadata.class); - ThreadPool threadPool = mock(ThreadPool.class); - - ThreadContext threadContext = new ThreadContext(Settings.EMPTY); - when(client.threadPool()).thenReturn(threadPool); - when(threadPool.getThreadContext()).thenReturn(threadContext); - when(clusterService.state()).thenReturn(clusterState); - when(clusterState.metadata()).thenReturn(metadata); - + public void testGetRuleWithAttributes() { + IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); + ActionListener listener = mock(ActionListener.class); + SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + Client client = rulePersistenceService.getClient(); when(client.prepareSearch(TEST_INDEX_NAME)).thenReturn(searchRequestBuilder); when(searchRequestBuilder.setQuery(any(QueryBuilder.class))).thenReturn(searchRequestBuilder); when(searchRequestBuilder.setSize(anyInt())).thenReturn(searchRequestBuilder); - - return client; + rulePersistenceService.getRuleFromIndex(null, ATTRIBUTE_MAP, null, listener); + verify(client).prepareSearch(TEST_INDEX_NAME); + verify(searchRequestBuilder).setQuery(any()); } public void testDeleteRule_successful() { diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java new file mode 100644 index 0000000000000..b23014d4d145e --- /dev/null +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java @@ -0,0 +1,54 @@ +/* + * 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.rule.utils; + +import org.opensearch.autotagging.Rule; +import org.opensearch.rule.RuleTestUtils; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.time.Instant; +import java.util.Locale; + +import static org.opensearch.rule.RuleTestUtils.DESCRIPTION_ONE; + +public class IndexStoredRuleParserTests extends OpenSearchTestCase { + public static final String VALID_JSON = String.format(Locale.ROOT, """ + { + "description": "%s", + "mock_feature_type": "feature value", + "mock_attribute_one": ["attribute_value_one", "attribute_value_two"], + "updated_at": "%s" + } + """, DESCRIPTION_ONE, Instant.now().toString()); + + private static final String INVALID_JSON = """ + { + "name": "TestRule", + "description": "A test rule for unit testing", + "mock_attribute_three": ["attribute_value_one", "attribute_value_two"] + } + """; + + public void testParseRule_Success() throws IOException { + Rule parsedRule = IndexStoredRuleParser.parseRule(VALID_JSON, RuleTestUtils.MockRuleFeatureType.INSTANCE); + assertNotNull(parsedRule); + assertEquals("A test rule for unit testing", parsedRule.getDescription()); + assertEquals(RuleTestUtils.MockRuleFeatureType.INSTANCE, parsedRule.getFeatureType()); + } + + public void testParseRule_InvalidJson() { + Exception exception = assertThrows( + RuntimeException.class, + () -> IndexStoredRuleParser.parseRule(INVALID_JSON, RuleTestUtils.MockRuleFeatureType.INSTANCE) + ); + assertTrue(exception.getMessage().contains("mock_attribute_three is not a valid attribute within the mock_feature_type feature.")); + } + +} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java new file mode 100644 index 0000000000000..9093d8f40cc38 --- /dev/null +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java @@ -0,0 +1,42 @@ +/* + * 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.rule.utils; + +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.rule.RuleTestUtils; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.HashMap; + +import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_MAP; +import static org.opensearch.rule.RuleTestUtils.PATTERN_ONE; +import static org.opensearch.rule.RuleTestUtils._ID_ONE; + +public class IndexStoredRuleUtilsTests extends OpenSearchTestCase { + public void testBuildGetRuleQuery_WithId() { + BoolQueryBuilder query = IndexStoredRuleUtils.buildGetRuleQuery( + _ID_ONE, + new HashMap<>(), + RuleTestUtils.MockRuleFeatureType.INSTANCE + ); + assertNotNull(query); + assertEquals(1, query.must().size()); + QueryBuilder idQuery = query.must().get(0); + assertTrue(idQuery.toString().contains(_ID_ONE)); + } + + public void testBuildGetRuleQuery_WithAttributes() { + BoolQueryBuilder query = IndexStoredRuleUtils.buildGetRuleQuery(null, ATTRIBUTE_MAP, RuleTestUtils.MockRuleFeatureType.INSTANCE); + assertNotNull(query); + assertTrue(query.must().size() == 1); + assertTrue(query.toString().contains(RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE.getName())); + assertTrue(query.toString().contains(PATTERN_ONE)); + } +} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java deleted file mode 100644 index 6ec75e4b942ff..0000000000000 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java +++ /dev/null @@ -1,79 +0,0 @@ -/* - * 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.rule.utils; - -import org.opensearch.rule.autotagging.Attribute; -import org.opensearch.rule.autotagging.AutoTaggingRegistry; -import org.opensearch.rule.autotagging.FeatureType; - -import java.util.Map; -import java.util.Set; - -public class RuleTestUtils { - public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; - public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; - public static final String ATTRIBUTE_VALUE_TWO = "mock_attribute_two"; - public static final String DESCRIPTION_ONE = "description_1"; - public static final String FEATURE_TYPE_NAME = "mock_feature_type"; - public static final String TEST_INDEX_NAME = ".test_index_for_rule"; - public static final Map> ATTRIBUTE_MAP = Map.of( - MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, - Set.of(ATTRIBUTE_VALUE_ONE) - ); - - public static final String INVALID_ATTRIBUTE = "invalid_attribute"; - - public static class MockRuleFeatureType implements FeatureType { - - public static final MockRuleFeatureType INSTANCE = new MockRuleFeatureType(); - - private MockRuleFeatureType() {} - - static { - INSTANCE.registerFeatureType(); - } - - @Override - public String getName() { - return FEATURE_TYPE_NAME; - } - - @Override - public Map getAllowedAttributesRegistry() { - return Map.of( - ATTRIBUTE_VALUE_ONE, - MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, - ATTRIBUTE_VALUE_TWO, - MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO - ); - } - - @Override - public void registerFeatureType() { - AutoTaggingRegistry.registerFeatureType(INSTANCE); - } - } - - public enum MockRuleAttributes implements Attribute { - MOCK_RULE_ATTRIBUTE_ONE(ATTRIBUTE_VALUE_ONE), - MOCK_RULE_ATTRIBUTE_TWO(ATTRIBUTE_VALUE_TWO), - INVALID_ATTRIBUTE(RuleTestUtils.INVALID_ATTRIBUTE); - - private final String name; - - MockRuleAttributes(String name) { - this.name = name; - } - - @Override - public String getName() { - return name; - } - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index a82a0404eeb75..dedc02a06e0c6 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -63,8 +63,6 @@ import java.util.List; import java.util.function.Supplier; -import static org.opensearch.plugin.wlm.rule.service.RulePersistenceService.RULES_INDEX; - /** * Plugin class for WorkloadManagement */ @@ -127,8 +125,7 @@ public List getActionFilters() { new ActionPlugin.ActionHandler<>(CreateWorkloadGroupAction.INSTANCE, TransportCreateWorkloadGroupAction.class), new ActionPlugin.ActionHandler<>(GetWorkloadGroupAction.INSTANCE, TransportGetWorkloadGroupAction.class), new ActionPlugin.ActionHandler<>(DeleteWorkloadGroupAction.INSTANCE, TransportDeleteWorkloadGroupAction.class), - new ActionPlugin.ActionHandler<>(UpdateWorkloadGroupAction.INSTANCE, TransportUpdateWorkloadGroupAction.class), - new ActionPlugin.ActionHandler<>(GetRuleAction.INSTANCE, TransportGetRuleAction.class) + new ActionPlugin.ActionHandler<>(UpdateWorkloadGroupAction.INSTANCE, TransportUpdateWorkloadGroupAction.class) ); } @@ -151,8 +148,7 @@ public List getRestHandlers( new RestCreateWorkloadGroupAction(), new RestGetWorkloadGroupAction(), new RestDeleteWorkloadGroupAction(), - new RestUpdateWorkloadGroupAction(), - new RestGetRuleAction() + new RestUpdateWorkloadGroupAction() ); } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java index e4bb603f869a7..0f8f3c3f4628c 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java @@ -14,7 +14,6 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.rule.action.GetRuleRequest; import org.opensearch.rule.action.GetRuleResponse; -import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.rule.service.RulePersistenceService; import org.opensearch.tasks.Task; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java deleted file mode 100644 index 6a96918633e9b..0000000000000 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WlmRuleTestUtils.java +++ /dev/null @@ -1,91 +0,0 @@ -/* - * 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.rule; - -import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.Rule; - -import java.util.Map; -import java.util.Set; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertTrue; - -public class WlmRuleTestUtils { - public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; - public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; - public static final String FEATURE_VALUE_ONE = "feature_value_one"; - public static final String FEATURE_VALUE_TWO = "feature_value_two"; - public static final String PATTERN_ONE = "pattern_1"; - public static final String PATTERN_TWO = "pattern_2"; - public static final String DESCRIPTION_ONE = "description_1"; - public static final String DESCRIPTION_TWO = "description_2"; - public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; - public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; - public static final String SEARCH_AFTER = "search_after_id"; - public static final Map> ATTRIBUTE_MAP = Map.of(QueryGroupAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE)); - public static final Rule ruleOne = Rule.builder() - .description(DESCRIPTION_ONE) - .featureType(QueryGroupFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_ONE) - .attributeMap(Map.of(QueryGroupAttribute.INDEX_PATTERN, Set.of(PATTERN_ONE))) - .updatedAt(TIMESTAMP_ONE) - .build(); - - public static final Rule ruleTwo = Rule.builder() - .description(DESCRIPTION_TWO) - .featureType(QueryGroupFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_TWO) - .attributeMap(Map.of(QueryGroupAttribute.INDEX_PATTERN, Set.of(PATTERN_TWO))) - .updatedAt(TIMESTAMP_TWO) - .build(); - - public static Map ruleMap() { - return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); - } - - // public static WlmRulePersistenceService setUpRulePersistenceService(Map queryGroupMap) { - // Client client = mock(Client.class); - // ClusterService clusterService = mock(ClusterService.class); - // ClusterState clusterState = mock(ClusterState.class); - // Metadata metadata = mock(Metadata.class); - // ThreadPool threadPool = mock(ThreadPool.class); - // - // ThreadContext threadContext = new ThreadContext(Settings.EMPTY); - // when(client.threadPool()).thenReturn(threadPool); - // when(threadPool.getThreadContext()).thenReturn(threadContext); - // when(clusterService.state()).thenReturn(clusterState); - // when(clusterState.metadata()).thenReturn(metadata); - // when(metadata.queryGroups()).thenReturn(queryGroupMap); - // return new WlmRulePersistenceService(clusterService, client); - // } - - public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { - assertEquals(mapOne.size(), mapTwo.size()); - for (Map.Entry entry : mapOne.entrySet()) { - String id = entry.getKey(); - assertTrue(mapTwo.containsKey(id)); - Rule one = mapOne.get(id); - Rule two = mapTwo.get(id); - assertEqualRule(one, two, ruleUpdated); - } - } - - public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { - if (ruleUpdated) { - assertEquals(one.getDescription(), two.getDescription()); - assertEquals(one.getFeatureType(), two.getFeatureType()); - assertEquals(one.getFeatureValue(), two.getFeatureValue()); - assertEquals(one.getAttributeMap(), two.getAttributeMap()); - assertEquals(one.getAttributeMap(), two.getAttributeMap()); - } else { - assertEquals(one, two); - } - } -} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java deleted file mode 100644 index 6df56b15ed875..0000000000000 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/service/WlmRulePersistenceServiceTests.java +++ /dev/null @@ -1,67 +0,0 @@ -/* - * 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.rule.service; - -import org.opensearch.test.OpenSearchTestCase; - -@SuppressWarnings("unchecked") -public class WlmRulePersistenceServiceTests extends OpenSearchTestCase { - // public void testBuildGetRuleQuery_WithId() { - // WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); - // BoolQueryBuilder query = rulePersistenceService.buildGetRuleQuery(_ID_ONE, new HashMap<>()); - // assertTrue(query.hasClauses()); - // assertEquals(QueryBuilders.termQuery(_ID_STRING, _ID_ONE).toString(), query.must().get(0).toString()); - // } - // - // public void testBuildGetRuleQuery_WithFilters() { - // WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); - // BoolQueryBuilder query = rulePersistenceService.buildGetRuleQuery(null, ATTRIBUTE_MAP); - // assertTrue(query.hasClauses()); - // assertEquals(1, query.must().size()); - // assertTrue(query.filter().contains(QueryBuilders.existsQuery(QueryGroupFeatureType.NAME))); - // } - // - // public void testGetRule_WithId() { - // WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); - // Client client = rulePersistenceService.getClient(); - // ActionListener listener = mock(ActionListener.class); - // SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - // SetupMocksForGetRule(client, searchRequestBuilder); - // - // rulePersistenceService.getRule(_ID_ONE, new HashMap<>(), null, listener); - // verify(client).prepareSearch(WlmRulePersistenceService.RULES_INDEX); - // verify(searchRequestBuilder).setQuery(any()); - // verify(searchRequestBuilder).execute(any()); - // } - // - // public void testGetRule_WithSearchAfter() { - // WlmRulePersistenceService rulePersistenceService = setUpRulePersistenceService(new HashMap<>()); - // Client client = rulePersistenceService.getClient(); - // ActionListener listener = mock(ActionListener.class); - // SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - // SetupMocksForGetRule(client, searchRequestBuilder); - // when(searchRequestBuilder.addSort(anyString(), any(SortOrder.class))).thenReturn(searchRequestBuilder); - // when(searchRequestBuilder.searchAfter(any())).thenReturn(searchRequestBuilder); - // - // rulePersistenceService.getRule(null, new HashMap<>(), _ID_TWO, listener); - // verify(searchRequestBuilder).addSort(_ID_STRING, SortOrder.ASC); - // verify(searchRequestBuilder).searchAfter(new Object[] { _ID_TWO }); - // } - // - // public void SetupMocksForGetRule(Client client, SearchRequestBuilder searchRequestBuilder) { - // when(client.prepareSearch(anyString())).thenReturn(searchRequestBuilder); - // when(searchRequestBuilder.setQuery(any())).thenReturn(searchRequestBuilder); - // when(searchRequestBuilder.setSize(anyInt())).thenReturn(searchRequestBuilder); - // doAnswer(invocation -> { - // ActionListener searchListener = invocation.getArgument(0); - // searchListener.onResponse(mock(SearchResponse.class)); - // return null; - // }).when(searchRequestBuilder).execute(any()); - // } -} From 1ab33c1aebaaa69d9755c94dd1e1f600be687c5c Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Thu, 3 Apr 2025 16:22:29 -0700 Subject: [PATCH 10/20] Add create rule api Signed-off-by: Ruirui Zhang --- CHANGELOG.md | 1 + .../rule/action/CreateRuleRequest.java | 60 +++++++ .../rule/action/CreateRuleResponse.java | 81 ++++++++++ .../rule/action/GetRuleRequest.java | 98 ------------ .../rule/rest/RestCreateRuleAction.java | 82 ++++++++++ .../rule/rest/RestGetRuleAction.java | 104 ------------- .../IndexStoredRulePersistenceService.java | 147 +++++++++++++++++- .../rule/service/RulePersistenceService.java | 10 +- .../rule/utils/IndexStoredRuleUtils.java | 17 ++ .../org/opensearch/rule/RuleTestUtils.java | 26 ++-- .../rule/action/CreateRuleRequestTests.java | 33 ++++ .../rule/action/CreateRuleResponseTests.java | 63 ++++++++ ...ndexStoredRulePersistenceServiceTests.java | 116 +++++++++----- .../utils/IndexStoredRuleParserTests.java | 3 +- .../rule/utils/IndexStoredRuleUtilsTests.java | 25 ++- .../wlm/rule/action/CreateWlmRuleAction.java | 36 +++++ .../wlm/rule/action/GetWlmRuleAction.java | 36 ----- .../action/TransportCreateWlmRuleAction.java | 57 +++++++ .../action/TransportGetWlmRuleAction.java | 49 ------ .../rule/rest/RestCreateWlmRuleAction.java | 63 ++++++++ .../wlm/rule/rest/RestGetWlmRuleAction.java | 64 -------- ...java => RestCreateWlmRuleActionTests.java} | 15 +- 22 files changed, 757 insertions(+), 429 deletions(-) create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java create mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java create mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java create mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/CreateWlmRuleAction.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleAction.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/{RestGetWlmRuleActionTests.java => RestCreateWlmRuleActionTests.java} (62%) diff --git a/CHANGELOG.md b/CHANGELOG.md index b11ffb7e3bd19..6a31b5e21d607 100644 --- 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/), - [Rule based auto-tagging] Add get rule API ([#17336](https://github.com/opensearch-project/OpenSearch/pull/17336)) - [Rule based auto-tagging] Add Delete Rule API ([#18184](https://github.com/opensearch-project/OpenSearch/pull/18184)) - Add paginated wlm/stats API ([#17638](https://github.com/opensearch-project/OpenSearch/pull/17638)) +- [Rule based auto-tagging] Add Create rule API ([#17792](https://github.com/opensearch-project/OpenSearch/pull/17792)) - Implement parallel shard refresh behind cluster settings ([#17782](https://github.com/opensearch-project/OpenSearch/pull/17782)) - Bump OpenSearch Core main branch to 3.0.0 ([#18039](https://github.com/opensearch-project/OpenSearch/pull/18039)) - [Rule based Auto-tagging] Add wlm `ActionFilter` ([#17791](https://github.com/opensearch-project/OpenSearch/pull/17791)) diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java new file mode 100644 index 0000000000000..a2956d1a65bad --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java @@ -0,0 +1,60 @@ +/* + * 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.rule.action; + +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.autotagging.Rule; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +import java.io.IOException; + +/** + * A request for create Rule + * @opensearch.experimental + */ +public class CreateRuleRequest extends ActionRequest { + private final Rule rule; + + /** + * constructor for CreateRuleRequest + * @param rule the rule to create + */ + public CreateRuleRequest(Rule rule) { + this.rule = rule; + } + + /** + * Constructs a CreateRuleRequest from a StreamInput for deserialization + * @param in - The {@link StreamInput} instance to read from. + */ + public CreateRuleRequest(StreamInput in) throws IOException { + super(in); + rule = new Rule(in); + } + + @Override + public ActionRequestValidationException validate() { + return null; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + rule.writeTo(out); + } + + /** + * rule getter + */ + public Rule getRule() { + return rule; + } +} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java new file mode 100644 index 0000000000000..aef3936c554c3 --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java @@ -0,0 +1,81 @@ +/* + * 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.rule.action; + +import org.opensearch.autotagging.Rule; +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.Map; + +import static org.opensearch.autotagging.Rule._ID_STRING; + +/** + * Response for the create API for Rule + * @opensearch.experimental + */ +public class CreateRuleResponse extends ActionResponse implements ToXContent, ToXContentObject { + private final String _id; + private final Rule rule; + private final RestStatus restStatus; + + /** + * contructor for CreateRuleResponse + * @param id - the id for the rule created + * @param rule - the rule created + * @param restStatus - the status of CreateRuleResponse + */ + public CreateRuleResponse(String id, final Rule rule, RestStatus restStatus) { + this._id = id; + this.rule = rule; + this.restStatus = restStatus; + } + + /** + * Constructs a CreateRuleResponse from a StreamInput for deserialization + * @param in - The {@link StreamInput} instance to read from. + */ + public CreateRuleResponse(StreamInput in) throws IOException { + _id = in.readString(); + rule = new Rule(in); + restStatus = RestStatus.readFrom(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeString(_id); + rule.writeTo(out); + RestStatus.writeTo(out, restStatus); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + return rule.toXContent(builder, new MapParams(Map.of(_ID_STRING, _id))); + } + + /** + * rule getter + */ + public Rule getRule() { + return rule; + } + + /** + * restStatus getter + */ + public RestStatus getRestStatus() { + return restStatus; + } +} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java deleted file mode 100644 index 03d1899a2e6a8..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleRequest.java +++ /dev/null @@ -1,98 +0,0 @@ -/* - * 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.rule.action; - -import org.opensearch.action.ActionRequest; -import org.opensearch.action.ActionRequestValidationException; -import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.core.common.io.stream.StreamOutput; - -import java.io.IOException; -import java.util.HashSet; -import java.util.Map; -import java.util.Set; - -/** - * A request for get Rule - * Example Request: - * The endpoint "localhost:9200/_wlm/rule" is specific to the Workload Management feature to manage rules - * curl -X GET "localhost:9200/_wlm/rule" - get all rules - * curl -X GET "localhost:9200/_wlm/rule/{_id}" - get single rule by id - * curl -X GET "localhost:9200/_wlm/rule?index_pattern=a,b" - get all rules containing attribute index_pattern as a or b - * @opensearch.experimental - */ -public class GetRuleRequest extends ActionRequest { - private final String id; - private final Map> attributeFilters; - private final String searchAfter; - private final FeatureType featureType; - - /** - * Constructor for GetRuleRequest - * @param id - Rule id to get - * @param attributeFilters - Rules will be filtered based on the attribute-value mappings. - * @param searchAfter - The sort value used for pagination. - * @param featureType - The feature type related to rule. - */ - public GetRuleRequest(String id, Map> attributeFilters, String searchAfter, FeatureType featureType) { - this.id = id; - this.attributeFilters = attributeFilters; - this.searchAfter = searchAfter; - this.featureType = featureType; - } - - /** - * Constructs a GetRuleRequest from a StreamInput for deserialization. - * @param in - The {@link StreamInput} instance to read from. - */ - public GetRuleRequest(StreamInput in) throws IOException { - super(in); - id = in.readOptionalString(); - featureType = FeatureType.from(in); - attributeFilters = in.readMap(i -> Attribute.from(i, featureType), i -> new HashSet<>(i.readStringList())); - searchAfter = in.readOptionalString(); - } - - @Override - public ActionRequestValidationException validate() { - return null; - } - - @Override - public void writeTo(StreamOutput out) throws IOException { - super.writeTo(out); - out.writeOptionalString(id); - featureType.writeTo(out); - out.writeMap(attributeFilters, (output, attribute) -> attribute.writeTo(output), StreamOutput::writeStringCollection); - out.writeOptionalString(searchAfter); - } - - /** - * id getter - */ - public String getId() { - return id; - } - - /** - * attributeFilters getter - */ - public Map> getAttributeFilters() { - return attributeFilters; - } - - /** - * searchAfter getter - */ - public String getSearchAfter() { - return searchAfter; - } -} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java new file mode 100644 index 0000000000000..a096ac517cb40 --- /dev/null +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.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.rule.rest; + +import org.opensearch.action.ActionType; +import org.opensearch.autotagging.FeatureType; +import org.opensearch.autotagging.Rule; +import org.opensearch.autotagging.Rule.Builder; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentParser; +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 org.opensearch.rule.action.CreateRuleRequest; +import org.opensearch.rule.action.CreateRuleResponse; +import org.opensearch.transport.client.node.NodeClient; +import org.joda.time.Instant; + +import java.io.IOException; +import java.util.List; + +/** + * Rest action to create a Rule + * @opensearch.experimental + */ +public abstract class RestCreateRuleAction extends BaseRestHandler { + /** + * constructor for RestCreateRuleAction + */ + public RestCreateRuleAction() {} + + @Override + public abstract String getName(); + + @Override + public abstract List routes(); + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + try (XContentParser parser = request.contentParser()) { + Builder builder = Builder.fromXContent(parser, retrieveFeatureTypeInstance()); + CreateRuleRequest createRuleRequest = buildCreateRuleRequest(builder.updatedAt(Instant.now().toString()).build()); + return channel -> client.execute(retrieveCreateRuleActionInstance(), createRuleRequest, createRuleResponse(channel)); + } + } + + private RestResponseListener createRuleResponse(final RestChannel channel) { + return new RestResponseListener<>(channel) { + @Override + public RestResponse buildResponse(final CreateRuleResponse response) throws Exception { + return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS)); + } + }; + } + + /** + * Abstract method for subclasses to provide specific ActionType Instance + */ + protected abstract > T retrieveCreateRuleActionInstance(); + + /** + * Abstract method for subclasses to provide specific FeatureType Instance + */ + protected abstract FeatureType retrieveFeatureTypeInstance(); + + /** + * Abstract method for subclasses to construct a {@link CreateRuleRequest}. This method allows subclasses + * to define their own request-building logic depending on their specific needs. + * @param rule - the rule to create + */ + protected abstract CreateRuleRequest buildCreateRuleRequest(Rule rule); +} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java deleted file mode 100644 index f27bdfbbb4e59..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java +++ /dev/null @@ -1,104 +0,0 @@ -/* - * 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.rule.rest; - -import org.opensearch.action.ActionType; -import org.opensearch.autotagging.Attribute; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.core.xcontent.ToXContent; -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 org.opensearch.rule.action.GetRuleRequest; -import org.opensearch.rule.action.GetRuleResponse; -import org.opensearch.transport.client.node.NodeClient; - -import java.io.IOException; -import java.util.Arrays; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; - -import static org.opensearch.autotagging.Rule._ID_STRING; - -/** - * Rest action to get a Rule - * @opensearch.experimental - */ -public abstract class RestGetRuleAction extends BaseRestHandler { - /** - * field name used for pagination - */ - public static final String SEARCH_AFTER_STRING = "search_after"; - - /** - * Constructor for RestGetRuleAction - */ - public RestGetRuleAction() {} - - @Override - public abstract String getName(); - - @Override - public abstract List routes(); - - @Override - protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - final Map> attributeFilters = new HashMap<>(); - for (String attributeName : request.params().keySet()) { - if (attributeName.equals(_ID_STRING) || attributeName.equals(SEARCH_AFTER_STRING)) { - continue; - } - String[] valuesArray = request.param(attributeName).split(","); - attributeFilters.put(getAttributeFromName(attributeName), new HashSet<>(Arrays.asList(valuesArray))); - } - final GetRuleRequest getRuleRequest = buildGetRuleRequest( - request.param(_ID_STRING), - attributeFilters, - request.param(SEARCH_AFTER_STRING) - ); - return channel -> client.execute(retrieveGetRuleActionInstance(), getRuleRequest, getRuleResponse(channel)); - } - - private RestResponseListener getRuleResponse(final RestChannel channel) { - return new RestResponseListener<>(channel) { - @Override - public RestResponse buildResponse(final GetRuleResponse response) throws Exception { - return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS)); - } - }; - } - - /** - * Abstract method for subclasses to retrieve the Attribute corresponding - * to the attribute name. - * @param name - The name of the attribute to retrieve. - */ - protected abstract Attribute getAttributeFromName(String name); - - /** - * Abstract method for subclasses to provide specific ActionType Instance - */ - protected abstract > T retrieveGetRuleActionInstance(); - - /** - * Abstract method for subclasses to construct a {@link GetRuleRequest}. This method allows subclasses - * to define their own request-building logic depending on their specific needs. - * - * @param id - The ID of the rule to retrieve. - * @param attributeFilters - A map of {@link Attribute} keys to sets of string values for filtering. - * @param searchAfter - The pagination value to fetch the next set of results. - */ - protected abstract GetRuleRequest buildGetRuleRequest(String id, Map> attributeFilters, String searchAfter); -} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java index 52cd6e97a75b7..75d8141d9a139 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -10,17 +10,25 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.ResourceAlreadyExistsException; import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.admin.indices.create.CreateIndexRequest; +import org.opensearch.action.admin.indices.create.CreateIndexResponse; +import org.opensearch.action.index.IndexRequest; import org.opensearch.action.search.SearchRequestBuilder; import org.opensearch.action.search.SearchResponse; import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.FeatureType; import org.opensearch.autotagging.Rule; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; import org.opensearch.index.query.BoolQueryBuilder; -import org.opensearch.rule.action.GetRuleRequest; +import org.opensearch.rule.action.CreateRuleRequest; +import org.opensearch.rule.action.CreateRuleResponse; import org.opensearch.rule.action.GetRuleResponse; import org.opensearch.rule.utils.IndexStoredRuleParser; import org.opensearch.rule.utils.IndexStoredRuleUtils; @@ -28,6 +36,7 @@ import org.opensearch.search.sort.SortOrder; import org.opensearch.transport.client.Client; +import java.io.IOException; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -45,37 +54,152 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService * The system index name used for storing rules */ private final String indexName; + private final ClusterService clusterService; private final Client client; private final FeatureType featureType; private final int maxRulesPerGetRequest; private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class); + private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); /** * Constructs an instance of {@link IndexStoredRulePersistenceService} with the specified parameters. * This service handles persistence and retrieval of stored rules within an OpenSearch index. * @param indexName - The name of the OpenSearch index where the rules are stored. + * @param clusterService - The clusterService used in IndexStoredRulePersistenceService. * @param client - The OpenSearch client used to interact with the OpenSearch cluster. * @param featureType - The feature type associated with the stored rules. * @param maxRulesPerGetRequest - The maximum number of rules that can be returned in a single get request. */ - public IndexStoredRulePersistenceService(String indexName, Client client, FeatureType featureType, int maxRulesPerGetRequest) { + public IndexStoredRulePersistenceService( + String indexName, + ClusterService clusterService, + Client client, + FeatureType featureType, + int maxRulesPerGetRequest + ) { this.indexName = indexName; + this.clusterService = clusterService; this.client = client; this.featureType = featureType; this.maxRulesPerGetRequest = maxRulesPerGetRequest; } /** - * Entry point for the get rule api logic in persistence service. - * @param getRuleRequest the getRuleRequest to process. - * @param listener the listener for GetRuleResponse. + * Entry point for the create rule api logic in persistence service + * @param request - The CreateRuleRequest + * @param listener - ActionListener for CreateRuleResponse + */ + public void createRule(CreateRuleRequest request, ActionListener listener) { + try (ThreadContext.StoredContext ctx = getContext()) { + createIndexIfAbsent(new ActionListener<>() { + @Override + public void onResponse(Boolean indexCreated) { + if (!indexCreated) { + listener.onFailure(new IllegalStateException(indexName + " index creation failed and rule cannot be persisted")); + return; + } + checkDuplicateRule(request.getRule(), listener); + } + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + }); + } + } + + /** + * Check if there's an existing Rule with the same attributes. + * For example, if there's an existing Rule with the attribute index_pattern: ["a", "b", "c"], + * then we cannot create another Rule with only one attribute index_pattern: ["b"], because the value "b" + * already exists under another Rule. Note that the conflict exists only when we have the exact same attribute + * names in the two rules (That is, a Rule with attribute "index_pattern" won't create a conflict with another + * Rule that has "index_pattern" and some other attributes). + * @param rule - The rule to update. + * @param listener - ActionListener for CreateRuleResponse */ - public void getRule(GetRuleRequest getRuleRequest, ActionListener listener) { - getRuleFromIndex(getRuleRequest.getId(), getRuleRequest.getAttributeFilters(), getRuleRequest.getSearchAfter(), listener); + public void checkDuplicateRule(Rule rule, ActionListener listener) { + try (ThreadContext.StoredContext ctx = getContext()) { + getRuleFromIndex(null, rule.getAttributeMap(), null, new ActionListener<>() { + @Override + public void onResponse(GetRuleResponse getRuleResponse) { + String duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(rule.getAttributeMap(), getRuleResponse.getRules()); + if (duplicateRuleId != null) { + listener.onFailure( + new IllegalArgumentException( + "A rule that has the same attribute values already exists under rule id " + duplicateRuleId + ) + ); + return; + } + persistRule(rule, listener); + } + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + }); + } } /** - * Get rules from index. If id is provided, we only get a single rule. + * Persist the rule in the index + * @param rule - The rule to update. + * @param listener - ActionListener for CreateRuleResponse + */ + public void persistRule(Rule rule, ActionListener listener) { + try (ThreadContext.StoredContext ctx = getContext()) { + IndexRequest indexRequest = new IndexRequest(indexName).source( + rule.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS) + ); + client.index(indexRequest, ActionListener.wrap(indexResponse -> { + listener.onResponse(new CreateRuleResponse(indexResponse.getId(), rule, RestStatus.OK)); + }, e -> { + logger.warn("Failed to save Rule object due to error: {}", e.getMessage()); + listener.onFailure(e); + })); + } catch (IOException e) { + logger.error("Error saving rule to index: {}", indexName, e); + listener.onFailure(new RuntimeException("Failed to save rule to index.")); + } + } + + /** + * Creates the system index .rules if it doesn't exist + * @param listener - ActionListener for CreateRuleResponse + */ + void createIndexIfAbsent(ActionListener listener) { + try (ThreadContext.StoredContext ctx = getContext()) { + if (clusterService.state().metadata().hasIndex(indexName)) { + listener.onResponse(true); + return; + } + final CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName).settings(indexSettings); + client.admin().indices().create(createIndexRequest, new ActionListener<>() { + @Override + public void onResponse(CreateIndexResponse response) { + logger.info("Index {} created?: {}", indexName, response.isAcknowledged()); + listener.onResponse(response.isAcknowledged()); + } + + @Override + public void onFailure(Exception e) { + if (e instanceof ResourceAlreadyExistsException) { + logger.info("Index {} already exists", indexName); + listener.onResponse(true); + } else { + logger.error("Failed to create index {}: {}", indexName, e.getMessage()); + listener.onFailure(e); + } + } + }); + } + } + + /** + * Entry point for the get rule api logic in persistence service. If id is provided, we only get a single rule. * Otherwise, we get all rules that satisfy the attributeFilters. * @param id - The id of the rule to get. * @param attributeFilters - A map containing the attributes that user want to filter on @@ -132,4 +256,11 @@ private ThreadContext.StoredContext getContext() { public Client getClient() { return client; } + + /** + * clusterService getter + */ + public ClusterService getClusterService() { + return clusterService; + } } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java index adcccfc429210..6b814a271db60 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java @@ -9,8 +9,8 @@ package org.opensearch.rule.service; import org.opensearch.core.action.ActionListener; -import org.opensearch.rule.action.GetRuleRequest; -import org.opensearch.rule.action.GetRuleResponse; +import org.opensearch.rule.action.CreateRuleRequest; +import org.opensearch.rule.action.CreateRuleResponse; /** * Interface for a service that handles rule persistence CRUD operations. @@ -19,9 +19,9 @@ public interface RulePersistenceService { /** - * Get rules based on the provided request. - * @param request The request containing the details for retrieving the rule. + * Create rules based on the provided request. + * @param request The request containing the details for creating the rule. * @param listener The listener that will handle the response or failure. */ - void getRule(GetRuleRequest request, ActionListener listener); + void createRule(CreateRuleRequest request, ActionListener listener); } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java index d43557a81b4f5..27af328c07e60 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java @@ -10,6 +10,7 @@ import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.FeatureType; +import org.opensearch.autotagging.Rule; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilders; @@ -54,4 +55,20 @@ public static BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeMapToValidate, Map ruleMap) { + for (Map.Entry entry : ruleMap.entrySet()) { + String ruleId = entry.getKey(); + Rule currRule = entry.getValue(); + if (attributeMapToValidate.size() == currRule.getAttributeMap().size()) { + return ruleId; + } + } + return null; + } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java index 14388d65a5e20..c2c5f12b5aa1c 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java @@ -35,15 +35,19 @@ public class RuleTestUtils { public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; public static final String FEATURE_VALUE_ONE = "feature_value_one"; public static final String FEATURE_VALUE_TWO = "feature_value_two"; - public static final String PATTERN_ONE = "pattern_1"; - public static final String PATTERN_TWO = "pattern_2"; + public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; + public static final String ATTRIBUTE_VALUE_TWO = "mock_attribute_two"; public static final String DESCRIPTION_ONE = "description_1"; public static final String DESCRIPTION_TWO = "description_2"; + public static final String SEARCH_AFTER = "search_after_id"; public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; - public static final String SEARCH_AFTER = "search_after_id"; + public static final String FEATURE_TYPE_NAME = "mock_feature_type"; public static final String TEST_INDEX_NAME = ".test_index_for_rule"; - public static final Map> ATTRIBUTE_MAP = Map.of(MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, Set.of(PATTERN_ONE)); + public static final Map> ATTRIBUTE_MAP = Map.of( + MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, + Set.of(ATTRIBUTE_VALUE_ONE) + ); public static final Rule ruleOne = Rule.builder() .description(DESCRIPTION_ONE) .featureType(MockRuleFeatureType.INSTANCE) @@ -56,7 +60,7 @@ public class RuleTestUtils { .description(DESCRIPTION_TWO) .featureType(MockRuleFeatureType.INSTANCE) .featureValue(FEATURE_VALUE_TWO) - .attributeMap(Map.of(MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(PATTERN_TWO))) + .attributeMap(Map.of(MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(ATTRIBUTE_VALUE_TWO))) .updatedAt(TIMESTAMP_TWO) .build(); @@ -77,7 +81,7 @@ public static IndexStoredRulePersistenceService setUpIndexStoredRulePersistenceS when(clusterService.state()).thenReturn(clusterState); when(clusterState.metadata()).thenReturn(metadata); when(metadata.queryGroups()).thenReturn(queryGroupMap); - return new IndexStoredRulePersistenceService(TEST_INDEX_NAME, client, MockRuleFeatureType.INSTANCE, 50); + return new IndexStoredRulePersistenceService(TEST_INDEX_NAME, clusterService, client, MockRuleFeatureType.INSTANCE, 50); } public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { @@ -115,15 +119,15 @@ private MockRuleFeatureType() {} @Override public String getName() { - return "mock_feature_type"; + return FEATURE_TYPE_NAME; } @Override public Map getAllowedAttributesRegistry() { return Map.of( - "mock_attribute_one", + ATTRIBUTE_VALUE_ONE, MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, - "mock_attribute_two", + ATTRIBUTE_VALUE_TWO, MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO ); } @@ -135,8 +139,8 @@ public void registerFeatureType() { } public enum MockRuleAttributes implements Attribute { - MOCK_RULE_ATTRIBUTE_ONE("mock_attribute_one"), - MOCK_RULE_ATTRIBUTE_TWO("mock_attribute_two"); + MOCK_RULE_ATTRIBUTE_ONE(ATTRIBUTE_VALUE_ONE), + MOCK_RULE_ATTRIBUTE_TWO(ATTRIBUTE_VALUE_TWO); ; private final String name; diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java new file mode 100644 index 0000000000000..ec415fbd9c0a2 --- /dev/null +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java @@ -0,0 +1,33 @@ +/* + * 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.rule.action; + +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; + +import static org.opensearch.rule.RuleTestUtils.assertEqualRule; +import static org.opensearch.rule.RuleTestUtils.ruleOne; + +public class CreateRuleRequestTests extends OpenSearchTestCase { + + /** + * Test case to verify the serialization and deserialization of CreateRuleRequest. + */ + public void testSerialization() throws IOException { + CreateRuleRequest request = new CreateRuleRequest(ruleOne); + BytesStreamOutput out = new BytesStreamOutput(); + request.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + CreateRuleRequest otherRequest = new CreateRuleRequest(streamInput); + assertEqualRule(ruleOne, otherRequest.getRule(), false); + } +} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java new file mode 100644 index 0000000000000..8ec92ccada46f --- /dev/null +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java @@ -0,0 +1,63 @@ +/* + * 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.rule.action; + +import org.opensearch.autotagging.Rule; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.Map; + +import static org.opensearch.rule.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.RuleTestUtils.assertEqualRules; +import static org.opensearch.rule.RuleTestUtils.ruleOne; +import static org.mockito.Mockito.mock; + +public class CreateRuleResponseTests extends OpenSearchTestCase { + + /** + * Test case to verify serialization and deserialization of CreateRuleResponse + */ + public void testSerialization() throws IOException { + CreateRuleResponse response = new CreateRuleResponse(_ID_ONE, ruleOne, RestStatus.OK); + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + StreamInput streamInput = out.bytes().streamInput(); + CreateRuleResponse otherResponse = new CreateRuleResponse(streamInput); + assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); + Rule responseRule = response.getRule(); + Rule otherResponseRule = otherResponse.getRule(); + assertEqualRules(Map.of(_ID_ONE, responseRule), Map.of(_ID_ONE, otherResponseRule), false); + } + + /** + * Test case to validate the toXContent method of CreateRuleResponse + */ + public void testToXContentCreateRule() throws IOException { + XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); + CreateRuleResponse response = new CreateRuleResponse(_ID_ONE, ruleOne, RestStatus.OK); + String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); + String expected = "{\n" + + " \"_id\" : \"AgfUO5Ja9yfvhdONlYi3TQ==\",\n" + + " \"description\" : \"description_1\",\n" + + " \"mock_attribute_one\" : [\n" + + " \"pattern_1\"\n" + + " ],\n" + + " \"mock_feature_type\" : \"feature_value_one\",\n" + + " \"updated_at\" : \"2024-01-26T08:58:57.558Z\"\n" + + "}"; + assertEquals(expected, actual); + } +} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java index 73a3b1cd962af..c0de380d05580 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java @@ -32,17 +32,21 @@ import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.RuleQueryMapper; import org.opensearch.rule.autotagging.Rule; -import org.opensearch.rule.utils.RuleTestUtils; +import org.opensearch.rule.RuleTestUtils; +import org.opensearch.action.admin.indices.create.CreateIndexRequest; +import org.opensearch.action.admin.indices.create.CreateIndexResponse; +import org.opensearch.action.index.IndexRequest; +import org.opensearch.action.index.IndexResponse; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.core.action.ActionListener; -import org.opensearch.core.common.bytes.BytesArray; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.rule.action.GetRuleResponse; -import org.opensearch.search.SearchHit; -import org.opensearch.search.SearchHits; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.rule.action.CreateRuleResponse; import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.transport.client.AdminClient; import org.opensearch.transport.client.Client; +import org.opensearch.transport.client.IndicesAdminClient; +import java.io.IOException; import java.util.HashMap; import org.mockito.ArgumentCaptor; @@ -55,61 +59,87 @@ import static org.mockito.ArgumentMatchers.argThat; import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_MAP; import static org.opensearch.rule.RuleTestUtils.TEST_INDEX_NAME; -import static org.opensearch.rule.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.RuleTestUtils.ruleOne; import static org.opensearch.rule.RuleTestUtils.setUpIndexStoredRulePersistenceService; -import static org.opensearch.rule.utils.IndexStoredRuleParserTests.VALID_JSON; import static org.mockito.Mockito.any; -import static org.mockito.Mockito.anyInt; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @SuppressWarnings("unchecked") public class IndexStoredRulePersistenceServiceTests extends OpenSearchTestCase { - public void testGetRuleByIdSuccess() { + public void testCreateIndexIfAbsent_IndexExists() { IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - SearchResponse searchResponse = mock(SearchResponse.class); - SearchHits searchHits = new SearchHits(new SearchHit[] { new SearchHit(1) }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1.0f); - when(searchResponse.getHits()).thenReturn(searchHits); - SearchHit hit = searchHits.getHits()[0]; - hit.sourceRef(new BytesArray(VALID_JSON)); - - ActionListener listener = mock(ActionListener.class); - rulePersistenceService.handleGetRuleResponse(_ID_ONE, searchResponse, listener); - - ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(GetRuleResponse.class); - verify(listener).onResponse(responseCaptor.capture()); - GetRuleResponse response = responseCaptor.getValue(); - assertEquals(response.getRules().size(), 1); - assertEquals(RestStatus.OK, response.getRestStatus()); + ClusterService clusterService = rulePersistenceService.getClusterService(); + when(clusterService.state().metadata().hasIndex(TEST_INDEX_NAME)).thenReturn(true); + ActionListener listener = mock(ActionListener.class); + rulePersistenceService.createIndexIfAbsent(listener); + verify(listener).onResponse(true); + verifyNoMoreInteractions(listener); } - public void testGetRuleByIdNotFound() { + public void testCreateIndexIfAbsent() { IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - SearchResponse searchResponse = mock(SearchResponse.class); - when(searchResponse.getHits()).thenReturn(new SearchHits(new SearchHit[] {}, new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1.0f)); + Client client = rulePersistenceService.getClient(); + IndicesAdminClient indicesAdminClient = mock(IndicesAdminClient.class); + when(client.admin()).thenReturn(mock(AdminClient.class)); + when(client.admin().indices()).thenReturn(indicesAdminClient); + doAnswer(invocation -> { + ActionListener listener = invocation.getArgument(1); + listener.onResponse(new CreateIndexResponse(true, true, TEST_INDEX_NAME)); // Assuming the index creation was successful + return null; + }).when(indicesAdminClient).create(any(CreateIndexRequest.class), any(ActionListener.class)); + rulePersistenceService.createIndexIfAbsent(new ActionListener<>() { + @Override + public void onResponse(Boolean indexCreated) { + assertTrue(indexCreated); + } + + @Override + public void onFailure(Exception e) { + fail("Index creation failed: " + e.getMessage()); + } + }); + verify(indicesAdminClient).create(any(CreateIndexRequest.class), any(ActionListener.class)); + } - ActionListener listener = mock(ActionListener.class); - rulePersistenceService.handleGetRuleResponse(_ID_ONE, searchResponse, listener); + public void testPersistRuleSuccess() { + IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); + Client client = rulePersistenceService.getClient(); + ActionListener listener = mock(ActionListener.class); + IndexResponse indexResponse = new IndexResponse(new ShardId(TEST_INDEX_NAME, "uuid", 0), "id", 1, 1, 1, true); + doAnswer(invocation -> { + ActionListener actionListener = invocation.getArgument(1); + actionListener.onResponse(indexResponse); + return null; + }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); + + rulePersistenceService.persistRule(ruleOne, listener); + verify(client).index(any(IndexRequest.class), any(ActionListener.class)); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(CreateRuleResponse.class); + verify(listener).onResponse(responseCaptor.capture()); - ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); - verify(listener).onFailure(exceptionCaptor.capture()); - Exception exception = exceptionCaptor.getValue(); - assertTrue(exception instanceof ResourceNotFoundException); + CreateRuleResponse createRuleResponse = responseCaptor.getValue(); + assertNotNull(createRuleResponse); + assertEquals(ruleOne, createRuleResponse.getRule()); } - public void testGetRuleWithAttributes() { + public void testPersistRuleFailure() throws IOException { IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - ActionListener listener = mock(ActionListener.class); - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); Client client = rulePersistenceService.getClient(); - when(client.prepareSearch(TEST_INDEX_NAME)).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.setQuery(any(QueryBuilder.class))).thenReturn(searchRequestBuilder); - when(searchRequestBuilder.setSize(anyInt())).thenReturn(searchRequestBuilder); - rulePersistenceService.getRuleFromIndex(null, ATTRIBUTE_MAP, null, listener); - verify(client).prepareSearch(TEST_INDEX_NAME); - verify(searchRequestBuilder).setQuery(any()); + ActionListener listener = mock(ActionListener.class); + doAnswer(invocation -> { + ActionListener actionListener = invocation.getArgument(1); + actionListener.onFailure(new RuntimeException("Indexing failed")); + return null; + }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); + + rulePersistenceService.persistRule(ruleOne, listener); + verify(client).index(any(IndexRequest.class), any(ActionListener.class)); + verify(listener).onFailure(any(RuntimeException.class)); } public void testDeleteRule_successful() { diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java index b23014d4d145e..277fce2e2dc72 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java @@ -39,7 +39,7 @@ public class IndexStoredRuleParserTests extends OpenSearchTestCase { public void testParseRule_Success() throws IOException { Rule parsedRule = IndexStoredRuleParser.parseRule(VALID_JSON, RuleTestUtils.MockRuleFeatureType.INSTANCE); assertNotNull(parsedRule); - assertEquals("A test rule for unit testing", parsedRule.getDescription()); + assertEquals(DESCRIPTION_ONE, parsedRule.getDescription()); assertEquals(RuleTestUtils.MockRuleFeatureType.INSTANCE, parsedRule.getFeatureType()); } @@ -50,5 +50,4 @@ public void testParseRule_InvalidJson() { ); assertTrue(exception.getMessage().contains("mock_attribute_three is not a valid attribute within the mock_feature_type feature.")); } - } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java index 9093d8f40cc38..b8bafa7e7d33c 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java @@ -8,16 +8,21 @@ package org.opensearch.rule.utils; +import org.opensearch.autotagging.Attribute; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.rule.RuleTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.util.HashMap; +import java.util.Map; +import java.util.Set; import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_MAP; -import static org.opensearch.rule.RuleTestUtils.PATTERN_ONE; +import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_VALUE_ONE; +import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_VALUE_TWO; import static org.opensearch.rule.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.RuleTestUtils.ruleOne; public class IndexStoredRuleUtilsTests extends OpenSearchTestCase { public void testBuildGetRuleQuery_WithId() { @@ -37,6 +42,22 @@ public void testBuildGetRuleQuery_WithAttributes() { assertNotNull(query); assertTrue(query.must().size() == 1); assertTrue(query.toString().contains(RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE.getName())); - assertTrue(query.toString().contains(PATTERN_ONE)); + assertTrue(query.toString().contains(ATTRIBUTE_VALUE_ONE)); + } + + public void testGetDuplicateRuleId_Found() { + String duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(ATTRIBUTE_MAP, Map.of(_ID_ONE, ruleOne)); + assertEquals(_ID_ONE, duplicateRuleId); + } + + public void testGetDuplicateRuleId_NotFound() { + Map> map = Map.of( + RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, + Set.of(ATTRIBUTE_VALUE_ONE), + RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, + Set.of(ATTRIBUTE_VALUE_TWO) + ); + String duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(map, Map.of(_ID_ONE, ruleOne)); + assertNull(duplicateRuleId); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/CreateWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/CreateWlmRuleAction.java new file mode 100644 index 0000000000000..8321ed7a89b17 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/CreateWlmRuleAction.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.rule.action; + +import org.opensearch.action.ActionType; +import org.opensearch.rule.action.CreateRuleResponse; + +/** + * Action type for creating a Rule in workload management + * @opensearch.experimental + */ +public class CreateWlmRuleAction extends ActionType { + + /** + * An instance of CreateWlmRuleAction + */ + public static final CreateWlmRuleAction INSTANCE = new CreateWlmRuleAction(); + + /** + * Name for CreateWlmRuleAction + */ + public static final String NAME = "cluster:admin/opensearch/wlm/rule/_create"; + + /** + * Default constructor + */ + private CreateWlmRuleAction() { + super(NAME, CreateRuleResponse::new); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java deleted file mode 100644 index 080d417f1319d..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/GetWlmRuleAction.java +++ /dev/null @@ -1,36 +0,0 @@ -/* - * 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.rule.action; - -import org.opensearch.action.ActionType; -import org.opensearch.rule.action.GetRuleResponse; - -/** - * Action type for getting Rules in workload management - * @opensearch.experimental - */ -public class GetWlmRuleAction extends ActionType { - - /** - * An instance of GetWlmRuleAction - */ - public static final GetWlmRuleAction INSTANCE = new GetWlmRuleAction(); - - /** - * Name for GetWlmRuleAction - */ - public static final String NAME = "cluster:admin/opensearch/wlm/rule/_get"; - - /** - * Default constructor for GetWlmRuleAction - */ - private GetWlmRuleAction() { - super(NAME, GetRuleResponse::new); - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java new file mode 100644 index 0000000000000..613a6f6760ab8 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java @@ -0,0 +1,57 @@ +/* + * 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.rule.action; + +import org.opensearch.ResourceNotFoundException; +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.rule.action.CreateRuleRequest; +import org.opensearch.rule.action.CreateRuleResponse; +import org.opensearch.rule.service.IndexStoredRulePersistenceService; +import org.opensearch.rule.service.RulePersistenceService; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +/** + * Transport action to create Rule in workload management + * @opensearch.experimental + */ +public class TransportCreateWlmRuleAction extends HandledTransportAction { + + private final IndexStoredRulePersistenceService rulePersistenceService; + + /** + * Constructor for TransportCreateWlmRuleAction + * + * @param transportService - a {@link TransportService} object + * @param actionFilters - a {@link ActionFilters} object + * @param rulePersistenceService - a {@link RulePersistenceService} object + */ + @Inject + public TransportCreateWlmRuleAction( + TransportService transportService, + ActionFilters actionFilters, + IndexStoredRulePersistenceService rulePersistenceService + ) { + super(CreateWlmRuleAction.NAME, transportService, actionFilters, CreateRuleRequest::new); + this.rulePersistenceService = rulePersistenceService; + } + + @Override + protected void doExecute(Task task, CreateRuleRequest request, ActionListener listener) { + String queryGroupId = request.getRule().getFeatureValue(); + if (!rulePersistenceService.getClusterService().state().metadata().queryGroups().containsKey(queryGroupId)) { + listener.onFailure(new ResourceNotFoundException("Couldn't find an existing query group with id: " + queryGroupId)); + return; + } + rulePersistenceService.createRule(request, listener); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java deleted file mode 100644 index 0f8f3c3f4628c..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportGetWlmRuleAction.java +++ /dev/null @@ -1,49 +0,0 @@ -/* - * 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.rule.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.rule.action.GetRuleRequest; -import org.opensearch.rule.action.GetRuleResponse; -import org.opensearch.rule.service.RulePersistenceService; -import org.opensearch.tasks.Task; -import org.opensearch.transport.TransportService; - -/** - * Transport action to get workload management Rules - * @opensearch.experimental - */ -public class TransportGetWlmRuleAction extends HandledTransportAction { - - private final RulePersistenceService rulePersistenceService; - - /** - * Constructor for TransportGetWlmRuleAction - * @param transportService - a {@link TransportService} object - * @param actionFilters - a {@link ActionFilters} object - * @param rulePersistenceService - a {@link RulePersistenceService} object - */ - @Inject - public TransportGetWlmRuleAction( - TransportService transportService, - ActionFilters actionFilters, - RulePersistenceService rulePersistenceService - ) { - super(GetWlmRuleAction.NAME, transportService, actionFilters, GetRuleRequest::new); - this.rulePersistenceService = rulePersistenceService; - } - - @Override - protected void doExecute(Task task, GetRuleRequest request, ActionListener listener) { - rulePersistenceService.getRule(request, listener); - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleAction.java new file mode 100644 index 0000000000000..d94b51a19be96 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleAction.java @@ -0,0 +1,63 @@ +/* + * 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.rule.rest; + +import org.opensearch.action.ActionType; +import org.opensearch.autotagging.FeatureType; +import org.opensearch.autotagging.Rule; +import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; +import org.opensearch.plugin.wlm.rule.action.CreateWlmRuleAction; +import org.opensearch.rule.action.CreateRuleRequest; +import org.opensearch.rule.action.CreateRuleResponse; +import org.opensearch.rule.rest.RestCreateRuleAction; + +import java.util.List; + +import static org.opensearch.rest.RestRequest.Method.POST; +import static org.opensearch.rest.RestRequest.Method.PUT; + +/** + * Rest action to create a Rule in workload management + * @opensearch.experimental + */ +public class RestCreateWlmRuleAction extends RestCreateRuleAction { + + /** + * Constructor for RestCreateWlmRuleAction + */ + public RestCreateWlmRuleAction() { + super(); + } + + @Override + public String getName() { + return "create_rule"; + } + + @Override + public List routes() { + return List.of(new Route(POST, "_wlm/rule/"), new Route(PUT, "_wlm/rule/")); + } + + @Override + @SuppressWarnings("unchecked") + protected > T retrieveCreateRuleActionInstance() { + return (T) CreateWlmRuleAction.INSTANCE; + } + + @Override + protected FeatureType retrieveFeatureTypeInstance() { + return QueryGroupFeatureType.INSTANCE; + } + + @Override + protected CreateRuleRequest buildCreateRuleRequest(Rule rule) { + return new CreateRuleRequest(rule); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java deleted file mode 100644 index 437dd8b246091..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleAction.java +++ /dev/null @@ -1,64 +0,0 @@ -/* - * 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.rule.rest; - -import org.opensearch.action.ActionType; -import org.opensearch.autotagging.Attribute; -import org.opensearch.plugin.wlm.rule.QueryGroupAttribute; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.plugin.wlm.rule.action.GetWlmRuleAction; -import org.opensearch.rule.action.GetRuleRequest; -import org.opensearch.rule.action.GetRuleResponse; -import org.opensearch.rule.rest.RestGetRuleAction; - -import java.util.List; -import java.util.Map; -import java.util.Set; - -import static org.opensearch.rest.RestRequest.Method.GET; - -/** - * Rest action to get workload management Rules - * @opensearch.experimental - */ -public class RestGetWlmRuleAction extends RestGetRuleAction { - - /** - * Constructor for RestGetWlmRuleAction - */ - public RestGetWlmRuleAction() { - super(); - } - - @Override - public String getName() { - return "get_rule"; - } - - @Override - public List routes() { - return List.of(new Route(GET, "_wlm/rule/"), new Route(GET, "_wlm/rule/{_id}")); - } - - @Override - protected Attribute getAttributeFromName(String name) { - return QueryGroupAttribute.fromName(name); - } - - @Override - @SuppressWarnings("unchecked") - protected > T retrieveGetRuleActionInstance() { - return (T) GetWlmRuleAction.INSTANCE; - } - - @Override - protected GetRuleRequest buildGetRuleRequest(String id, Map> attributeFilters, String searchAfter) { - return new GetRuleRequest(id, attributeFilters, searchAfter, QueryGroupFeatureType.INSTANCE); - } -} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleActionTests.java similarity index 62% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleActionTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleActionTests.java index a0f61cbaf3c8b..ac4515910422c 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestGetWlmRuleActionTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleActionTests.java @@ -14,23 +14,24 @@ import java.util.List; -import static org.opensearch.rest.RestRequest.Method.GET; +import static org.opensearch.rest.RestRequest.Method.POST; +import static org.opensearch.rest.RestRequest.Method.PUT; -public class RestGetWlmRuleActionTests extends OpenSearchTestCase { +public class RestCreateWlmRuleActionTests extends OpenSearchTestCase { /** * Test case to validate the construction for RestGetRuleAction */ public void testConstruction() { - RestGetWlmRuleAction action = new RestGetWlmRuleAction(); + RestCreateWlmRuleAction action = new RestCreateWlmRuleAction(); assertNotNull(action); - assertEquals("get_rule", action.getName()); + assertEquals("create_rule", action.getName()); List routes = action.routes(); assertEquals(2, routes.size()); RestHandler.Route route = routes.get(0); - assertEquals(GET, route.getMethod()); + assertEquals(POST, route.getMethod()); assertEquals("_wlm/rule/", route.getPath()); route = routes.get(1); - assertEquals(GET, route.getMethod()); - assertEquals("_wlm/rule/{_id}", route.getPath()); + assertEquals(PUT, route.getMethod()); + assertEquals("_wlm/rule/", route.getPath()); } } From b888ec4bd4a435d327d82bae7393126232e5067d Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Fri, 4 Apr 2025 00:55:59 -0700 Subject: [PATCH 11/20] modify based on comments Signed-off-by: Ruirui Zhang --- .../rule/action/CreateRuleRequest.java | 8 + .../rule/action/CreateRuleResponse.java | 23 +-- .../rule/rest/RestCreateRuleAction.java | 52 +++--- .../IndexStoredRulePersistenceService.java | 77 ++++----- .../rule/utils/IndexStoredRuleParser.java | 2 +- .../rule/utils/IndexStoredRuleUtils.java | 27 ++- .../org/opensearch/rule/RuleTestUtils.java | 2 +- .../rule/action/CreateRuleRequestTests.java | 4 +- .../rule/action/CreateRuleResponseTests.java | 14 +- ...ndexStoredRulePersistenceServiceTests.java | 71 ++++++++ .../utils/IndexStoredRuleParserTests.java | 3 +- .../rule/utils/IndexStoredRuleUtilsTests.java | 26 +-- .../opensearch/rule/utils/RuleTestUtils.java | 157 ++++++++++++++++++ .../action/TransportCreateWlmRuleAction.java | 12 +- .../rule/rest/RestCreateWlmRuleAction.java | 63 ------- .../rest/RestCreateWlmRuleActionTests.java | 37 ----- 16 files changed, 366 insertions(+), 212 deletions(-) create mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleAction.java delete mode 100644 plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleActionTests.java diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java index a2956d1a65bad..3eb6e4e9f086d 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java @@ -18,6 +18,14 @@ /** * A request for create Rule + * Example request: + * Note that the endpoint below is for wlm rules specifically and serves only as an example + * curl -XPUT "localhost:9200/_wlm/rule/" -H 'Content-Type: application/json' -d ' + * { + * "description": "description1", + * "index_pattern": ["log*", "event*"], + * "query_group": "poOiU851RwyLYvV5lbvv5w" + * }' * @opensearch.experimental */ public class CreateRuleRequest extends ActionRequest { diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java index aef3936c554c3..ba6c83dc2d8c9 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java @@ -12,7 +12,6 @@ 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; @@ -24,23 +23,28 @@ /** * Response for the create API for Rule + * Example response: + * { + * "_id":"wi6VApYBoX5wstmtU_8l", + * "description":"description1", + * "index_pattern":["log*", "uvent*"], + * "query_group":"poOiU851RwyLYvV5lbvv5w", + * "updated_at":"2025-04-04T20:54:22.406Z" + * } * @opensearch.experimental */ public class CreateRuleResponse extends ActionResponse implements ToXContent, ToXContentObject { private final String _id; private final Rule rule; - private final RestStatus restStatus; /** * contructor for CreateRuleResponse * @param id - the id for the rule created * @param rule - the rule created - * @param restStatus - the status of CreateRuleResponse */ - public CreateRuleResponse(String id, final Rule rule, RestStatus restStatus) { + public CreateRuleResponse(String id, final Rule rule) { this._id = id; this.rule = rule; - this.restStatus = restStatus; } /** @@ -50,14 +54,12 @@ public CreateRuleResponse(String id, final Rule rule, RestStatus restStatus) { public CreateRuleResponse(StreamInput in) throws IOException { _id = in.readString(); rule = new Rule(in); - restStatus = RestStatus.readFrom(in); } @Override public void writeTo(StreamOutput out) throws IOException { out.writeString(_id); rule.writeTo(out); - RestStatus.writeTo(out, restStatus); } @Override @@ -71,11 +73,4 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws public Rule getRule() { return rule; } - - /** - * restStatus getter - */ - public RestStatus getRestStatus() { - return restStatus; - } } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java index a096ac517cb40..a6a6a13609553 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java @@ -10,7 +10,6 @@ import org.opensearch.action.ActionType; import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule; import org.opensearch.autotagging.Rule.Builder; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; @@ -33,24 +32,42 @@ * Rest action to create a Rule * @opensearch.experimental */ -public abstract class RestCreateRuleAction extends BaseRestHandler { +public class RestCreateRuleAction extends BaseRestHandler { + private final String name; + private final List routes; + private final FeatureType featureType; + private final ActionType createRuleAction; + /** - * constructor for RestCreateRuleAction + * constructor for RestUpdateRuleAction + * @param name - RestUpdateRuleAction name + * @param routes the list of REST routes this action handles + * @param featureType the feature type associated with the rule + * @param createRuleAction the action to execute for updating a rule */ - public RestCreateRuleAction() {} + public RestCreateRuleAction(String name, List routes, FeatureType featureType, ActionType createRuleAction) { + this.name = name; + this.routes = routes; + this.featureType = featureType; + this.createRuleAction = createRuleAction; + } @Override - public abstract String getName(); + public String getName() { + return name; + } @Override - public abstract List routes(); + public List routes() { + return routes; + } @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { try (XContentParser parser = request.contentParser()) { - Builder builder = Builder.fromXContent(parser, retrieveFeatureTypeInstance()); - CreateRuleRequest createRuleRequest = buildCreateRuleRequest(builder.updatedAt(Instant.now().toString()).build()); - return channel -> client.execute(retrieveCreateRuleActionInstance(), createRuleRequest, createRuleResponse(channel)); + Builder builder = Builder.fromXContent(parser, featureType); + CreateRuleRequest createRuleRequest = new CreateRuleRequest(builder.updatedAt(Instant.now().toString()).build()); + return channel -> client.execute(createRuleAction, createRuleRequest, createRuleResponse(channel)); } } @@ -62,21 +79,4 @@ public RestResponse buildResponse(final CreateRuleResponse response) throws Exce } }; } - - /** - * Abstract method for subclasses to provide specific ActionType Instance - */ - protected abstract > T retrieveCreateRuleActionInstance(); - - /** - * Abstract method for subclasses to provide specific FeatureType Instance - */ - protected abstract FeatureType retrieveFeatureTypeInstance(); - - /** - * Abstract method for subclasses to construct a {@link CreateRuleRequest}. This method allows subclasses - * to define their own request-building logic depending on their specific needs. - * @param rule - the rule to create - */ - protected abstract CreateRuleRequest buildCreateRuleRequest(Rule rule); } diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java index 75d8141d9a139..6f8d67e51be6a 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -40,6 +40,7 @@ import java.util.Arrays; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; @@ -57,7 +58,7 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService private final ClusterService clusterService; private final Client client; private final FeatureType featureType; - private final int maxRulesPerGetRequest; + private final int maxRulesPerPage; private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class); private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); @@ -68,26 +69,27 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService * @param clusterService - The clusterService used in IndexStoredRulePersistenceService. * @param client - The OpenSearch client used to interact with the OpenSearch cluster. * @param featureType - The feature type associated with the stored rules. - * @param maxRulesPerGetRequest - The maximum number of rules that can be returned in a single get request. + * @param maxRulesPerPage - The maximum number of rules that can be returned in a single get request. */ public IndexStoredRulePersistenceService( String indexName, ClusterService clusterService, Client client, FeatureType featureType, - int maxRulesPerGetRequest + int maxRulesPerPage ) { this.indexName = indexName; this.clusterService = clusterService; this.client = client; this.featureType = featureType; - this.maxRulesPerGetRequest = maxRulesPerGetRequest; + this.maxRulesPerPage = maxRulesPerPage; } /** - * Entry point for the create rule api logic in persistence service - * @param request - The CreateRuleRequest - * @param listener - ActionListener for CreateRuleResponse + * Entry point for the create rule API logic in persistence service. + * It ensures the index exists, validates for duplicate rules, and persists the new rule. + * @param request The CreateRuleRequest + * @param listener ActionListener for CreateRuleResponse */ public void createRule(CreateRuleRequest request, ActionListener listener) { try (ThreadContext.StoredContext ctx = getContext()) { @@ -98,7 +100,17 @@ public void onResponse(Boolean indexCreated) { listener.onFailure(new IllegalStateException(indexName + " index creation failed and rule cannot be persisted")); return; } - checkDuplicateRule(request.getRule(), listener); + validateNoDuplicateRule(request.getRule(), new ActionListener<>() { + @Override + public void onResponse(Void unused) { + persistRule(request.getRule(), listener); + } + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + }); } @Override @@ -110,30 +122,21 @@ public void onFailure(Exception e) { } /** - * Check if there's an existing Rule with the same attributes. - * For example, if there's an existing Rule with the attribute index_pattern: ["a", "b", "c"], - * then we cannot create another Rule with only one attribute index_pattern: ["b"], because the value "b" - * already exists under another Rule. Note that the conflict exists only when we have the exact same attribute - * names in the two rules (That is, a Rule with attribute "index_pattern" won't create a conflict with another - * Rule that has "index_pattern" and some other attributes). - * @param rule - The rule to update. - * @param listener - ActionListener for CreateRuleResponse + * Validates that no duplicate rule exists with the same attribute map. + * If a conflict is found, fails the listener + * @param rule - the rule we check duplicate against + * @param listener - listener for validateNoDuplicateRule response */ - public void checkDuplicateRule(Rule rule, ActionListener listener) { + private void validateNoDuplicateRule(Rule rule, ActionListener listener) { try (ThreadContext.StoredContext ctx = getContext()) { getRuleFromIndex(null, rule.getAttributeMap(), null, new ActionListener<>() { @Override public void onResponse(GetRuleResponse getRuleResponse) { - String duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(rule.getAttributeMap(), getRuleResponse.getRules()); - if (duplicateRuleId != null) { - listener.onFailure( - new IllegalArgumentException( - "A rule that has the same attribute values already exists under rule id " + duplicateRuleId - ) - ); - return; - } - persistRule(rule, listener); + Optional duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(rule, getRuleResponse.getRules()); + duplicateRuleId.ifPresentOrElse( + id -> listener.onFailure(new IllegalArgumentException("Rule already exists under rule id " + id)), + () -> listener.onResponse(null) + ); } @Override @@ -149,28 +152,28 @@ public void onFailure(Exception e) { * @param rule - The rule to update. * @param listener - ActionListener for CreateRuleResponse */ - public void persistRule(Rule rule, ActionListener listener) { + private void persistRule(Rule rule, ActionListener listener) { try (ThreadContext.StoredContext ctx = getContext()) { IndexRequest indexRequest = new IndexRequest(indexName).source( rule.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS) ); client.index(indexRequest, ActionListener.wrap(indexResponse -> { - listener.onResponse(new CreateRuleResponse(indexResponse.getId(), rule, RestStatus.OK)); + listener.onResponse(new CreateRuleResponse(indexResponse.getId(), rule)); }, e -> { logger.warn("Failed to save Rule object due to error: {}", e.getMessage()); listener.onFailure(e); })); } catch (IOException e) { - logger.error("Error saving rule to index: {}", indexName, e); + logger.error("Error saving rule to index: {}", indexName); listener.onFailure(new RuntimeException("Failed to save rule to index.")); } } /** - * Creates the system index .rules if it doesn't exist + * Creates the system index if it doesn't exist * @param listener - ActionListener for CreateRuleResponse */ - void createIndexIfAbsent(ActionListener listener) { + private void createIndexIfAbsent(ActionListener listener) { try (ThreadContext.StoredContext ctx = getContext()) { if (clusterService.state().metadata().hasIndex(indexName)) { listener.onResponse(true); @@ -187,7 +190,7 @@ public void onResponse(CreateIndexResponse response) { @Override public void onFailure(Exception e) { if (e instanceof ResourceAlreadyExistsException) { - logger.info("Index {} already exists", indexName); + logger.trace("Index {} already exists", indexName); listener.onResponse(true); } else { logger.error("Failed to create index {}: {}", indexName, e.getMessage()); @@ -206,7 +209,7 @@ public void onFailure(Exception e) { * @param searchAfter - The sort values from the last document of the previous page, used for pagination * @param listener - ActionListener for GetRuleResponse */ - public void getRuleFromIndex( + private void getRuleFromIndex( String id, Map> attributeFilters, String searchAfter, @@ -217,7 +220,7 @@ public void getRuleFromIndex( // actions within this block are trusted and executed with system-level privileges. try (ThreadContext.StoredContext context = getContext()) { BoolQueryBuilder boolQuery = IndexStoredRuleUtils.buildGetRuleQuery(id, attributeFilters, featureType); - SearchRequestBuilder searchRequest = client.prepareSearch(indexName).setQuery(boolQuery).setSize(maxRulesPerGetRequest); + SearchRequestBuilder searchRequest = client.prepareSearch(indexName).setQuery(boolQuery).setSize(maxRulesPerPage); if (searchAfter != null) { searchRequest.addSort(_ID_STRING, SortOrder.ASC).searchAfter(new Object[] { searchAfter }); } @@ -233,11 +236,11 @@ public void getRuleFromIndex( * @param searchResponse - Response received from index * @param listener - ActionListener for GetRuleResponse */ - void handleGetRuleResponse(String id, SearchResponse searchResponse, ActionListener listener) { + private void handleGetRuleResponse(String id, SearchResponse searchResponse, ActionListener listener) { List hits = Arrays.asList(searchResponse.getHits().getHits()); if (id != null && hits.isEmpty()) { logger.error("Rule with ID " + id + " not found."); - listener.onFailure(new ResourceNotFoundException("Rule with ID " + id + " doesn't exist in the .rules index.")); + listener.onFailure(new ResourceNotFoundException("Rule with ID " + id + " doesn't exist in the index.")); return; } Map ruleMap = hits.stream() diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java index 4d2e808a22c2d..ccf48cceb645a 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java @@ -28,7 +28,7 @@ public class IndexStoredRuleParser { /** * constructor for IndexStoredRuleParser */ - public IndexStoredRuleParser() {} + private IndexStoredRuleParser() {} private static final Logger logger = LogManager.getLogger(IndexStoredRuleParser.class); diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java index 27af328c07e60..4b52f85f0cd24 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java +++ b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java @@ -15,6 +15,7 @@ import org.opensearch.index.query.QueryBuilders; import java.util.Map; +import java.util.Optional; import java.util.Set; import static org.opensearch.autotagging.Rule._ID_STRING; @@ -28,7 +29,7 @@ public class IndexStoredRuleUtils { /** * constructor for IndexStoredRuleUtils */ - public IndexStoredRuleUtils() {} + private IndexStoredRuleUtils() {} /** * Builds a Boolean query to retrieve a rule by its ID or attribute filters. @@ -57,18 +58,30 @@ public static BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeMapToValidate, Map ruleMap) { + public static Optional getDuplicateRuleId(Rule rule, Map ruleMap) { + Map> attributeMapToValidate = rule.getAttributeMap(); for (Map.Entry entry : ruleMap.entrySet()) { String ruleId = entry.getKey(); Rule currRule = entry.getValue(); + // Compare the size of the attribute maps to ensure we only check for duplicates with the same number of attributes. if (attributeMapToValidate.size() == currRule.getAttributeMap().size()) { - return ruleId; + return Optional.of(ruleId); } } - return null; + return Optional.empty(); } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java index c2c5f12b5aa1c..fd9a228370af4 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.rule; +package org.opensearch.rule.utils; import org.opensearch.autotagging.Attribute; import org.opensearch.autotagging.AutoTaggingRegistry; diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java index ec415fbd9c0a2..4ebdf296cf1c0 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java @@ -14,8 +14,8 @@ import java.io.IOException; -import static org.opensearch.rule.RuleTestUtils.assertEqualRule; -import static org.opensearch.rule.RuleTestUtils.ruleOne; +import static org.opensearch.rule.utils.RuleTestUtils.assertEqualRule; +import static org.opensearch.rule.utils.RuleTestUtils.ruleOne; public class CreateRuleRequestTests extends OpenSearchTestCase { diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java index 8ec92ccada46f..17de247c09564 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java @@ -12,7 +12,6 @@ import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.test.OpenSearchTestCase; @@ -20,9 +19,9 @@ import java.io.IOException; import java.util.Map; -import static org.opensearch.rule.RuleTestUtils._ID_ONE; -import static org.opensearch.rule.RuleTestUtils.assertEqualRules; -import static org.opensearch.rule.RuleTestUtils.ruleOne; +import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.assertEqualRules; +import static org.opensearch.rule.utils.RuleTestUtils.ruleOne; import static org.mockito.Mockito.mock; public class CreateRuleResponseTests extends OpenSearchTestCase { @@ -31,12 +30,11 @@ public class CreateRuleResponseTests extends OpenSearchTestCase { * Test case to verify serialization and deserialization of CreateRuleResponse */ public void testSerialization() throws IOException { - CreateRuleResponse response = new CreateRuleResponse(_ID_ONE, ruleOne, RestStatus.OK); + CreateRuleResponse response = new CreateRuleResponse(_ID_ONE, ruleOne); BytesStreamOutput out = new BytesStreamOutput(); response.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); CreateRuleResponse otherResponse = new CreateRuleResponse(streamInput); - assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); Rule responseRule = response.getRule(); Rule otherResponseRule = otherResponse.getRule(); assertEqualRules(Map.of(_ID_ONE, responseRule), Map.of(_ID_ONE, otherResponseRule), false); @@ -47,13 +45,13 @@ public void testSerialization() throws IOException { */ public void testToXContentCreateRule() throws IOException { XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); - CreateRuleResponse response = new CreateRuleResponse(_ID_ONE, ruleOne, RestStatus.OK); + CreateRuleResponse response = new CreateRuleResponse(_ID_ONE, ruleOne); String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); String expected = "{\n" + " \"_id\" : \"AgfUO5Ja9yfvhdONlYi3TQ==\",\n" + " \"description\" : \"description_1\",\n" + " \"mock_attribute_one\" : [\n" - + " \"pattern_1\"\n" + + " \"mock_attribute_one\"\n" + " ],\n" + " \"mock_feature_type\" : \"feature_value_one\",\n" + " \"updated_at\" : \"2024-01-26T08:58:57.558Z\"\n" diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java index c0de380d05580..a813fe3ce7880 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java @@ -214,4 +214,75 @@ public void testDeleteRule_notFound() { verify(listener).onFailure(any(ResourceNotFoundException.class)); } + + // public void testCreateIndexIfAbsent_IndexExists() { + // IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); + // ClusterService clusterService = rulePersistenceService.getClusterService(); + // when(clusterService.state().metadata().hasIndex(TEST_INDEX_NAME)).thenReturn(true); + // ActionListener listener = mock(ActionListener.class); + // rulePersistenceService.createIndexIfAbsent(listener); + // verify(listener).onResponse(true); + // verifyNoMoreInteractions(listener); + // } + // + // public void testCreateIndexIfAbsent() { + // IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); + // Client client = rulePersistenceService.getClient(); + // IndicesAdminClient indicesAdminClient = mock(IndicesAdminClient.class); + // when(client.admin()).thenReturn(mock(AdminClient.class)); + // when(client.admin().indices()).thenReturn(indicesAdminClient); + // doAnswer(invocation -> { + // ActionListener listener = invocation.getArgument(1); + // listener.onResponse(new CreateIndexResponse(true, true, TEST_INDEX_NAME)); // Assuming the index creation was successful + // return null; + // }).when(indicesAdminClient).create(any(CreateIndexRequest.class), any(ActionListener.class)); + // rulePersistenceService.createIndexIfAbsent(new ActionListener<>() { + // @Override + // public void onResponse(Boolean indexCreated) { + // assertTrue(indexCreated); + // } + // + // @Override + // public void onFailure(Exception e) { + // fail("Index creation failed: " + e.getMessage()); + // } + // }); + // verify(indicesAdminClient).create(any(CreateIndexRequest.class), any(ActionListener.class)); + // } + // + // public void testPersistRuleSuccess() { + // IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); + // Client client = rulePersistenceService.getClient(); + // ActionListener listener = mock(ActionListener.class); + // IndexResponse indexResponse = new IndexResponse(new ShardId(TEST_INDEX_NAME, "uuid", 0), "id", 1, 1, 1, true); + // doAnswer(invocation -> { + // ActionListener actionListener = invocation.getArgument(1); + // actionListener.onResponse(indexResponse); + // return null; + // }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); + // + // rulePersistenceService.persistRule(ruleOne, listener); + // verify(client).index(any(IndexRequest.class), any(ActionListener.class)); + // ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(CreateRuleResponse.class); + // verify(listener).onResponse(responseCaptor.capture()); + // + // CreateRuleResponse createRuleResponse = responseCaptor.getValue(); + // assertNotNull(createRuleResponse); + // assertEquals(ruleOne, createRuleResponse.getRule()); + // } + // + // public void testPersistRuleFailure() throws IOException { + // IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); + // Client client = rulePersistenceService.getClient(); + // ActionListener listener = mock(ActionListener.class); + // doAnswer(invocation -> { + // ActionListener actionListener = invocation.getArgument(1); + // actionListener.onFailure(new RuntimeException("Indexing failed")); + // return null; + // }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); + // + // rulePersistenceService.persistRule(ruleOne, listener); + // verify(client).index(any(IndexRequest.class), any(ActionListener.class)); + // verify(listener).onFailure(any(RuntimeException.class)); + // } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java index 277fce2e2dc72..18e1bb99926b7 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java @@ -9,14 +9,13 @@ package org.opensearch.rule.utils; import org.opensearch.autotagging.Rule; -import org.opensearch.rule.RuleTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; import java.time.Instant; import java.util.Locale; -import static org.opensearch.rule.RuleTestUtils.DESCRIPTION_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.DESCRIPTION_ONE; public class IndexStoredRuleParserTests extends OpenSearchTestCase { public static final String VALID_JSON = String.format(Locale.ROOT, """ diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java index b8bafa7e7d33c..54c6745566c11 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java @@ -9,20 +9,23 @@ package org.opensearch.rule.utils; import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.Rule; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilder; -import org.opensearch.rule.RuleTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.util.HashMap; import java.util.Map; +import java.util.Optional; import java.util.Set; -import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_MAP; -import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_VALUE_ONE; -import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_VALUE_TWO; -import static org.opensearch.rule.RuleTestUtils._ID_ONE; -import static org.opensearch.rule.RuleTestUtils.ruleOne; +import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_MAP; +import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_VALUE_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_VALUE_TWO; +import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.ruleOne; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class IndexStoredRuleUtilsTests extends OpenSearchTestCase { public void testBuildGetRuleQuery_WithId() { @@ -46,18 +49,21 @@ public void testBuildGetRuleQuery_WithAttributes() { } public void testGetDuplicateRuleId_Found() { - String duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(ATTRIBUTE_MAP, Map.of(_ID_ONE, ruleOne)); - assertEquals(_ID_ONE, duplicateRuleId); + Optional duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(ruleOne, Map.of(_ID_ONE, ruleOne)); + assertFalse(duplicateRuleId.isEmpty()); + assertEquals(_ID_ONE, duplicateRuleId.get()); } public void testGetDuplicateRuleId_NotFound() { + Rule rule = mock(Rule.class); Map> map = Map.of( RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, Set.of(ATTRIBUTE_VALUE_ONE), RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(ATTRIBUTE_VALUE_TWO) ); - String duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(map, Map.of(_ID_ONE, ruleOne)); - assertNull(duplicateRuleId); + when(rule.getAttributeMap()).thenReturn(map); + Optional duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(rule, Map.of(_ID_ONE, ruleOne)); + assertTrue(duplicateRuleId.isEmpty()); } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java new file mode 100644 index 0000000000000..fd9a228370af4 --- /dev/null +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java @@ -0,0 +1,157 @@ +/* + * 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.rule.utils; + +import org.opensearch.autotagging.Attribute; +import org.opensearch.autotagging.AutoTaggingRegistry; +import org.opensearch.autotagging.FeatureType; +import org.opensearch.autotagging.Rule; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.rule.service.IndexStoredRulePersistenceService; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.client.Client; + +import java.util.Map; +import java.util.Set; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class RuleTestUtils { + public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; + public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; + public static final String FEATURE_VALUE_ONE = "feature_value_one"; + public static final String FEATURE_VALUE_TWO = "feature_value_two"; + public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; + public static final String ATTRIBUTE_VALUE_TWO = "mock_attribute_two"; + public static final String DESCRIPTION_ONE = "description_1"; + public static final String DESCRIPTION_TWO = "description_2"; + public static final String SEARCH_AFTER = "search_after_id"; + public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; + public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; + public static final String FEATURE_TYPE_NAME = "mock_feature_type"; + public static final String TEST_INDEX_NAME = ".test_index_for_rule"; + public static final Map> ATTRIBUTE_MAP = Map.of( + MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, + Set.of(ATTRIBUTE_VALUE_ONE) + ); + public static final Rule ruleOne = Rule.builder() + .description(DESCRIPTION_ONE) + .featureType(MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_ONE) + .attributeMap(ATTRIBUTE_MAP) + .updatedAt(TIMESTAMP_ONE) + .build(); + + public static final Rule ruleTwo = Rule.builder() + .description(DESCRIPTION_TWO) + .featureType(MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_TWO) + .attributeMap(Map.of(MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(ATTRIBUTE_VALUE_TWO))) + .updatedAt(TIMESTAMP_TWO) + .build(); + + public static Map ruleMap() { + return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); + } + + public static IndexStoredRulePersistenceService setUpIndexStoredRulePersistenceService(Map queryGroupMap) { + Client client = mock(Client.class); + ClusterService clusterService = mock(ClusterService.class); + ClusterState clusterState = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + ThreadPool threadPool = mock(ThreadPool.class); + + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + when(client.threadPool()).thenReturn(threadPool); + when(threadPool.getThreadContext()).thenReturn(threadContext); + when(clusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + when(metadata.queryGroups()).thenReturn(queryGroupMap); + return new IndexStoredRulePersistenceService(TEST_INDEX_NAME, clusterService, client, MockRuleFeatureType.INSTANCE, 50); + } + + public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { + assertEquals(mapOne.size(), mapTwo.size()); + for (Map.Entry entry : mapOne.entrySet()) { + String id = entry.getKey(); + assertTrue(mapTwo.containsKey(id)); + Rule one = mapOne.get(id); + Rule two = mapTwo.get(id); + assertEqualRule(one, two, ruleUpdated); + } + } + + public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { + if (ruleUpdated) { + assertEquals(one.getDescription(), two.getDescription()); + assertEquals(one.getFeatureType(), two.getFeatureType()); + assertEquals(one.getFeatureValue(), two.getFeatureValue()); + assertEquals(one.getAttributeMap(), two.getAttributeMap()); + assertEquals(one.getAttributeMap(), two.getAttributeMap()); + } else { + assertEquals(one, two); + } + } + + public static class MockRuleFeatureType implements FeatureType { + + public static final MockRuleFeatureType INSTANCE = new MockRuleFeatureType(); + + private MockRuleFeatureType() {} + + static { + INSTANCE.registerFeatureType(); + } + + @Override + public String getName() { + return FEATURE_TYPE_NAME; + } + + @Override + public Map getAllowedAttributesRegistry() { + return Map.of( + ATTRIBUTE_VALUE_ONE, + MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, + ATTRIBUTE_VALUE_TWO, + MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO + ); + } + + @Override + public void registerFeatureType() { + AutoTaggingRegistry.registerFeatureType(INSTANCE); + } + } + + public enum MockRuleAttributes implements Attribute { + MOCK_RULE_ATTRIBUTE_ONE(ATTRIBUTE_VALUE_ONE), + MOCK_RULE_ATTRIBUTE_TWO(ATTRIBUTE_VALUE_TWO); + ; + + private final String name; + + MockRuleAttributes(String name) { + this.name = name; + } + + @Override + public String getName() { + return name; + } + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java index 613a6f6760ab8..07c6e8a0d1ec0 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java @@ -11,11 +11,11 @@ import org.opensearch.ResourceNotFoundException; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; import org.opensearch.rule.action.CreateRuleRequest; import org.opensearch.rule.action.CreateRuleResponse; -import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.rule.service.RulePersistenceService; import org.opensearch.tasks.Task; import org.opensearch.transport.TransportService; @@ -26,29 +26,33 @@ */ public class TransportCreateWlmRuleAction extends HandledTransportAction { - private final IndexStoredRulePersistenceService rulePersistenceService; + private final RulePersistenceService rulePersistenceService; + private final ClusterService clusterService; /** * Constructor for TransportCreateWlmRuleAction * * @param transportService - a {@link TransportService} object * @param actionFilters - a {@link ActionFilters} object + * @param clusterService - a {@link ClusterService} object * @param rulePersistenceService - a {@link RulePersistenceService} object */ @Inject public TransportCreateWlmRuleAction( TransportService transportService, ActionFilters actionFilters, - IndexStoredRulePersistenceService rulePersistenceService + ClusterService clusterService, + RulePersistenceService rulePersistenceService ) { super(CreateWlmRuleAction.NAME, transportService, actionFilters, CreateRuleRequest::new); + this.clusterService = clusterService; this.rulePersistenceService = rulePersistenceService; } @Override protected void doExecute(Task task, CreateRuleRequest request, ActionListener listener) { String queryGroupId = request.getRule().getFeatureValue(); - if (!rulePersistenceService.getClusterService().state().metadata().queryGroups().containsKey(queryGroupId)) { + if (!clusterService.state().metadata().queryGroups().containsKey(queryGroupId)) { listener.onFailure(new ResourceNotFoundException("Couldn't find an existing query group with id: " + queryGroupId)); return; } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleAction.java deleted file mode 100644 index d94b51a19be96..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleAction.java +++ /dev/null @@ -1,63 +0,0 @@ -/* - * 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.rule.rest; - -import org.opensearch.action.ActionType; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule; -import org.opensearch.plugin.wlm.rule.QueryGroupFeatureType; -import org.opensearch.plugin.wlm.rule.action.CreateWlmRuleAction; -import org.opensearch.rule.action.CreateRuleRequest; -import org.opensearch.rule.action.CreateRuleResponse; -import org.opensearch.rule.rest.RestCreateRuleAction; - -import java.util.List; - -import static org.opensearch.rest.RestRequest.Method.POST; -import static org.opensearch.rest.RestRequest.Method.PUT; - -/** - * Rest action to create a Rule in workload management - * @opensearch.experimental - */ -public class RestCreateWlmRuleAction extends RestCreateRuleAction { - - /** - * Constructor for RestCreateWlmRuleAction - */ - public RestCreateWlmRuleAction() { - super(); - } - - @Override - public String getName() { - return "create_rule"; - } - - @Override - public List routes() { - return List.of(new Route(POST, "_wlm/rule/"), new Route(PUT, "_wlm/rule/")); - } - - @Override - @SuppressWarnings("unchecked") - protected > T retrieveCreateRuleActionInstance() { - return (T) CreateWlmRuleAction.INSTANCE; - } - - @Override - protected FeatureType retrieveFeatureTypeInstance() { - return QueryGroupFeatureType.INSTANCE; - } - - @Override - protected CreateRuleRequest buildCreateRuleRequest(Rule rule) { - return new CreateRuleRequest(rule); - } -} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleActionTests.java deleted file mode 100644 index ac4515910422c..0000000000000 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/rest/RestCreateWlmRuleActionTests.java +++ /dev/null @@ -1,37 +0,0 @@ - -/* - * 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.rule.rest; - -import org.opensearch.rest.RestHandler; -import org.opensearch.test.OpenSearchTestCase; - -import java.util.List; - -import static org.opensearch.rest.RestRequest.Method.POST; -import static org.opensearch.rest.RestRequest.Method.PUT; - -public class RestCreateWlmRuleActionTests extends OpenSearchTestCase { - /** - * Test case to validate the construction for RestGetRuleAction - */ - public void testConstruction() { - RestCreateWlmRuleAction action = new RestCreateWlmRuleAction(); - assertNotNull(action); - assertEquals("create_rule", action.getName()); - List routes = action.routes(); - assertEquals(2, routes.size()); - RestHandler.Route route = routes.get(0); - assertEquals(POST, route.getMethod()); - assertEquals("_wlm/rule/", route.getPath()); - route = routes.get(1); - assertEquals(PUT, route.getMethod()); - assertEquals("_wlm/rule/", route.getPath()); - } -} From 1706b8aa4072e983b4763e5747547dfb88be5a24 Mon Sep 17 00:00:00 2001 From: Lingxi Chen Date: Tue, 15 Apr 2025 06:02:06 -0700 Subject: [PATCH 12/20] rebase issues Signed-off-by: Lingxi Chen --- .../plugin/wlm/WorkloadManagementPluginModule.java | 2 +- .../action/CreateWorkloadGroupAction.java | 2 +- .../action/CreateWorkloadGroupRequest.java | 2 +- .../action/CreateWorkloadGroupResponse.java | 2 +- .../action/DeleteWorkloadGroupAction.java | 2 +- .../action/DeleteWorkloadGroupRequest.java | 2 +- .../action/GetWorkloadGroupAction.java | 2 +- .../action/GetWorkloadGroupRequest.java | 2 +- .../action/GetWorkloadGroupResponse.java | 2 +- .../action/TransportCreateWorkloadGroupAction.java | 6 +----- .../action/TransportDeleteWorkloadGroupAction.java | 6 +----- .../action/TransportGetWorkloadGroupAction.java | 6 +----- .../action/TransportUpdateWorkloadGroupAction.java | 6 +----- .../action/UpdateWorkloadGroupAction.java | 2 +- .../action/UpdateWorkloadGroupRequest.java | 0 .../action/UpdateWorkloadGroupResponse.java | 2 +- .../wlm/{querygroup => }/action/package-info.java | 2 +- .../rest/RestCreateWorkloadGroupAction.java | 8 +------- .../rest/RestDeleteWorkloadGroupAction.java | 2 +- .../rest/RestGetWorkloadGroupAction.java | 8 +------- .../rest/RestUpdateWorkloadGroupAction.java | 8 +------- .../plugin/wlm/{querygroup => }/rest/package-info.java | 2 +- .../wlm/service/WorkloadGroupPersistenceService.java | 10 +--------- .../wlm/{querygroup => }/service/package-info.java | 2 +- .../opensearch/plugin/wlm/WorkloadGroupTestUtils.java | 2 +- .../action/CreateWorkloadGroupRequestTests.java | 2 +- .../action/CreateWorkloadGroupResponseTests.java | 2 +- .../action/DeleteWorkloadGroupRequestTests.java | 2 +- .../action/GetWorkloadGroupRequestTests.java | 2 +- .../action/GetWorkloadGroupResponseTests.java | 2 +- .../TransportDeleteWorkloadGroupActionTests.java | 6 +----- .../action/TransportGetWorkloadGroupActionTests.java | 2 +- .../action/UpdateWorkloadGroupRequestTests.java | 2 +- .../action/UpdateWorkloadGroupResponseTests.java | 2 +- .../action/WorkloadGroupActionTestUtils.java | 0 .../rest/RestDeleteWorkloadGroupActionTests.java | 2 +- .../service/WorkloadGroupPersistenceServiceTests.java | 0 37 files changed, 34 insertions(+), 80 deletions(-) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/CreateWorkloadGroupAction.java (94%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/CreateWorkloadGroupRequest.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/CreateWorkloadGroupResponse.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/DeleteWorkloadGroupAction.java (94%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/DeleteWorkloadGroupRequest.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/GetWorkloadGroupAction.java (94%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/GetWorkloadGroupRequest.java (96%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/GetWorkloadGroupResponse.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/TransportCreateWorkloadGroupAction.java (87%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/TransportDeleteWorkloadGroupAction.java (88%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/TransportGetWorkloadGroupAction.java (89%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/TransportUpdateWorkloadGroupAction.java (87%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/UpdateWorkloadGroupAction.java (94%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/UpdateWorkloadGroupRequest.java (100%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/UpdateWorkloadGroupResponse.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/action/package-info.java (84%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/rest/RestCreateWorkloadGroupAction.java (82%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/rest/RestDeleteWorkloadGroupAction.java (97%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/rest/RestGetWorkloadGroupAction.java (80%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/rest/RestUpdateWorkloadGroupAction.java (82%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/rest/package-info.java (85%) rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/{querygroup => }/service/package-info.java (84%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/action/CreateWorkloadGroupRequestTests.java (96%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/action/CreateWorkloadGroupResponseTests.java (98%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/action/DeleteWorkloadGroupRequestTests.java (96%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/action/GetWorkloadGroupRequestTests.java (97%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/action/GetWorkloadGroupResponseTests.java (99%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/action/TransportDeleteWorkloadGroupActionTests.java (84%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/action/TransportGetWorkloadGroupActionTests.java (97%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/action/UpdateWorkloadGroupRequestTests.java (98%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/action/UpdateWorkloadGroupResponseTests.java (98%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/action/WorkloadGroupActionTestUtils.java (100%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/rest/RestDeleteWorkloadGroupActionTests.java (98%) rename plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/{querygroup => }/service/WorkloadGroupPersistenceServiceTests.java (100%) diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java index 77d938a85fdaa..f3ccf9795e93b 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java @@ -10,7 +10,7 @@ import org.opensearch.common.inject.AbstractModule; import org.opensearch.common.inject.Singleton; -import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService;; +import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; import org.opensearch.plugin.wlm.rule.service.WlmRulePersistenceService; import org.opensearch.plugin.wlm.rule.service.WlmRuleProcessingService; import org.opensearch.plugin.wlm.rule.service.WlmRuleResponseBuilder; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupAction.java similarity index 94% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupAction.java index 433305f7e9bf6..ca9784ebc7e4b 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionType; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequest.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequest.java index b31cc6f0fb1d2..ad932667b25e8 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequest.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponse.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponse.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponse.java index d765b3ed14b65..b33214e042398 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponse.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.core.action.ActionResponse; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupAction.java similarity index 94% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupAction.java index 54755758a15a3..39b47d69776f4 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionType; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequest.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequest.java index 73a6cc0e1dcd7..940a3815b1662 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequest.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.action.support.clustermanager.AcknowledgedRequest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupAction.java similarity index 94% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupAction.java index 28f35ba2f18b4..ee1b40a2f9bbc 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionType; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequest.java similarity index 96% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequest.java index 28ab53849824e..4b8a5f85fd236 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequest.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequest.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.action.support.clustermanager.ClusterManagerNodeReadRequest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponse.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponse.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponse.java index 22b66bca9ae56..ab8f773088a37 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponse.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.core.action.ActionResponse; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportCreateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java similarity index 87% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportCreateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java index ae401af9769ee..2039f1cb590ff 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportCreateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction; @@ -17,11 +17,7 @@ import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; -<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportCreateWorkloadGroupAction.java import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; -======== -import org.opensearch.plugin.wlm.querygroup.service.QueryGroupPersistenceService; ->>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportUpdateQueryGroupAction.java import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupAction.java similarity index 88% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupAction.java index 40c82b9c0a56f..2bfbadba4d51d 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; @@ -19,11 +19,7 @@ import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; -<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupAction.java import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; -======== -import org.opensearch.plugin.wlm.querygroup.service.QueryGroupPersistenceService; ->>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteQueryGroupAction.java import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupAction.java similarity index 89% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupAction.java index ce64e737337cd..bb2fbab047343 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -23,11 +23,7 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.rest.RestStatus; -<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupAction.java import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; -======== -import org.opensearch.plugin.wlm.querygroup.service.QueryGroupPersistenceService; ->>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetQueryGroupAction.java import org.opensearch.search.pipeline.SearchPipelineService; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportUpdateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java similarity index 87% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportUpdateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java index b29896c0a0b1f..ef639d44b4155 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportUpdateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction; @@ -17,11 +17,7 @@ import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; -<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportUpdateWorkloadGroupAction.java import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; -======== -import org.opensearch.plugin.wlm.querygroup.service.QueryGroupPersistenceService; ->>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/TransportCreateQueryGroupAction.java import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupAction.java similarity index 94% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupAction.java index 6e49008ff2fae..b4f8e1ce90126 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionType; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupRequest.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupRequest.java similarity index 100% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupRequest.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupRequest.java diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponse.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponse.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponse.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponse.java index 2f9d543d5b1f9..9b8fccbdb5346 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponse.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponse.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.core.action.ActionResponse; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java similarity index 84% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/package-info.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java index 472b41716b44d..df30f55a99b3c 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/action/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/package-info.java @@ -9,4 +9,4 @@ /** * Package for the action classes related to query groups in WorkloadManagementPlugin */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestCreateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateWorkloadGroupAction.java similarity index 82% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestCreateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateWorkloadGroupAction.java index 94d4d82cfbb8e..5ef59602f7893 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestCreateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestCreateWorkloadGroupAction.java @@ -6,20 +6,14 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.rest; +package org.opensearch.plugin.wlm.rest; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentParser; -<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestCreateWorkloadGroupAction.java import org.opensearch.plugin.wlm.action.CreateWorkloadGroupAction; import org.opensearch.plugin.wlm.action.CreateWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.CreateWorkloadGroupResponse; -======== -import org.opensearch.plugin.wlm.querygroup.action.UpdateQueryGroupAction; -import org.opensearch.plugin.wlm.querygroup.action.UpdateQueryGroupRequest; -import org.opensearch.plugin.wlm.querygroup.action.UpdateQueryGroupResponse; ->>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestUpdateQueryGroupAction.java import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupAction.java similarity index 97% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupAction.java index e3b5b8c30ea3b..d0d82f43679fa 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupAction.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.rest; +package org.opensearch.plugin.wlm.rest; import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupAction; import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupRequest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestGetWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetWorkloadGroupAction.java similarity index 80% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestGetWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetWorkloadGroupAction.java index 50dd7e0ea50db..818531352f4d3 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestGetWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestGetWorkloadGroupAction.java @@ -6,19 +6,13 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.rest; +package org.opensearch.plugin.wlm.rest; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; -<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestGetWorkloadGroupAction.java import org.opensearch.plugin.wlm.action.GetWorkloadGroupAction; import org.opensearch.plugin.wlm.action.GetWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.GetWorkloadGroupResponse; -======== -import org.opensearch.plugin.wlm.querygroup.action.GetQueryGroupAction; -import org.opensearch.plugin.wlm.querygroup.action.GetQueryGroupRequest; -import org.opensearch.plugin.wlm.querygroup.action.GetQueryGroupResponse; ->>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestGetQueryGroupAction.java import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestUpdateWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateWorkloadGroupAction.java similarity index 82% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestUpdateWorkloadGroupAction.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateWorkloadGroupAction.java index 1294529d857cb..db77dc5963037 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestUpdateWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/RestUpdateWorkloadGroupAction.java @@ -6,20 +6,14 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.rest; +package org.opensearch.plugin.wlm.rest; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentParser; -<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestUpdateWorkloadGroupAction.java import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupAction; import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupResponse; -======== -import org.opensearch.plugin.wlm.querygroup.action.CreateQueryGroupAction; -import org.opensearch.plugin.wlm.querygroup.action.CreateQueryGroupRequest; -import org.opensearch.plugin.wlm.querygroup.action.CreateQueryGroupResponse; ->>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestCreateQueryGroupAction.java import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java similarity index 85% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/package-info.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java index a51d67a6cb3b7..889f3e107db07 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rest/package-info.java @@ -9,4 +9,4 @@ /** * Package for the rest classes related to query groups in WorkloadManagementPlugin */ -package org.opensearch.plugin.wlm.querygroup.rest; +package org.opensearch.plugin.wlm.rest; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java index 02d4b5d7c97d2..8fd5fe5dfcfed 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.service; +package org.opensearch.plugin.wlm.service; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -27,19 +27,11 @@ import org.opensearch.common.settings.Settings; import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; -<<<<<<<< HEAD:plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java import org.opensearch.plugin.wlm.action.CreateWorkloadGroupResponse; import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupRequest; import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupResponse; import org.opensearch.wlm.MutableWorkloadGroupFragment; -======== -import org.opensearch.plugin.wlm.querygroup.action.CreateQueryGroupResponse; -import org.opensearch.plugin.wlm.querygroup.action.DeleteQueryGroupRequest; -import org.opensearch.plugin.wlm.querygroup.action.UpdateQueryGroupRequest; -import org.opensearch.plugin.wlm.querygroup.action.UpdateQueryGroupResponse; -import org.opensearch.wlm.MutableQueryGroupFragment; ->>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/service/QueryGroupPersistenceService.java import org.opensearch.wlm.ResourceType; import java.util.Collection; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/service/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java similarity index 84% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/service/package-info.java rename to plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java index 3758c9fcd9b81..e8c88ee656dc7 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/service/package-info.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/package-info.java @@ -9,4 +9,4 @@ /** * Package for the service classes related to query groups in WorkloadManagementPlugin */ -package org.opensearch.plugin.wlm.querygroup.service; +package org.opensearch.plugin.wlm.service; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadGroupTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadGroupTestUtils.java index dac26c3c0f929..bac644a172c1e 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadGroupTestUtils.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadGroupTestUtils.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup; +package org.opensearch.plugin.wlm; import org.opensearch.cluster.ClusterName; import org.opensearch.cluster.ClusterState; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequestTests.java similarity index 96% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequestTests.java index 272f4db09109e..31d3ea00b7bda 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupRequestTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.common.io.stream.BytesStreamOutput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponseTests.java similarity index 98% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponseTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponseTests.java index 033e1a67cc012..d25050341f997 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/CreateWorkloadGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateWorkloadGroupResponseTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.common.io.stream.BytesStreamOutput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequestTests.java similarity index 96% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequestTests.java index 2df70252f9f55..a7fa0939583c5 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/DeleteWorkloadGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/DeleteWorkloadGroupRequestTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.ActionRequestValidationException; import org.opensearch.common.io.stream.BytesStreamOutput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequestTests.java similarity index 97% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequestTests.java index 52e1880cf9a4a..832761d5084bb 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupRequestTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponseTests.java similarity index 99% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponseTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponseTests.java index f9d77de883e73..dc0aeabc7a033 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/GetWorkloadGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/GetWorkloadGroupResponseTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.common.io.stream.BytesStreamOutput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupActionTests.java similarity index 84% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupActionTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupActionTests.java index 00aa957371e74..7ffa33aa8a80a 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupActionTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupActionTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; @@ -14,11 +14,7 @@ import org.opensearch.cluster.metadata.IndexNameExpressionResolver; import org.opensearch.cluster.service.ClusterService; import org.opensearch.core.action.ActionListener; -<<<<<<<< HEAD:plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteWorkloadGroupActionTests.java import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; -======== -import org.opensearch.plugin.wlm.querygroup.service.QueryGroupPersistenceService; ->>>>>>>> eaf11b7f929 (add get rule api logic):plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportDeleteQueryGroupActionTests.java import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupActionTests.java similarity index 97% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupActionTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupActionTests.java index b3aa0dceee78e..cf12d9f6408cf 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/TransportGetWorkloadGroupActionTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportGetWorkloadGroupActionTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.ResourceNotFoundException; import org.opensearch.action.support.ActionFilters; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupRequestTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupRequestTests.java similarity index 98% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupRequestTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupRequestTests.java index c982d5837d8eb..e8d883da5c6eb 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupRequestTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupRequestTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponseTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponseTests.java similarity index 98% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponseTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponseTests.java index ff27c74becc5b..97b9b9029373f 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/UpdateWorkloadGroupResponseTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/UpdateWorkloadGroupResponseTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.action; +package org.opensearch.plugin.wlm.action; import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.common.io.stream.BytesStreamOutput; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/WorkloadGroupActionTestUtils.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/WorkloadGroupActionTestUtils.java similarity index 100% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/action/WorkloadGroupActionTestUtils.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/WorkloadGroupActionTestUtils.java diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupActionTests.java similarity index 98% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupActionTests.java index df59ba3103db5..8ce5c869f4481 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteWorkloadGroupActionTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rest/RestDeleteWorkloadGroupActionTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.querygroup.rest; +package org.opensearch.plugin.wlm.rest; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; import org.opensearch.common.CheckedConsumer; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/service/WorkloadGroupPersistenceServiceTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceServiceTests.java similarity index 100% rename from plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/querygroup/service/WorkloadGroupPersistenceServiceTests.java rename to plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceServiceTests.java From 11041376910dbc51f1bef1983f99d3889682ac54 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Fri, 25 Apr 2025 01:04:50 -0700 Subject: [PATCH 13/20] rebase after get api is merged Signed-off-by: Ruirui Zhang --- .../rule/action/GetRuleResponse.java | 103 ------- .../opensearch/rule/action/package-info.java | 12 - .../opensearch/rule/rest/package-info.java | 12 - .../IndexStoredRulePersistenceService.java | 269 ------------------ .../rule/service/RulePersistenceService.java | 27 -- .../opensearch/rule/service/package-info.java | 12 - .../rule/utils/IndexStoredRuleParser.java | 51 ---- .../rule/utils/IndexStoredRuleUtils.java | 87 ------ .../opensearch/rule/utils/package-info.java | 12 - modules/autotagging-commons/build.gradle | 1 - .../opensearch/rule}/CreateRuleRequest.java | 13 +- .../opensearch/rule}/CreateRuleResponse.java | 6 +- .../opensearch/rule/RuleDuplicateChecker.java | 28 ++ .../rule/RulePersistenceService.java | 7 + .../rule/autotagging/FeatureType.java | 14 +- .../autotagging/FeatureValueValidator.java | 21 ++ .../IndexStoredRulePersistenceService.java | 147 +++++++++- .../IndexBasedRuleDuplicateChecker.java | 59 ++++ .../storage/IndexBasedRuleQueryMapper.java | 1 + .../org/opensearch/rule/RuleTestUtils.java | 157 ---------- .../rule/action/CreateRuleRequestTests.java | 5 +- .../rule/action/CreateRuleResponseTests.java | 7 +- .../rule/action/GetRuleRequestTests.java | 89 +++++- .../rule/action/GetRuleResponseTests.java | 51 ++-- .../rule/autotagging/RuleTests.java | 6 +- ...ndexStoredRulePersistenceServiceTests.java | 256 +++++++---------- .../utils/IndexStoredRuleParserTests.java | 52 ---- .../rule/utils/IndexStoredRuleUtilsTests.java | 69 ----- .../opensearch/rule/utils/RuleTestUtils.java | 97 +------ .../opensearch/rule/RuleFrameworkPlugin.java | 13 +- .../rule/action/CreateRuleAction.java | 18 +- .../action/TransportCreateRuleAction.java | 54 ++++ .../rule/rest/RestCreateRuleAction.java | 38 +-- .../InMemoryRuleProcessingServiceTests.java | 7 +- .../plugin/wlm/WorkloadManagementPlugin.java | 18 +- .../wlm/WorkloadManagementPluginModule.java | 11 +- .../plugin/wlm/rule/QueryGroupAttribute.java | 61 ---- .../wlm/rule/QueryGroupFeatureType.java | 64 ----- .../wlm/rule/WorkloadGroupFeatureType.java | 19 +- .../WorkloadGroupFeatureValueValidator.java | 40 +++ .../action/TransportCreateWlmRuleAction.java | 61 ---- .../plugin/wlm/rule/action/package-info.java | 12 - .../plugin/wlm/rule/rest/package-info.java | 12 - .../rule/WorkloadGroupFeatureTypeTests.java | 2 +- 44 files changed, 680 insertions(+), 1421 deletions(-) delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/action/package-info.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/package-info.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/service/package-info.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java delete mode 100644 libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/package-info.java rename {libs/autotagging-commons/src/main/java/org/opensearch/rule/action => modules/autotagging-commons/common/src/main/java/org/opensearch/rule}/CreateRuleRequest.java (78%) rename {libs/autotagging-commons/src/main/java/org/opensearch/rule/action => modules/autotagging-commons/common/src/main/java/org/opensearch/rule}/CreateRuleResponse.java (93%) create mode 100644 modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleDuplicateChecker.java create mode 100644 modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureValueValidator.java create mode 100644 modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleDuplicateChecker.java delete mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java delete mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java delete mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java rename plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/CreateWlmRuleAction.java => modules/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleAction.java (51%) create mode 100644 modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java rename {libs => modules}/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java (63%) delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/package-info.java delete mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/package-info.java diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java deleted file mode 100644 index 110522ff86af6..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleResponse.java +++ /dev/null @@ -1,103 +0,0 @@ -/* - * 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.rule.action; - -import org.opensearch.autotagging.Rule; -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.Map; - -import static org.opensearch.autotagging.Rule._ID_STRING; - -/** - * Response for the get API for Rule. - * Example response: - * { - * "rules": [ - * { - * "_id": "z1MJApUB0zgMcDmz-UQq", - * "description": "Rule for tagging query_group_id to index123" - * "index_pattern": ["index123"], - * "query_group": "query_group_id", - * "updated_at": "2025-02-14T01:19:22.589Z" - * }, - * ... - * ], - * "search_after": ["z1MJApUB0zgMcDmz-UQq"] - * } - * @opensearch.experimental - */ -public class GetRuleResponse extends ActionResponse implements ToXContent, ToXContentObject { - private final Map rules; - private final String searchAfter; - private final RestStatus restStatus; - - /** - * Constructor for GetRuleResponse - * @param rules - Rules get from the request - * @param searchAfter - The sort value used for pagination. - * @param restStatus - Status of the GetRuleResponse - */ - public GetRuleResponse(final Map rules, String searchAfter, RestStatus restStatus) { - this.rules = rules; - this.searchAfter = searchAfter; - this.restStatus = restStatus; - } - - /** - * Constructs a GetRuleResponse from a StreamInput for deserialization - * @param in - The {@link StreamInput} instance to read from. - */ - public GetRuleResponse(StreamInput in) throws IOException { - this(in.readMap(StreamInput::readString, Rule::new), in.readOptionalString(), RestStatus.readFrom(in)); - } - - @Override - public void writeTo(StreamOutput out) throws IOException { - out.writeMap(rules, StreamOutput::writeString, (outStream, rule) -> rule.writeTo(outStream)); - out.writeOptionalString(searchAfter); - RestStatus.writeTo(out, restStatus); - } - - @Override - public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { - builder.startObject(); - builder.startArray("rules"); - for (Map.Entry entry : rules.entrySet()) { - entry.getValue().toXContent(builder, new MapParams(Map.of(_ID_STRING, entry.getKey()))); - } - builder.endArray(); - if (searchAfter != null && !searchAfter.isEmpty()) { - builder.field("search_after", new Object[] { searchAfter }); - } - builder.endObject(); - return builder; - } - - /** - * rules getter - */ - public Map getRules() { - return rules; - } - - /** - * restStatus getter - */ - public RestStatus getRestStatus() { - return restStatus; - } -} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/package-info.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/package-info.java deleted file mode 100644 index 91913aff23eac..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/package-info.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * 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. - */ - -/** - * This package contains abstract action classes for rules - */ -package org.opensearch.rule.action; diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/package-info.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/package-info.java deleted file mode 100644 index c1000b90b1856..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/package-info.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * 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. - */ - -/** - * This package contains abstract rest classes for rules - */ -package org.opensearch.rule.rest; diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java deleted file mode 100644 index 6f8d67e51be6a..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ /dev/null @@ -1,269 +0,0 @@ -/* - * 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.rule.service; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.opensearch.ResourceAlreadyExistsException; -import org.opensearch.ResourceNotFoundException; -import org.opensearch.action.admin.indices.create.CreateIndexRequest; -import org.opensearch.action.admin.indices.create.CreateIndexResponse; -import org.opensearch.action.index.IndexRequest; -import org.opensearch.action.search.SearchRequestBuilder; -import org.opensearch.action.search.SearchResponse; -import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.common.xcontent.XContentFactory; -import org.opensearch.core.action.ActionListener; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.core.xcontent.ToXContent; -import org.opensearch.index.query.BoolQueryBuilder; -import org.opensearch.rule.action.CreateRuleRequest; -import org.opensearch.rule.action.CreateRuleResponse; -import org.opensearch.rule.action.GetRuleResponse; -import org.opensearch.rule.utils.IndexStoredRuleParser; -import org.opensearch.rule.utils.IndexStoredRuleUtils; -import org.opensearch.search.SearchHit; -import org.opensearch.search.sort.SortOrder; -import org.opensearch.transport.client.Client; - -import java.io.IOException; -import java.util.Arrays; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.Set; -import java.util.stream.Collectors; - -import static org.opensearch.autotagging.Rule._ID_STRING; - -/** - * This class encapsulates the logic to manage the lifecycle of rules at index level - * @opensearch.experimental - */ -public class IndexStoredRulePersistenceService implements RulePersistenceService { - /** - * The system index name used for storing rules - */ - private final String indexName; - private final ClusterService clusterService; - private final Client client; - private final FeatureType featureType; - private final int maxRulesPerPage; - private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class); - private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); - - /** - * Constructs an instance of {@link IndexStoredRulePersistenceService} with the specified parameters. - * This service handles persistence and retrieval of stored rules within an OpenSearch index. - * @param indexName - The name of the OpenSearch index where the rules are stored. - * @param clusterService - The clusterService used in IndexStoredRulePersistenceService. - * @param client - The OpenSearch client used to interact with the OpenSearch cluster. - * @param featureType - The feature type associated with the stored rules. - * @param maxRulesPerPage - The maximum number of rules that can be returned in a single get request. - */ - public IndexStoredRulePersistenceService( - String indexName, - ClusterService clusterService, - Client client, - FeatureType featureType, - int maxRulesPerPage - ) { - this.indexName = indexName; - this.clusterService = clusterService; - this.client = client; - this.featureType = featureType; - this.maxRulesPerPage = maxRulesPerPage; - } - - /** - * Entry point for the create rule API logic in persistence service. - * It ensures the index exists, validates for duplicate rules, and persists the new rule. - * @param request The CreateRuleRequest - * @param listener ActionListener for CreateRuleResponse - */ - public void createRule(CreateRuleRequest request, ActionListener listener) { - try (ThreadContext.StoredContext ctx = getContext()) { - createIndexIfAbsent(new ActionListener<>() { - @Override - public void onResponse(Boolean indexCreated) { - if (!indexCreated) { - listener.onFailure(new IllegalStateException(indexName + " index creation failed and rule cannot be persisted")); - return; - } - validateNoDuplicateRule(request.getRule(), new ActionListener<>() { - @Override - public void onResponse(Void unused) { - persistRule(request.getRule(), listener); - } - - @Override - public void onFailure(Exception e) { - listener.onFailure(e); - } - }); - } - - @Override - public void onFailure(Exception e) { - listener.onFailure(e); - } - }); - } - } - - /** - * Validates that no duplicate rule exists with the same attribute map. - * If a conflict is found, fails the listener - * @param rule - the rule we check duplicate against - * @param listener - listener for validateNoDuplicateRule response - */ - private void validateNoDuplicateRule(Rule rule, ActionListener listener) { - try (ThreadContext.StoredContext ctx = getContext()) { - getRuleFromIndex(null, rule.getAttributeMap(), null, new ActionListener<>() { - @Override - public void onResponse(GetRuleResponse getRuleResponse) { - Optional duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(rule, getRuleResponse.getRules()); - duplicateRuleId.ifPresentOrElse( - id -> listener.onFailure(new IllegalArgumentException("Rule already exists under rule id " + id)), - () -> listener.onResponse(null) - ); - } - - @Override - public void onFailure(Exception e) { - listener.onFailure(e); - } - }); - } - } - - /** - * Persist the rule in the index - * @param rule - The rule to update. - * @param listener - ActionListener for CreateRuleResponse - */ - private void persistRule(Rule rule, ActionListener listener) { - try (ThreadContext.StoredContext ctx = getContext()) { - IndexRequest indexRequest = new IndexRequest(indexName).source( - rule.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS) - ); - client.index(indexRequest, ActionListener.wrap(indexResponse -> { - listener.onResponse(new CreateRuleResponse(indexResponse.getId(), rule)); - }, e -> { - logger.warn("Failed to save Rule object due to error: {}", e.getMessage()); - listener.onFailure(e); - })); - } catch (IOException e) { - logger.error("Error saving rule to index: {}", indexName); - listener.onFailure(new RuntimeException("Failed to save rule to index.")); - } - } - - /** - * Creates the system index if it doesn't exist - * @param listener - ActionListener for CreateRuleResponse - */ - private void createIndexIfAbsent(ActionListener listener) { - try (ThreadContext.StoredContext ctx = getContext()) { - if (clusterService.state().metadata().hasIndex(indexName)) { - listener.onResponse(true); - return; - } - final CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName).settings(indexSettings); - client.admin().indices().create(createIndexRequest, new ActionListener<>() { - @Override - public void onResponse(CreateIndexResponse response) { - logger.info("Index {} created?: {}", indexName, response.isAcknowledged()); - listener.onResponse(response.isAcknowledged()); - } - - @Override - public void onFailure(Exception e) { - if (e instanceof ResourceAlreadyExistsException) { - logger.trace("Index {} already exists", indexName); - listener.onResponse(true); - } else { - logger.error("Failed to create index {}: {}", indexName, e.getMessage()); - listener.onFailure(e); - } - } - }); - } - } - - /** - * Entry point for the get rule api logic in persistence service. If id is provided, we only get a single rule. - * Otherwise, we get all rules that satisfy the attributeFilters. - * @param id - The id of the rule to get. - * @param attributeFilters - A map containing the attributes that user want to filter on - * @param searchAfter - The sort values from the last document of the previous page, used for pagination - * @param listener - ActionListener for GetRuleResponse - */ - private void getRuleFromIndex( - String id, - Map> attributeFilters, - String searchAfter, - ActionListener listener - ) { - // Stash the current thread context when interacting with system index to perform - // operations as the system itself, bypassing authorization checks. This ensures that - // actions within this block are trusted and executed with system-level privileges. - try (ThreadContext.StoredContext context = getContext()) { - BoolQueryBuilder boolQuery = IndexStoredRuleUtils.buildGetRuleQuery(id, attributeFilters, featureType); - SearchRequestBuilder searchRequest = client.prepareSearch(indexName).setQuery(boolQuery).setSize(maxRulesPerPage); - if (searchAfter != null) { - searchRequest.addSort(_ID_STRING, SortOrder.ASC).searchAfter(new Object[] { searchAfter }); - } - searchRequest.execute(ActionListener.wrap(searchResponse -> handleGetRuleResponse(id, searchResponse, listener), e -> { - logger.error("Failed to fetch all rules: {}", e.getMessage()); - listener.onFailure(e); - })); - } - } - - /** - * Process searchResponse from index and send a GetRuleResponse - * @param searchResponse - Response received from index - * @param listener - ActionListener for GetRuleResponse - */ - private void handleGetRuleResponse(String id, SearchResponse searchResponse, ActionListener listener) { - List hits = Arrays.asList(searchResponse.getHits().getHits()); - if (id != null && hits.isEmpty()) { - logger.error("Rule with ID " + id + " not found."); - listener.onFailure(new ResourceNotFoundException("Rule with ID " + id + " doesn't exist in the index.")); - return; - } - Map ruleMap = hits.stream() - .collect(Collectors.toMap(SearchHit::getId, hit -> IndexStoredRuleParser.parseRule(hit.getSourceAsString(), featureType))); - String nextSearchAfter = hits.isEmpty() ? null : hits.get(hits.size() - 1).getId(); - listener.onResponse(new GetRuleResponse(ruleMap, nextSearchAfter, RestStatus.OK)); - } - - private ThreadContext.StoredContext getContext() { - return client.threadPool().getThreadContext().stashContext(); - } - - /** - * client getter - */ - public Client getClient() { - return client; - } - - /** - * clusterService getter - */ - public ClusterService getClusterService() { - return clusterService; - } -} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java deleted file mode 100644 index 6b814a271db60..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/RulePersistenceService.java +++ /dev/null @@ -1,27 +0,0 @@ -/* - * 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.rule.service; - -import org.opensearch.core.action.ActionListener; -import org.opensearch.rule.action.CreateRuleRequest; -import org.opensearch.rule.action.CreateRuleResponse; - -/** - * Interface for a service that handles rule persistence CRUD operations. - * @opensearch.experimental - */ -public interface RulePersistenceService { - - /** - * Create rules based on the provided request. - * @param request The request containing the details for creating the rule. - * @param listener The listener that will handle the response or failure. - */ - void createRule(CreateRuleRequest request, ActionListener listener); -} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/package-info.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/package-info.java deleted file mode 100644 index f51b281884dc6..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/service/package-info.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * 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. - */ - -/** - * This package contains abstract service classes for rules - */ -package org.opensearch.rule.service; diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java deleted file mode 100644 index ccf48cceb645a..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleParser.java +++ /dev/null @@ -1,51 +0,0 @@ -/* - * 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.rule.utils; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule; -import org.opensearch.core.xcontent.DeprecationHandler; -import org.opensearch.core.xcontent.MediaTypeRegistry; -import org.opensearch.core.xcontent.NamedXContentRegistry; -import org.opensearch.core.xcontent.XContentParser; - -import java.io.IOException; - -/** - * Utility class for parsing index stored rules into Rule objects. - * @opensearch.experimental - */ -public class IndexStoredRuleParser { - - /** - * constructor for IndexStoredRuleParser - */ - private IndexStoredRuleParser() {} - - private static final Logger logger = LogManager.getLogger(IndexStoredRuleParser.class); - - /** - * Parses a source string into a Rule object - * @param source - The raw source string representing the rule to be parsed - * @param featureType - The feature type to associate with the parsed rule - */ - public static Rule parseRule(String source, FeatureType featureType) { - try ( - XContentParser parser = MediaTypeRegistry.JSON.xContent() - .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, source) - ) { - return Rule.Builder.fromXContent(parser, featureType).build(); - } catch (IOException e) { - logger.info("Issue met when parsing rule {}: {}", source, e.getMessage()); - throw new RuntimeException("Cannot parse rule from index: " + source); - } - } -} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java deleted file mode 100644 index 4b52f85f0cd24..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/IndexStoredRuleUtils.java +++ /dev/null @@ -1,87 +0,0 @@ -/* - * 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.rule.utils; - -import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule; -import org.opensearch.index.query.BoolQueryBuilder; -import org.opensearch.index.query.QueryBuilders; - -import java.util.Map; -import java.util.Optional; -import java.util.Set; - -import static org.opensearch.autotagging.Rule._ID_STRING; - -/** - * Utility class that provides methods for the lifecycle of rules. - * @opensearch.experimental - */ -public class IndexStoredRuleUtils { - - /** - * constructor for IndexStoredRuleUtils - */ - private IndexStoredRuleUtils() {} - - /** - * Builds a Boolean query to retrieve a rule by its ID or attribute filters. - * @param id The ID of the rule to search for. If null, no ID-based filtering is applied. - * @param attributeFilters A map of attributes and their corresponding filter values. This allows filtering by specific attribute values. - * @param featureType The feature type that is required in the query. - */ - public static BoolQueryBuilder buildGetRuleQuery(String id, Map> attributeFilters, FeatureType featureType) { - BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); - if (id != null) { - return boolQuery.must(QueryBuilders.termQuery(_ID_STRING, id)); - } - for (Map.Entry> entry : attributeFilters.entrySet()) { - Attribute attribute = entry.getKey(); - Set values = entry.getValue(); - if (values != null && !values.isEmpty()) { - BoolQueryBuilder attributeQuery = QueryBuilders.boolQuery(); - for (String value : values) { - attributeQuery.should(QueryBuilders.matchQuery(attribute.getName(), value)); - } - boolQuery.must(attributeQuery); - } - } - boolQuery.filter(QueryBuilders.existsQuery(featureType.getName())); - return boolQuery; - } - - /** - * Checks if a duplicate rule exists based on the attribute map. - * A rule is considered a duplicate when the attribute value already exists in the index, and the number of - * attributes in the new rule is equal to the number of attributes in an existing rule. - * - * For example, if an existing rule has: - * attribute1 = ['a'] and attribute2 = ['c'] - * And we are creating a new rule with: - * attribute1 = ['a'] - * Then it's not a duplicate because the existing rule has attribute2 and is more granular - * - * @param rule The rule to be validated against ruleMap. - * @param ruleMap This map entries are Rules that contain the attribute values from rule, meaning they - * have a partial or complete overlap with the new rule being created. - */ - public static Optional getDuplicateRuleId(Rule rule, Map ruleMap) { - Map> attributeMapToValidate = rule.getAttributeMap(); - for (Map.Entry entry : ruleMap.entrySet()) { - String ruleId = entry.getKey(); - Rule currRule = entry.getValue(); - // Compare the size of the attribute maps to ensure we only check for duplicates with the same number of attributes. - if (attributeMapToValidate.size() == currRule.getAttributeMap().size()) { - return Optional.of(ruleId); - } - } - return Optional.empty(); - } -} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/package-info.java b/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/package-info.java deleted file mode 100644 index 5b2961fb3fb2e..0000000000000 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/utils/package-info.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * 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 for the service classes related to query groups in WorkloadManagementPlugin - */ -package org.opensearch.rule.utils; diff --git a/modules/autotagging-commons/build.gradle b/modules/autotagging-commons/build.gradle index 5c996fc85bea9..8b9c4fbb7d409 100644 --- a/modules/autotagging-commons/build.gradle +++ b/modules/autotagging-commons/build.gradle @@ -6,7 +6,6 @@ * compatible open source license. */ - opensearchplugin { name = "rule-framework" description = 'OpenSearch Rule Framework plugin' diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java similarity index 78% rename from libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java rename to modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java index 3eb6e4e9f086d..fea2a9eda6e27 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleRequest.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java @@ -6,13 +6,13 @@ * compatible open source license. */ -package org.opensearch.rule.action; +package org.opensearch.rule; import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionRequestValidationException; -import org.opensearch.autotagging.Rule; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.rule.autotagging.Rule; import java.io.IOException; @@ -50,7 +50,14 @@ public CreateRuleRequest(StreamInput in) throws IOException { @Override public ActionRequestValidationException validate() { - return null; + try { + rule.getFeatureType().getFeatureValueValidator().validate(rule.getFeatureValue()); + return null; + } catch (Exception e) { + ActionRequestValidationException validationException = new ActionRequestValidationException(); + validationException.addValidationError(e.getMessage()); + return validationException; + } } @Override diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleResponse.java similarity index 93% rename from libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java rename to modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleResponse.java index ba6c83dc2d8c9..9c95c30ca13c5 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleResponse.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleResponse.java @@ -6,20 +6,20 @@ * compatible open source license. */ -package org.opensearch.rule.action; +package org.opensearch.rule; -import org.opensearch.autotagging.Rule; 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.xcontent.ToXContent; import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.rule.autotagging.Rule; import java.io.IOException; import java.util.Map; -import static org.opensearch.autotagging.Rule._ID_STRING; +import static org.opensearch.rule.autotagging.Rule._ID_STRING; /** * Response for the create API for Rule diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleDuplicateChecker.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleDuplicateChecker.java new file mode 100644 index 0000000000000..0e5eb1742de0f --- /dev/null +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleDuplicateChecker.java @@ -0,0 +1,28 @@ +/* + * 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.rule; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.rule.autotagging.Rule; + +import java.util.Map; +import java.util.Optional; + +/** + * Interface to check for rule duplication. + */ +@ExperimentalApi +public interface RuleDuplicateChecker { + /** + * Checks if the given rule already exists in the provided rule map. + * @param rule the rule to check for duplication + * @param ruleMap a map of existing rules, keyed by rule ID + */ + Optional getDuplicateRuleId(Rule rule, Map ruleMap); +} diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RulePersistenceService.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RulePersistenceService.java index 674e77a990849..b29323da421e7 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RulePersistenceService.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RulePersistenceService.java @@ -17,6 +17,13 @@ */ public interface RulePersistenceService { + /** + * Create rules based on the provided request. + * @param request The request containing the details for creating the rule. + * @param listener The listener that will handle the response or failure. + */ + void createRule(CreateRuleRequest request, ActionListener listener); + /** * Get rules based on the provided request. * @param request The request containing the details for retrieving the rule. diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureType.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureType.java index c752f917264de..f56d6d3da5589 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureType.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureType.java @@ -17,13 +17,11 @@ /** * Represents a feature type within the auto-tagging feature. Feature types define different categories of - * characteristics that can be used for tagging and classification. Implementations of this interface are - * responsible for registering feature types in {@link AutoTaggingRegistry}. Implementations must ensure that + * characteristics that can be used for tagging and classification. Implementations must ensure that * feature types are uniquely identifiable by their class and name. * * Implementers should follow these guidelines: * Feature types should be singletons and managed centrally to avoid duplicates. - * {@link #registerFeatureType()} must be called during initialization to ensure the feature type is available. * * @opensearch.experimental */ @@ -49,6 +47,11 @@ public interface FeatureType extends Writeable { */ Map getAllowedAttributesRegistry(); + /** + * returns the validator for feature value + */ + FeatureValueValidator getFeatureValueValidator(); + /** * returns max attribute values * @return @@ -65,11 +68,6 @@ default int getMaxCharLengthPerAttributeValue() { return DEFAULT_MAX_ATTRIBUTE_VALUE_LENGTH; } - /** - * makes the feature type usable and available to framework plugin - */ - void registerFeatureType(); - /** * checks the validity of the input attribute * @param attribute diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureValueValidator.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureValueValidator.java new file mode 100644 index 0000000000000..40b97e9ddd843 --- /dev/null +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureValueValidator.java @@ -0,0 +1,21 @@ +/* + * 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.rule.autotagging; + +/** + * Validates a feature value for a specific feature type + * @opensearch.experimental + */ +public interface FeatureValueValidator { + /** + * Validates the given feature value. + * @param featureValue the value to validate + */ + void validate(String featureValue); +} diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java index f42ea1e951e4f..3bdf8d9338411 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -10,19 +10,31 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.ResourceAlreadyExistsException; import org.opensearch.ResourceNotFoundException; import org.opensearch.action.DocWriteResponse; import org.opensearch.action.delete.DeleteRequest; import org.opensearch.action.search.SearchRequestBuilder; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; +import org.opensearch.action.admin.indices.create.CreateIndexRequest; +import org.opensearch.action.admin.indices.create.CreateIndexResponse; +import org.opensearch.action.index.IndexRequest; +import org.opensearch.action.search.SearchRequestBuilder; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.action.ActionListener; import org.opensearch.index.engine.DocumentMissingException; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.rule.DeleteRuleRequest; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.rule.CreateRuleRequest; +import org.opensearch.rule.CreateRuleResponse; import org.opensearch.rule.GetRuleRequest; import org.opensearch.rule.GetRuleResponse; +import org.opensearch.rule.RuleDuplicateChecker; import org.opensearch.rule.RuleEntityParser; import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.RuleQueryMapper; @@ -31,9 +43,11 @@ import org.opensearch.search.sort.SortOrder; import org.opensearch.transport.client.Client; +import java.io.IOException; import java.util.Arrays; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.stream.Collectors; import static org.opensearch.rule.autotagging.Rule._ID_STRING; @@ -48,32 +62,160 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService */ private final String indexName; private final Client client; + private final ClusterService clusterService; private final int maxRulesPerPage; private final RuleEntityParser parser; private final RuleQueryMapper queryBuilder; + private final RuleDuplicateChecker ruleDuplicateChecker; private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class); + private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); /** * Constructs an instance of {@link IndexStoredRulePersistenceService} with the specified parameters. * This service handles persistence and retrieval of stored rules within an OpenSearch index. * @param indexName - The name of the OpenSearch index where the rules are stored. * @param client - The OpenSearch client used to interact with the OpenSearch cluster. + * @param clusterService * @param maxRulesPerPage - The maximum number of rules that can be returned in a single get request. * @param parser * @param queryBuilder + * @param ruleDuplicateChecker */ public IndexStoredRulePersistenceService( String indexName, Client client, + ClusterService clusterService, int maxRulesPerPage, RuleEntityParser parser, - RuleQueryMapper queryBuilder + RuleQueryMapper queryBuilder, + RuleDuplicateChecker ruleDuplicateChecker ) { this.indexName = indexName; this.client = client; + this.clusterService = clusterService; this.maxRulesPerPage = maxRulesPerPage; this.parser = parser; this.queryBuilder = queryBuilder; + this.ruleDuplicateChecker = ruleDuplicateChecker; + } + + /** + * Entry point for the create rule API logic in persistence service. + * It ensures the index exists, validates for duplicate rules, and persists the new rule. + * @param request The CreateRuleRequest + * @param listener ActionListener for CreateRuleResponse + */ + public void createRule(CreateRuleRequest request, ActionListener listener) { + try (ThreadContext.StoredContext ctx = getContext()) { + createIndexIfAbsent(new ActionListener<>() { + @Override + public void onResponse(Void unused) { + validateNoDuplicateRule(request.getRule(), new ActionListener<>() { + @Override + public void onResponse(Void unused) { + persistRule(request.getRule(), listener); + } + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + }); + } + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + }); + } + } + + /** + * Creates the system index if it doesn't exist + * @param listener - ActionListener for CreateRuleResponse + */ + private void createIndexIfAbsent(ActionListener listener) { + try (ThreadContext.StoredContext ctx = getContext()) { + if (clusterService.state().metadata().hasIndex(indexName)) { + listener.onResponse(null); + return; + } + + final CreateIndexRequest request = new CreateIndexRequest(indexName).settings(indexSettings); + client.admin().indices().create(request, new ActionListener<>() { + @Override + public void onResponse(CreateIndexResponse response) { + if (!response.isAcknowledged()) { + logger.error("Index creation not acknowledged: {}", indexName); + listener.onFailure(new IllegalStateException(indexName + " index creation failed and rule cannot be persisted")); + } else { + logger.info("Index {} created", indexName); + listener.onResponse(null); + } + } + + @Override + public void onFailure(Exception e) { + if (e instanceof ResourceAlreadyExistsException) { + logger.trace("Index {} already exists", indexName); + listener.onResponse(null); + } else { + logger.error("Failed to create index {}: {}", indexName, e.getMessage()); + listener.onFailure(e); + } + } + }); + } + } + + /** + * Validates that no duplicate rule exists with the same attribute map. + * If a conflict is found, fails the listener + * @param rule - the rule we check duplicate against + * @param listener - listener for validateNoDuplicateRule response + */ + private void validateNoDuplicateRule(Rule rule, ActionListener listener) { + try (ThreadContext.StoredContext ctx = getContext()) { + QueryBuilder query = queryBuilder.from(new GetRuleRequest(null, rule.getAttributeMap(), null, rule.getFeatureType())); + getRuleFromIndex(null, query, null, new ActionListener<>() { + @Override + public void onResponse(GetRuleResponse getRuleResponse) { + Optional duplicateRuleId = ruleDuplicateChecker.getDuplicateRuleId(rule, getRuleResponse.getRules()); + duplicateRuleId.ifPresentOrElse( + id -> listener.onFailure(new IllegalArgumentException("Duplicate rule exists under id " + id)), + () -> listener.onResponse(null) + ); + } + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + }); + } + } + + /** + * Persist the rule in the index + * @param rule - The rule to update. + * @param listener - ActionListener for CreateRuleResponse + */ + private void persistRule(Rule rule, ActionListener listener) { + try (ThreadContext.StoredContext ctx = getContext()) { + IndexRequest indexRequest = new IndexRequest(indexName).source( + rule.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS) + ); + client.index(indexRequest, ActionListener.wrap(indexResponse -> { + listener.onResponse(new CreateRuleResponse(indexResponse.getId(), rule)); + }, e -> { + logger.warn("Failed to save Rule object due to error: {}", e.getMessage()); + listener.onFailure(e); + })); + } catch (IOException e) { + logger.error("Error saving rule to index: {}", indexName); + listener.onFailure(new RuntimeException("Failed to save rule to index.")); + } } /** @@ -82,8 +224,7 @@ public IndexStoredRulePersistenceService( * @param listener the listener for GetRuleResponse. */ public void getRule(GetRuleRequest getRuleRequest, ActionListener listener) { - final QueryBuilder getQueryBuilder = queryBuilder.from(getRuleRequest) - .filter(QueryBuilders.existsQuery(getRuleRequest.getFeatureType().getName())); + final QueryBuilder getQueryBuilder = queryBuilder.from(getRuleRequest); getRuleFromIndex(getRuleRequest.getId(), getQueryBuilder, getRuleRequest.getSearchAfter(), listener); } diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleDuplicateChecker.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleDuplicateChecker.java new file mode 100644 index 0000000000000..dd0800884b8a4 --- /dev/null +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleDuplicateChecker.java @@ -0,0 +1,59 @@ +/* + * 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.rule.storage; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.rule.RuleDuplicateChecker; +import org.opensearch.rule.autotagging.Attribute; +import org.opensearch.rule.autotagging.Rule; + +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** + * This class is used to check rule duplication for indexed based rules. + */ +@ExperimentalApi +public class IndexBasedRuleDuplicateChecker implements RuleDuplicateChecker { + + /** + * Default constructor + */ + public IndexBasedRuleDuplicateChecker() {} + + /** + * Checks if a duplicate rule exists based on the attribute map. + * A rule is considered a duplicate when the attribute value already exists in the index, and the number of + * attributes in the new rule is equal to the number of attributes in an existing rule. + * + * For example, if an existing rule has: + * attribute1 = ['a'] and attribute2 = ['c'] + * And we are creating a new rule with: + * attribute1 = ['a'] + * Then it's not a duplicate because the existing rule has attribute2 and is more granular + * + * @param rule The rule to be validated against ruleMap. + * @param ruleMap This map entries are Rules that contain the attribute values from rule, meaning they + * have a partial or complete overlap with the new rule being created. + */ + @Override + public Optional getDuplicateRuleId(Rule rule, Map ruleMap) { + Map> attributeMapToValidate = rule.getAttributeMap(); + for (Map.Entry entry : ruleMap.entrySet()) { + String ruleId = entry.getKey(); + Rule currRule = entry.getValue(); + // Compare the size of the attribute maps to ensure we only check for duplicates with the same number of attributes. + if (attributeMapToValidate.size() == currRule.getAttributeMap().size()) { + return Optional.of(ruleId); + } + } + return Optional.empty(); + } +} diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleQueryMapper.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleQueryMapper.java index 6f57ef478b577..96b77e0b4b643 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleQueryMapper.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleQueryMapper.java @@ -38,6 +38,7 @@ public QueryBuilder from(GetRuleRequest request) { final Map> attributeFilters = request.getAttributeFilters(); final String id = request.getId(); + boolQuery.filter(QueryBuilders.existsQuery(request.getFeatureType().getName())); if (id != null) { return boolQuery.must(QueryBuilders.termQuery(_ID_STRING, id)); } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java deleted file mode 100644 index fd9a228370af4..0000000000000 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleTestUtils.java +++ /dev/null @@ -1,157 +0,0 @@ -/* - * 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.rule.utils; - -import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.AutoTaggingRegistry; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule; -import org.opensearch.cluster.ClusterState; -import org.opensearch.cluster.metadata.Metadata; -import org.opensearch.cluster.metadata.QueryGroup; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.settings.Settings; -import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.rule.service.IndexStoredRulePersistenceService; -import org.opensearch.threadpool.ThreadPool; -import org.opensearch.transport.client.Client; - -import java.util.Map; -import java.util.Set; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertTrue; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public class RuleTestUtils { - public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; - public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; - public static final String FEATURE_VALUE_ONE = "feature_value_one"; - public static final String FEATURE_VALUE_TWO = "feature_value_two"; - public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; - public static final String ATTRIBUTE_VALUE_TWO = "mock_attribute_two"; - public static final String DESCRIPTION_ONE = "description_1"; - public static final String DESCRIPTION_TWO = "description_2"; - public static final String SEARCH_AFTER = "search_after_id"; - public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; - public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; - public static final String FEATURE_TYPE_NAME = "mock_feature_type"; - public static final String TEST_INDEX_NAME = ".test_index_for_rule"; - public static final Map> ATTRIBUTE_MAP = Map.of( - MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, - Set.of(ATTRIBUTE_VALUE_ONE) - ); - public static final Rule ruleOne = Rule.builder() - .description(DESCRIPTION_ONE) - .featureType(MockRuleFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_ONE) - .attributeMap(ATTRIBUTE_MAP) - .updatedAt(TIMESTAMP_ONE) - .build(); - - public static final Rule ruleTwo = Rule.builder() - .description(DESCRIPTION_TWO) - .featureType(MockRuleFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_TWO) - .attributeMap(Map.of(MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(ATTRIBUTE_VALUE_TWO))) - .updatedAt(TIMESTAMP_TWO) - .build(); - - public static Map ruleMap() { - return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); - } - - public static IndexStoredRulePersistenceService setUpIndexStoredRulePersistenceService(Map queryGroupMap) { - Client client = mock(Client.class); - ClusterService clusterService = mock(ClusterService.class); - ClusterState clusterState = mock(ClusterState.class); - Metadata metadata = mock(Metadata.class); - ThreadPool threadPool = mock(ThreadPool.class); - - ThreadContext threadContext = new ThreadContext(Settings.EMPTY); - when(client.threadPool()).thenReturn(threadPool); - when(threadPool.getThreadContext()).thenReturn(threadContext); - when(clusterService.state()).thenReturn(clusterState); - when(clusterState.metadata()).thenReturn(metadata); - when(metadata.queryGroups()).thenReturn(queryGroupMap); - return new IndexStoredRulePersistenceService(TEST_INDEX_NAME, clusterService, client, MockRuleFeatureType.INSTANCE, 50); - } - - public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { - assertEquals(mapOne.size(), mapTwo.size()); - for (Map.Entry entry : mapOne.entrySet()) { - String id = entry.getKey(); - assertTrue(mapTwo.containsKey(id)); - Rule one = mapOne.get(id); - Rule two = mapTwo.get(id); - assertEqualRule(one, two, ruleUpdated); - } - } - - public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { - if (ruleUpdated) { - assertEquals(one.getDescription(), two.getDescription()); - assertEquals(one.getFeatureType(), two.getFeatureType()); - assertEquals(one.getFeatureValue(), two.getFeatureValue()); - assertEquals(one.getAttributeMap(), two.getAttributeMap()); - assertEquals(one.getAttributeMap(), two.getAttributeMap()); - } else { - assertEquals(one, two); - } - } - - public static class MockRuleFeatureType implements FeatureType { - - public static final MockRuleFeatureType INSTANCE = new MockRuleFeatureType(); - - private MockRuleFeatureType() {} - - static { - INSTANCE.registerFeatureType(); - } - - @Override - public String getName() { - return FEATURE_TYPE_NAME; - } - - @Override - public Map getAllowedAttributesRegistry() { - return Map.of( - ATTRIBUTE_VALUE_ONE, - MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, - ATTRIBUTE_VALUE_TWO, - MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO - ); - } - - @Override - public void registerFeatureType() { - AutoTaggingRegistry.registerFeatureType(INSTANCE); - } - } - - public enum MockRuleAttributes implements Attribute { - MOCK_RULE_ATTRIBUTE_ONE(ATTRIBUTE_VALUE_ONE), - MOCK_RULE_ATTRIBUTE_TWO(ATTRIBUTE_VALUE_TWO); - ; - - private final String name; - - MockRuleAttributes(String name) { - this.name = name; - } - - @Override - public String getName() { - return name; - } - } -} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java index 4ebdf296cf1c0..4f6626df9af3c 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java @@ -10,12 +10,13 @@ import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.rule.CreateRuleRequest; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; -import static org.opensearch.rule.utils.RuleTestUtils.assertEqualRule; -import static org.opensearch.rule.utils.RuleTestUtils.ruleOne; +import static org.opensearch.rule.action.GetRuleRequestTests.assertEqualRule; +import static org.opensearch.rule.action.GetRuleRequestTests.ruleOne; public class CreateRuleRequestTests extends OpenSearchTestCase { diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java index 17de247c09564..e4bdf5fd94c53 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java @@ -8,20 +8,21 @@ package org.opensearch.rule.action; -import org.opensearch.autotagging.Rule; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.rule.CreateRuleResponse; +import org.opensearch.rule.autotagging.Rule; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; import java.util.Map; +import static org.opensearch.rule.action.GetRuleRequestTests.assertEqualRules; +import static org.opensearch.rule.action.GetRuleResponseTests.ruleOne; import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; -import static org.opensearch.rule.utils.RuleTestUtils.assertEqualRules; -import static org.opensearch.rule.utils.RuleTestUtils.ruleOne; import static org.mockito.Mockito.mock; public class CreateRuleResponseTests extends OpenSearchTestCase { diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java index 61f452e26f624..a451a58356606 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java @@ -10,25 +10,27 @@ import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.rule.RuleTestUtils; +import org.opensearch.rule.GetRuleRequest; +import org.opensearch.rule.autotagging.Attribute; +import org.opensearch.rule.autotagging.Rule; +import org.opensearch.rule.utils.RuleTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; import java.util.HashMap; - -import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_MAP; -import static org.opensearch.rule.RuleTestUtils.MockRuleFeatureType; -import static org.opensearch.rule.RuleTestUtils.SEARCH_AFTER; -import static org.opensearch.rule.RuleTestUtils._ID_ONE; +import java.util.Map; +import java.util.Set; public class GetRuleRequestTests extends OpenSearchTestCase { - /** * Test case to verify the serialization and deserialization of GetRuleRequest */ public void testSerialization() throws IOException { GetRuleRequest request = new GetRuleRequest(_ID_ONE, ATTRIBUTE_MAP, null, RuleTestUtils.MockRuleFeatureType.INSTANCE); assertEquals(_ID_ONE, request.getId()); + assertNull(request.validate()); + assertNull(request.getSearchAfter()); + assertEquals(RuleTestUtils.MockRuleFeatureType.INSTANCE, request.getFeatureType()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); StreamInput streamInput = out.bytes().streamInput(); @@ -41,7 +43,12 @@ public void testSerialization() throws IOException { * Test case to verify the serialization and deserialization of GetRuleRequest when name is null */ public void testSerializationWithNull() throws IOException { - GetRuleRequest request = new GetRuleRequest((String) null, new HashMap<>(), SEARCH_AFTER, MockRuleFeatureType.INSTANCE); + GetRuleRequest request = new GetRuleRequest( + (String) null, + new HashMap<>(), + SEARCH_AFTER, + RuleTestUtils.MockRuleFeatureType.INSTANCE + ); assertNull(request.getId()); BytesStreamOutput out = new BytesStreamOutput(); request.writeTo(out); @@ -50,4 +57,70 @@ public void testSerializationWithNull() throws IOException { assertEquals(request.getId(), otherRequest.getId()); assertEquals(request.getAttributeFilters(), otherRequest.getAttributeFilters()); } + + public void testValidate() { + GetRuleRequest request = new GetRuleRequest("", ATTRIBUTE_MAP, null, RuleTestUtils.MockRuleFeatureType.INSTANCE); + assertThrows(IllegalArgumentException.class, request::validate); + request = new GetRuleRequest(_ID_ONE, ATTRIBUTE_MAP, "", RuleTestUtils.MockRuleFeatureType.INSTANCE); + assertThrows(IllegalArgumentException.class, request::validate); + } + + public static final String _ID_ONE = "id_1"; + public static final String SEARCH_AFTER = "search_after"; + public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; + public static final String FEATURE_VALUE_ONE = "feature_value_one"; + public static final String FEATURE_VALUE_TWO = "feature_value_two"; + public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; + public static final String ATTRIBUTE_VALUE_TWO = "mock_attribute_two"; + public static final String DESCRIPTION_ONE = "description_1"; + public static final String DESCRIPTION_TWO = "description_2"; + public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; + public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; + public static final Map> ATTRIBUTE_MAP = Map.of( + RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, + Set.of(ATTRIBUTE_VALUE_ONE) + ); + + public static final Rule ruleOne = Rule.builder() + .description(DESCRIPTION_ONE) + .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_ONE) + .attributeMap(ATTRIBUTE_MAP) + .updatedAt(TIMESTAMP_ONE) + .build(); + + public static final Rule ruleTwo = Rule.builder() + .description(DESCRIPTION_TWO) + .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_TWO) + .attributeMap(Map.of(RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(ATTRIBUTE_VALUE_TWO))) + .updatedAt(TIMESTAMP_TWO) + .build(); + + public static Map ruleMap() { + return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); + } + + public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { + assertEquals(mapOne.size(), mapTwo.size()); + for (Map.Entry entry : mapOne.entrySet()) { + String id = entry.getKey(); + assertTrue(mapTwo.containsKey(id)); + Rule one = mapOne.get(id); + Rule two = mapTwo.get(id); + assertEqualRule(one, two, ruleUpdated); + } + } + + public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { + if (ruleUpdated) { + assertEquals(one.getDescription(), two.getDescription()); + assertEquals(one.getFeatureType(), two.getFeatureType()); + assertEquals(one.getFeatureValue(), two.getFeatureValue()); + assertEquals(one.getAttributeMap(), two.getAttributeMap()); + assertEquals(one.getAttributeMap(), two.getAttributeMap()); + } else { + assertEquals(one, two); + } + } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java index 5619bf4ee8567..f01bb94fab276 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java @@ -8,27 +8,45 @@ package org.opensearch.rule.action; -import org.opensearch.autotagging.Rule; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.rule.GetRuleResponse; +import org.opensearch.rule.autotagging.Attribute; +import org.opensearch.rule.autotagging.Rule; +import org.opensearch.rule.utils.RuleTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; import java.util.HashMap; import java.util.Map; +import java.util.Set; -import static org.opensearch.rule.RuleTestUtils.SEARCH_AFTER; -import static org.opensearch.rule.RuleTestUtils._ID_ONE; -import static org.opensearch.rule.RuleTestUtils.assertEqualRules; -import static org.opensearch.rule.RuleTestUtils.ruleMap; -import static org.opensearch.rule.RuleTestUtils.ruleOne; +import static org.opensearch.rule.action.GetRuleRequestTests.SEARCH_AFTER; +import static org.opensearch.rule.action.GetRuleRequestTests._ID_ONE; +import static org.opensearch.rule.action.GetRuleRequestTests.assertEqualRules; +import static org.opensearch.rule.action.GetRuleRequestTests.ruleMap; import static org.mockito.Mockito.mock; public class GetRuleResponseTests extends OpenSearchTestCase { + public static final String FEATURE_VALUE_ONE = "feature_value_one"; + public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; + public static final String DESCRIPTION_ONE = "description_1"; + public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; + static final Map> ATTRIBUTE_MAP = Map.of( + RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, + Set.of(ATTRIBUTE_VALUE_ONE) + ); + + public static final Rule ruleOne = Rule.builder() + .description(DESCRIPTION_ONE) + .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_ONE) + .attributeMap(ATTRIBUTE_MAP) + .updatedAt(TIMESTAMP_ONE) + .build(); /** * Test case to verify the serialization and deserialization of GetRuleResponse @@ -36,7 +54,7 @@ public class GetRuleResponseTests extends OpenSearchTestCase { public void testSerializationSingleRule() throws IOException { Map map = new HashMap<>(); map.put(_ID_ONE, ruleOne); - GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), null, RestStatus.OK); + GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), null); assertEquals(response.getRules(), map); BytesStreamOutput out = new BytesStreamOutput(); @@ -44,7 +62,6 @@ public void testSerializationSingleRule() throws IOException { StreamInput streamInput = out.bytes().streamInput(); GetRuleResponse otherResponse = new GetRuleResponse(streamInput); - assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); assertEqualRules(response.getRules(), otherResponse.getRules(), false); } @@ -52,7 +69,7 @@ public void testSerializationSingleRule() throws IOException { * Test case to verify the serialization and deserialization of GetRuleResponse when the result contains multiple rules */ public void testSerializationMultipleRule() throws IOException { - GetRuleResponse response = new GetRuleResponse(ruleMap(), SEARCH_AFTER, RestStatus.OK); + GetRuleResponse response = new GetRuleResponse(ruleMap(), SEARCH_AFTER); assertEquals(response.getRules(), ruleMap()); BytesStreamOutput out = new BytesStreamOutput(); @@ -60,7 +77,6 @@ public void testSerializationMultipleRule() throws IOException { StreamInput streamInput = out.bytes().streamInput(); GetRuleResponse otherResponse = new GetRuleResponse(streamInput); - assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); assertEquals(2, otherResponse.getRules().size()); assertEqualRules(response.getRules(), otherResponse.getRules(), false); } @@ -70,7 +86,7 @@ public void testSerializationMultipleRule() throws IOException { */ public void testSerializationNull() throws IOException { Map map = new HashMap<>(); - GetRuleResponse response = new GetRuleResponse(map, SEARCH_AFTER, RestStatus.OK); + GetRuleResponse response = new GetRuleResponse(map, SEARCH_AFTER); assertEquals(response.getRules(), map); BytesStreamOutput out = new BytesStreamOutput(); @@ -78,7 +94,6 @@ public void testSerializationNull() throws IOException { StreamInput streamInput = out.bytes().streamInput(); GetRuleResponse otherResponse = new GetRuleResponse(streamInput); - assertEquals(response.getRestStatus(), otherResponse.getRestStatus()); assertEquals(0, otherResponse.getRules().size()); } @@ -88,23 +103,23 @@ public void testSerializationNull() throws IOException { public void testToXContentGetSingleRule() throws IOException { Map map = new HashMap<>(); map.put(_ID_ONE, ruleOne); - GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), SEARCH_AFTER, RestStatus.OK); + GetRuleResponse response = new GetRuleResponse(Map.of(_ID_ONE, ruleOne), SEARCH_AFTER); XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); String actual = response.toXContent(builder, mock(ToXContent.Params.class)).toString(); String expected = "{\n" + " \"rules\" : [\n" + " {\n" - + " \"_id\" : \"AgfUO5Ja9yfvhdONlYi3TQ==\",\n" + + " \"_id\" : \"id_1\",\n" + " \"description\" : \"description_1\",\n" + " \"mock_attribute_one\" : [\n" - + " \"pattern_1\"\n" + + " \"mock_attribute_one\"\n" + " ],\n" + " \"mock_feature_type\" : \"feature_value_one\",\n" + " \"updated_at\" : \"2024-01-26T08:58:57.558Z\"\n" + " }\n" + " ],\n" + " \"search_after\" : [\n" - + " \"search_after_id\"\n" + + " \"search_after\"\n" + " ]\n" + "}"; assertEquals(expected, actual); @@ -115,7 +130,7 @@ public void testToXContentGetSingleRule() throws IOException { */ public void testToXContentGetZeroRule() throws IOException { XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint(); - GetRuleResponse otherResponse = new GetRuleResponse(new HashMap<>(), null, RestStatus.OK); + GetRuleResponse otherResponse = new GetRuleResponse(new HashMap<>(), null); String actual = otherResponse.toXContent(builder, mock(ToXContent.Params.class)).toString(); String expected = "{\n" + " \"rules\" : [ ]\n" + "}"; assertEquals(expected, actual); diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java index 4ae89b496a1b3..52687990f235b 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java @@ -92,7 +92,7 @@ public static class TestFeatureType implements FeatureType { public TestFeatureType() {} static { - INSTANCE.registerFeatureType(); + AutoTaggingRegistry.registerFeatureType(INSTANCE); } @Override @@ -116,8 +116,8 @@ public Map getAllowedAttributesRegistry() { } @Override - public void registerFeatureType() { - AutoTaggingRegistry.registerFeatureType(INSTANCE); + public FeatureValueValidator getFeatureValueValidator() { + return null; } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java index a813fe3ce7880..09a693dfd2b4e 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java @@ -32,114 +32,143 @@ import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.RuleQueryMapper; import org.opensearch.rule.autotagging.Rule; -import org.opensearch.rule.RuleTestUtils; -import org.opensearch.action.admin.indices.create.CreateIndexRequest; -import org.opensearch.action.admin.indices.create.CreateIndexResponse; -import org.opensearch.action.index.IndexRequest; -import org.opensearch.action.index.IndexResponse; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.core.action.ActionListener; -import org.opensearch.core.index.shard.ShardId; -import org.opensearch.rule.action.CreateRuleResponse; +import org.opensearch.rule.utils.RuleTestUtils; +import org.opensearch.search.SearchHit; +import org.opensearch.search.SearchHits; import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.transport.client.AdminClient; +import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.Client; -import org.opensearch.transport.client.IndicesAdminClient; -import java.io.IOException; import java.util.HashMap; import org.mockito.ArgumentCaptor; import static org.opensearch.rule.XContentRuleParserTests.VALID_JSON; -import static org.opensearch.rule.RuleTestUtils.TEST_INDEX_NAME; -import static org.opensearch.rule.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.TEST_INDEX_NAME; +import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.argThat; -import static org.opensearch.rule.RuleTestUtils.ATTRIBUTE_MAP; -import static org.opensearch.rule.RuleTestUtils.TEST_INDEX_NAME; -import static org.opensearch.rule.RuleTestUtils.ruleOne; -import static org.opensearch.rule.RuleTestUtils.setUpIndexStoredRulePersistenceService; import static org.mockito.Mockito.any; +import static org.mockito.Mockito.anyInt; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @SuppressWarnings("unchecked") public class IndexStoredRulePersistenceServiceTests extends OpenSearchTestCase { - public void testCreateIndexIfAbsent_IndexExists() { - IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - ClusterService clusterService = rulePersistenceService.getClusterService(); - when(clusterService.state().metadata().hasIndex(TEST_INDEX_NAME)).thenReturn(true); - ActionListener listener = mock(ActionListener.class); - rulePersistenceService.createIndexIfAbsent(listener); - verify(listener).onResponse(true); - verifyNoMoreInteractions(listener); - } + public static final int MAX_VALUES_PER_PAGE = 50; + + public void testGetRuleByIdSuccess() { + GetRuleRequest getRuleRequest = mock(GetRuleRequest.class); + when(getRuleRequest.getId()).thenReturn(_ID_ONE); + when(getRuleRequest.getAttributeFilters()).thenReturn(new HashMap<>()); + QueryBuilder queryBuilder = mock(QueryBuilder.class); + RuleQueryMapper mockRuleQueryMapper = mock(RuleQueryMapper.class); + RuleEntityParser mockRuleEntityParser = mock(RuleEntityParser.class); + Rule mockRule = mock(Rule.class); - public void testCreateIndexIfAbsent() { - IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - Client client = rulePersistenceService.getClient(); - IndicesAdminClient indicesAdminClient = mock(IndicesAdminClient.class); - when(client.admin()).thenReturn(mock(AdminClient.class)); - when(client.admin().indices()).thenReturn(indicesAdminClient); - doAnswer(invocation -> { - ActionListener listener = invocation.getArgument(1); - listener.onResponse(new CreateIndexResponse(true, true, TEST_INDEX_NAME)); // Assuming the index creation was successful + when(mockRuleEntityParser.parse(anyString())).thenReturn(mockRule); + when(mockRuleQueryMapper.from(getRuleRequest)).thenReturn(queryBuilder); + when(queryBuilder.filter(any())).thenReturn(queryBuilder); + + SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + Client client = setUpMockClient(searchRequestBuilder); + + RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( + TEST_INDEX_NAME, + client, + MAX_VALUES_PER_PAGE, + mockRuleEntityParser, + mockRuleQueryMapper + ); + + SearchResponse searchResponse = mock(SearchResponse.class); + SearchHits searchHits = new SearchHits(new SearchHit[] { new SearchHit(1) }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1.0f); + when(searchResponse.getHits()).thenReturn(searchHits); + SearchHit hit = searchHits.getHits()[0]; + hit.sourceRef(new BytesArray(VALID_JSON)); + + ActionListener listener = mock(ActionListener.class); + + doAnswer((invocation) -> { + ActionListener actionListener = invocation.getArgument(0); + actionListener.onResponse(searchResponse); return null; - }).when(indicesAdminClient).create(any(CreateIndexRequest.class), any(ActionListener.class)); - rulePersistenceService.createIndexIfAbsent(new ActionListener<>() { - @Override - public void onResponse(Boolean indexCreated) { - assertTrue(indexCreated); - } - - @Override - public void onFailure(Exception e) { - fail("Index creation failed: " + e.getMessage()); - } - }); - verify(indicesAdminClient).create(any(CreateIndexRequest.class), any(ActionListener.class)); + }).when(searchRequestBuilder).execute(any(ActionListener.class)); + + when(getRuleRequest.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); + rulePersistenceService.getRule(getRuleRequest, listener); + + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(GetRuleResponse.class); + verify(listener).onResponse(responseCaptor.capture()); + GetRuleResponse response = responseCaptor.getValue(); + assertEquals(response.getRules().size(), 1); } - public void testPersistRuleSuccess() { - IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - Client client = rulePersistenceService.getClient(); - ActionListener listener = mock(ActionListener.class); - IndexResponse indexResponse = new IndexResponse(new ShardId(TEST_INDEX_NAME, "uuid", 0), "id", 1, 1, 1, true); - doAnswer(invocation -> { - ActionListener actionListener = invocation.getArgument(1); - actionListener.onResponse(indexResponse); + public void testGetRuleByIdNotFound() { + GetRuleRequest getRuleRequest = mock(GetRuleRequest.class); + when(getRuleRequest.getId()).thenReturn(_ID_ONE); + QueryBuilder queryBuilder = mock(QueryBuilder.class); + RuleQueryMapper mockRuleQueryMapper = mock(RuleQueryMapper.class); + RuleEntityParser mockRuleEntityParser = mock(RuleEntityParser.class); + Rule mockRule = mock(Rule.class); + + when(mockRuleEntityParser.parse(anyString())).thenReturn(mockRule); + when(mockRuleQueryMapper.from(getRuleRequest)).thenReturn(queryBuilder); + when(queryBuilder.filter(any())).thenReturn(queryBuilder); + + SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + Client client = setUpMockClient(searchRequestBuilder); + + RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( + TEST_INDEX_NAME, + client, + MAX_VALUES_PER_PAGE, + mockRuleEntityParser, + mockRuleQueryMapper + ); + + SearchResponse searchResponse = mock(SearchResponse.class); + when(searchResponse.getHits()).thenReturn(new SearchHits(new SearchHit[] {}, new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1.0f)); + + ActionListener listener = mock(ActionListener.class); + + doAnswer(invocationOnMock -> { + ActionListener actionListener = invocationOnMock.getArgument(0); + actionListener.onResponse(searchResponse); return null; - }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); + }).when(searchRequestBuilder).execute(any(ActionListener.class)); - rulePersistenceService.persistRule(ruleOne, listener); - verify(client).index(any(IndexRequest.class), any(ActionListener.class)); - ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(CreateRuleResponse.class); - verify(listener).onResponse(responseCaptor.capture()); + when(getRuleRequest.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); + rulePersistenceService.getRule(getRuleRequest, listener); - CreateRuleResponse createRuleResponse = responseCaptor.getValue(); - assertNotNull(createRuleResponse); - assertEquals(ruleOne, createRuleResponse.getRule()); + ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); + verify(listener).onFailure(exceptionCaptor.capture()); + Exception exception = exceptionCaptor.getValue(); + assertTrue(exception instanceof ResourceNotFoundException); } - public void testPersistRuleFailure() throws IOException { - IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - Client client = rulePersistenceService.getClient(); - ActionListener listener = mock(ActionListener.class); - doAnswer(invocation -> { - ActionListener actionListener = invocation.getArgument(1); - actionListener.onFailure(new RuntimeException("Indexing failed")); - return null; - }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); + private Client setUpMockClient(SearchRequestBuilder searchRequestBuilder) { + Client client = mock(Client.class); + ClusterService clusterService = mock(ClusterService.class); + ClusterState clusterState = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + ThreadPool threadPool = mock(ThreadPool.class); - rulePersistenceService.persistRule(ruleOne, listener); - verify(client).index(any(IndexRequest.class), any(ActionListener.class)); - verify(listener).onFailure(any(RuntimeException.class)); + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + when(client.threadPool()).thenReturn(threadPool); + when(threadPool.getThreadContext()).thenReturn(threadContext); + when(clusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + + when(client.prepareSearch(TEST_INDEX_NAME)).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.setQuery(any(QueryBuilder.class))).thenReturn(searchRequestBuilder); + when(searchRequestBuilder.setSize(anyInt())).thenReturn(searchRequestBuilder); + + return client; } public void testDeleteRule_successful() { @@ -214,75 +243,4 @@ public void testDeleteRule_notFound() { verify(listener).onFailure(any(ResourceNotFoundException.class)); } - - // public void testCreateIndexIfAbsent_IndexExists() { - // IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - // ClusterService clusterService = rulePersistenceService.getClusterService(); - // when(clusterService.state().metadata().hasIndex(TEST_INDEX_NAME)).thenReturn(true); - // ActionListener listener = mock(ActionListener.class); - // rulePersistenceService.createIndexIfAbsent(listener); - // verify(listener).onResponse(true); - // verifyNoMoreInteractions(listener); - // } - // - // public void testCreateIndexIfAbsent() { - // IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - // Client client = rulePersistenceService.getClient(); - // IndicesAdminClient indicesAdminClient = mock(IndicesAdminClient.class); - // when(client.admin()).thenReturn(mock(AdminClient.class)); - // when(client.admin().indices()).thenReturn(indicesAdminClient); - // doAnswer(invocation -> { - // ActionListener listener = invocation.getArgument(1); - // listener.onResponse(new CreateIndexResponse(true, true, TEST_INDEX_NAME)); // Assuming the index creation was successful - // return null; - // }).when(indicesAdminClient).create(any(CreateIndexRequest.class), any(ActionListener.class)); - // rulePersistenceService.createIndexIfAbsent(new ActionListener<>() { - // @Override - // public void onResponse(Boolean indexCreated) { - // assertTrue(indexCreated); - // } - // - // @Override - // public void onFailure(Exception e) { - // fail("Index creation failed: " + e.getMessage()); - // } - // }); - // verify(indicesAdminClient).create(any(CreateIndexRequest.class), any(ActionListener.class)); - // } - // - // public void testPersistRuleSuccess() { - // IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - // Client client = rulePersistenceService.getClient(); - // ActionListener listener = mock(ActionListener.class); - // IndexResponse indexResponse = new IndexResponse(new ShardId(TEST_INDEX_NAME, "uuid", 0), "id", 1, 1, 1, true); - // doAnswer(invocation -> { - // ActionListener actionListener = invocation.getArgument(1); - // actionListener.onResponse(indexResponse); - // return null; - // }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); - // - // rulePersistenceService.persistRule(ruleOne, listener); - // verify(client).index(any(IndexRequest.class), any(ActionListener.class)); - // ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(CreateRuleResponse.class); - // verify(listener).onResponse(responseCaptor.capture()); - // - // CreateRuleResponse createRuleResponse = responseCaptor.getValue(); - // assertNotNull(createRuleResponse); - // assertEquals(ruleOne, createRuleResponse.getRule()); - // } - // - // public void testPersistRuleFailure() throws IOException { - // IndexStoredRulePersistenceService rulePersistenceService = setUpIndexStoredRulePersistenceService(new HashMap<>()); - // Client client = rulePersistenceService.getClient(); - // ActionListener listener = mock(ActionListener.class); - // doAnswer(invocation -> { - // ActionListener actionListener = invocation.getArgument(1); - // actionListener.onFailure(new RuntimeException("Indexing failed")); - // return null; - // }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); - // - // rulePersistenceService.persistRule(ruleOne, listener); - // verify(client).index(any(IndexRequest.class), any(ActionListener.class)); - // verify(listener).onFailure(any(RuntimeException.class)); - // } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java deleted file mode 100644 index 18e1bb99926b7..0000000000000 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleParserTests.java +++ /dev/null @@ -1,52 +0,0 @@ -/* - * 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.rule.utils; - -import org.opensearch.autotagging.Rule; -import org.opensearch.test.OpenSearchTestCase; - -import java.io.IOException; -import java.time.Instant; -import java.util.Locale; - -import static org.opensearch.rule.utils.RuleTestUtils.DESCRIPTION_ONE; - -public class IndexStoredRuleParserTests extends OpenSearchTestCase { - public static final String VALID_JSON = String.format(Locale.ROOT, """ - { - "description": "%s", - "mock_feature_type": "feature value", - "mock_attribute_one": ["attribute_value_one", "attribute_value_two"], - "updated_at": "%s" - } - """, DESCRIPTION_ONE, Instant.now().toString()); - - private static final String INVALID_JSON = """ - { - "name": "TestRule", - "description": "A test rule for unit testing", - "mock_attribute_three": ["attribute_value_one", "attribute_value_two"] - } - """; - - public void testParseRule_Success() throws IOException { - Rule parsedRule = IndexStoredRuleParser.parseRule(VALID_JSON, RuleTestUtils.MockRuleFeatureType.INSTANCE); - assertNotNull(parsedRule); - assertEquals(DESCRIPTION_ONE, parsedRule.getDescription()); - assertEquals(RuleTestUtils.MockRuleFeatureType.INSTANCE, parsedRule.getFeatureType()); - } - - public void testParseRule_InvalidJson() { - Exception exception = assertThrows( - RuntimeException.class, - () -> IndexStoredRuleParser.parseRule(INVALID_JSON, RuleTestUtils.MockRuleFeatureType.INSTANCE) - ); - assertTrue(exception.getMessage().contains("mock_attribute_three is not a valid attribute within the mock_feature_type feature.")); - } -} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java deleted file mode 100644 index 54c6745566c11..0000000000000 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/IndexStoredRuleUtilsTests.java +++ /dev/null @@ -1,69 +0,0 @@ -/* - * 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.rule.utils; - -import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.Rule; -import org.opensearch.index.query.BoolQueryBuilder; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.test.OpenSearchTestCase; - -import java.util.HashMap; -import java.util.Map; -import java.util.Optional; -import java.util.Set; - -import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_MAP; -import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_VALUE_ONE; -import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_VALUE_TWO; -import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; -import static org.opensearch.rule.utils.RuleTestUtils.ruleOne; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public class IndexStoredRuleUtilsTests extends OpenSearchTestCase { - public void testBuildGetRuleQuery_WithId() { - BoolQueryBuilder query = IndexStoredRuleUtils.buildGetRuleQuery( - _ID_ONE, - new HashMap<>(), - RuleTestUtils.MockRuleFeatureType.INSTANCE - ); - assertNotNull(query); - assertEquals(1, query.must().size()); - QueryBuilder idQuery = query.must().get(0); - assertTrue(idQuery.toString().contains(_ID_ONE)); - } - - public void testBuildGetRuleQuery_WithAttributes() { - BoolQueryBuilder query = IndexStoredRuleUtils.buildGetRuleQuery(null, ATTRIBUTE_MAP, RuleTestUtils.MockRuleFeatureType.INSTANCE); - assertNotNull(query); - assertTrue(query.must().size() == 1); - assertTrue(query.toString().contains(RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE.getName())); - assertTrue(query.toString().contains(ATTRIBUTE_VALUE_ONE)); - } - - public void testGetDuplicateRuleId_Found() { - Optional duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(ruleOne, Map.of(_ID_ONE, ruleOne)); - assertFalse(duplicateRuleId.isEmpty()); - assertEquals(_ID_ONE, duplicateRuleId.get()); - } - - public void testGetDuplicateRuleId_NotFound() { - Rule rule = mock(Rule.class); - Map> map = Map.of( - RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, - Set.of(ATTRIBUTE_VALUE_ONE), - RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, - Set.of(ATTRIBUTE_VALUE_TWO) - ); - when(rule.getAttributeMap()).thenReturn(map); - Optional duplicateRuleId = IndexStoredRuleUtils.getDuplicateRuleId(rule, Map.of(_ID_ONE, ruleOne)); - assertTrue(duplicateRuleId.isEmpty()); - } -} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java index fd9a228370af4..feefa61cb8c8f 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java @@ -8,104 +8,27 @@ package org.opensearch.rule.utils; -import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.AutoTaggingRegistry; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule; -import org.opensearch.cluster.ClusterState; -import org.opensearch.cluster.metadata.Metadata; -import org.opensearch.cluster.metadata.QueryGroup; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.settings.Settings; -import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.rule.service.IndexStoredRulePersistenceService; -import org.opensearch.threadpool.ThreadPool; -import org.opensearch.transport.client.Client; +import org.opensearch.rule.autotagging.Attribute; +import org.opensearch.rule.autotagging.AutoTaggingRegistry; +import org.opensearch.rule.autotagging.FeatureType; +import org.opensearch.rule.autotagging.FeatureValueValidator; import java.util.Map; import java.util.Set; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertTrue; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - public class RuleTestUtils { public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; - public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; - public static final String FEATURE_VALUE_ONE = "feature_value_one"; - public static final String FEATURE_VALUE_TWO = "feature_value_two"; public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; public static final String ATTRIBUTE_VALUE_TWO = "mock_attribute_two"; public static final String DESCRIPTION_ONE = "description_1"; - public static final String DESCRIPTION_TWO = "description_2"; - public static final String SEARCH_AFTER = "search_after_id"; - public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; - public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; public static final String FEATURE_TYPE_NAME = "mock_feature_type"; public static final String TEST_INDEX_NAME = ".test_index_for_rule"; public static final Map> ATTRIBUTE_MAP = Map.of( MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, Set.of(ATTRIBUTE_VALUE_ONE) ); - public static final Rule ruleOne = Rule.builder() - .description(DESCRIPTION_ONE) - .featureType(MockRuleFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_ONE) - .attributeMap(ATTRIBUTE_MAP) - .updatedAt(TIMESTAMP_ONE) - .build(); - - public static final Rule ruleTwo = Rule.builder() - .description(DESCRIPTION_TWO) - .featureType(MockRuleFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_TWO) - .attributeMap(Map.of(MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(ATTRIBUTE_VALUE_TWO))) - .updatedAt(TIMESTAMP_TWO) - .build(); - - public static Map ruleMap() { - return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); - } - - public static IndexStoredRulePersistenceService setUpIndexStoredRulePersistenceService(Map queryGroupMap) { - Client client = mock(Client.class); - ClusterService clusterService = mock(ClusterService.class); - ClusterState clusterState = mock(ClusterState.class); - Metadata metadata = mock(Metadata.class); - ThreadPool threadPool = mock(ThreadPool.class); - - ThreadContext threadContext = new ThreadContext(Settings.EMPTY); - when(client.threadPool()).thenReturn(threadPool); - when(threadPool.getThreadContext()).thenReturn(threadContext); - when(clusterService.state()).thenReturn(clusterState); - when(clusterState.metadata()).thenReturn(metadata); - when(metadata.queryGroups()).thenReturn(queryGroupMap); - return new IndexStoredRulePersistenceService(TEST_INDEX_NAME, clusterService, client, MockRuleFeatureType.INSTANCE, 50); - } - public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { - assertEquals(mapOne.size(), mapTwo.size()); - for (Map.Entry entry : mapOne.entrySet()) { - String id = entry.getKey(); - assertTrue(mapTwo.containsKey(id)); - Rule one = mapOne.get(id); - Rule two = mapTwo.get(id); - assertEqualRule(one, two, ruleUpdated); - } - } - - public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { - if (ruleUpdated) { - assertEquals(one.getDescription(), two.getDescription()); - assertEquals(one.getFeatureType(), two.getFeatureType()); - assertEquals(one.getFeatureValue(), two.getFeatureValue()); - assertEquals(one.getAttributeMap(), two.getAttributeMap()); - assertEquals(one.getAttributeMap(), two.getAttributeMap()); - } else { - assertEquals(one, two); - } - } + public static final String INVALID_ATTRIBUTE = "invalid_attribute"; public static class MockRuleFeatureType implements FeatureType { @@ -114,7 +37,7 @@ public static class MockRuleFeatureType implements FeatureType { private MockRuleFeatureType() {} static { - INSTANCE.registerFeatureType(); + AutoTaggingRegistry.registerFeatureType(INSTANCE); } @Override @@ -133,15 +56,15 @@ public Map getAllowedAttributesRegistry() { } @Override - public void registerFeatureType() { - AutoTaggingRegistry.registerFeatureType(INSTANCE); + public FeatureValueValidator getFeatureValueValidator() { + return null; } } public enum MockRuleAttributes implements Attribute { MOCK_RULE_ATTRIBUTE_ONE(ATTRIBUTE_VALUE_ONE), - MOCK_RULE_ATTRIBUTE_TWO(ATTRIBUTE_VALUE_TWO); - ; + MOCK_RULE_ATTRIBUTE_TWO(ATTRIBUTE_VALUE_TWO), + INVALID_ATTRIBUTE(RuleTestUtils.INVALID_ATTRIBUTE); private final String name; diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java index e386c0021475f..e97d13b35308c 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java @@ -28,6 +28,12 @@ import org.opensearch.rule.action.TransportGetRuleAction; import org.opensearch.rule.autotagging.AutoTaggingRegistry; import org.opensearch.rule.rest.RestDeleteRuleAction; +import org.opensearch.rule.action.CreateRuleAction; +import org.opensearch.rule.action.GetRuleAction; +import org.opensearch.rule.action.TransportCreateRuleAction; +import org.opensearch.rule.action.TransportGetRuleAction; +import org.opensearch.rule.autotagging.AutoTaggingRegistry; +import org.opensearch.rule.rest.RestCreateRuleAction; import org.opensearch.rule.rest.RestGetRuleAction; import org.opensearch.rule.spi.RuleFrameworkExtension; @@ -38,6 +44,8 @@ /** * This plugin provides the central APIs which can provide CRUD support to all consumers of Rule framework + * Plugins that define custom rule logic must implement {@link RuleFrameworkExtension}, which ensures + * their feature types and persistence services are automatically registered and available to the Rule Framework. */ public class RuleFrameworkPlugin extends Plugin implements ExtensiblePlugin, ActionPlugin { @@ -55,7 +63,8 @@ public RuleFrameworkPlugin() {} ruleFrameworkExtensions.forEach(this::consumeFrameworkExtension); return List.of( new ActionPlugin.ActionHandler<>(GetRuleAction.INSTANCE, TransportGetRuleAction.class), - new ActionPlugin.ActionHandler<>(DeleteRuleAction.INSTANCE, TransportDeleteRuleAction.class) + new ActionPlugin.ActionHandler<>(DeleteRuleAction.INSTANCE, TransportDeleteRuleAction.class), + new ActionPlugin.ActionHandler<>(CreateRuleAction.INSTANCE, TransportCreateRuleAction.class) ); } @@ -69,7 +78,7 @@ public List getRestHandlers( IndexNameExpressionResolver indexNameExpressionResolver, Supplier nodesInCluster ) { - return List.of(new RestGetRuleAction(), new RestDeleteRuleAction()); + return List.of(new RestGetRuleAction(), new RestDeleteRuleAction(), new RestCreateRuleAction()); } @Override diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/CreateWlmRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleAction.java similarity index 51% rename from plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/CreateWlmRuleAction.java rename to modules/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleAction.java index 8321ed7a89b17..fe7b424b31ca4 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/CreateWlmRuleAction.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/CreateRuleAction.java @@ -6,31 +6,31 @@ * compatible open source license. */ -package org.opensearch.plugin.wlm.rule.action; +package org.opensearch.rule.action; import org.opensearch.action.ActionType; -import org.opensearch.rule.action.CreateRuleResponse; +import org.opensearch.rule.CreateRuleResponse; /** - * Action type for creating a Rule in workload management + * Action type for creating a Rule * @opensearch.experimental */ -public class CreateWlmRuleAction extends ActionType { +public class CreateRuleAction extends ActionType { /** - * An instance of CreateWlmRuleAction + * An instance of CreateRuleAction */ - public static final CreateWlmRuleAction INSTANCE = new CreateWlmRuleAction(); + public static final CreateRuleAction INSTANCE = new CreateRuleAction(); /** - * Name for CreateWlmRuleAction + * Name for CreateRuleAction */ - public static final String NAME = "cluster:admin/opensearch/wlm/rule/_create"; + public static final String NAME = "cluster:admin/opensearch/rule/_create"; /** * Default constructor */ - private CreateWlmRuleAction() { + private CreateRuleAction() { super(NAME, CreateRuleResponse::new); } } diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java new file mode 100644 index 0000000000000..28dd658d107aa --- /dev/null +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java @@ -0,0 +1,54 @@ +/* + * 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.rule.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.rule.CreateRuleRequest; +import org.opensearch.rule.CreateRuleResponse; +import org.opensearch.rule.RulePersistenceService; +import org.opensearch.rule.RulePersistenceServiceRegistry; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +/** + * Transport action to create Rules + * @opensearch.experimental + */ +public class TransportCreateRuleAction extends HandledTransportAction { + + private final RulePersistenceServiceRegistry rulePersistenceServiceRegistry; + + /** + * Constructor for TransportCreateRuleAction + * + * @param transportService - a {@link TransportService} object + * @param actionFilters - a {@link ActionFilters} object + * @param rulePersistenceServiceRegistry - a {@link RulePersistenceServiceRegistry} object + */ + @Inject + public TransportCreateRuleAction( + TransportService transportService, + ActionFilters actionFilters, + RulePersistenceServiceRegistry rulePersistenceServiceRegistry + ) { + super(CreateRuleAction.NAME, transportService, actionFilters, CreateRuleRequest::new); + this.rulePersistenceServiceRegistry = rulePersistenceServiceRegistry; + } + + @Override + protected void doExecute(Task task, CreateRuleRequest request, ActionListener listener) { + final RulePersistenceService rulePersistenceService = rulePersistenceServiceRegistry.getRulePersistenceService( + request.getRule().getFeatureType() + ); + rulePersistenceService.createRule(request, listener); + } +} diff --git a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java similarity index 63% rename from libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java rename to modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java index a6a6a13609553..4aca4c9510b1c 100644 --- a/libs/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java @@ -8,66 +8,58 @@ package org.opensearch.rule.rest; -import org.opensearch.action.ActionType; -import org.opensearch.autotagging.FeatureType; -import org.opensearch.autotagging.Rule.Builder; import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; +import org.opensearch.rest.RestHandler; import org.opensearch.rest.RestRequest; import org.opensearch.rest.RestResponse; import org.opensearch.rest.action.RestResponseListener; -import org.opensearch.rule.action.CreateRuleRequest; -import org.opensearch.rule.action.CreateRuleResponse; +import org.opensearch.rule.CreateRuleRequest; +import org.opensearch.rule.CreateRuleResponse; +import org.opensearch.rule.action.CreateRuleAction; +import org.opensearch.rule.autotagging.FeatureType; +import org.opensearch.rule.autotagging.Rule.Builder; import org.opensearch.transport.client.node.NodeClient; import org.joda.time.Instant; import java.io.IOException; import java.util.List; +import static org.opensearch.rest.RestRequest.Method.POST; +import static org.opensearch.rest.RestRequest.Method.PUT; +import static org.opensearch.rule.rest.RestGetRuleAction.FEATURE_TYPE; + /** * Rest action to create a Rule * @opensearch.experimental */ public class RestCreateRuleAction extends BaseRestHandler { - private final String name; - private final List routes; - private final FeatureType featureType; - private final ActionType createRuleAction; - /** * constructor for RestUpdateRuleAction - * @param name - RestUpdateRuleAction name - * @param routes the list of REST routes this action handles - * @param featureType the feature type associated with the rule - * @param createRuleAction the action to execute for updating a rule */ - public RestCreateRuleAction(String name, List routes, FeatureType featureType, ActionType createRuleAction) { - this.name = name; - this.routes = routes; - this.featureType = featureType; - this.createRuleAction = createRuleAction; - } + public RestCreateRuleAction() {} @Override public String getName() { - return name; + return "create_rule"; } @Override public List routes() { - return routes; + return List.of(new RestHandler.Route(POST, "_rules/{featureType}/"), new RestHandler.Route(PUT, "_rules/{featureType}/")); } @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + final FeatureType featureType = FeatureType.from(request.param(FEATURE_TYPE)); try (XContentParser parser = request.contentParser()) { Builder builder = Builder.fromXContent(parser, featureType); CreateRuleRequest createRuleRequest = new CreateRuleRequest(builder.updatedAt(Instant.now().toString()).build()); - return channel -> client.execute(createRuleAction, createRuleRequest, createRuleResponse(channel)); + return channel -> client.execute(CreateRuleAction.INSTANCE, createRuleRequest, createRuleResponse(channel)); } } diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java index 20a53c345edf7..79eccc09ccdc1 100644 --- a/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java @@ -12,6 +12,7 @@ import org.opensearch.rule.autotagging.Attribute; import org.opensearch.rule.autotagging.AutoTaggingRegistry; import org.opensearch.rule.autotagging.FeatureType; +import org.opensearch.rule.autotagging.FeatureValueValidator; import org.opensearch.rule.autotagging.Rule; import org.opensearch.rule.storage.DefaultAttributeValueStore; import org.opensearch.test.OpenSearchTestCase; @@ -124,7 +125,7 @@ public enum WLMFeatureType implements FeatureType { WLM; static { - WLM.registerFeatureType(); + AutoTaggingRegistry.registerFeatureType(WLM); } @Override @@ -138,8 +139,8 @@ public Map getAllowedAttributesRegistry() { } @Override - public void registerFeatureType() { - AutoTaggingRegistry.registerFeatureType(WLM); + public FeatureValueValidator getFeatureValueValidator() { + return null; } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index dedc02a06e0c6..5342b378079b4 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -38,6 +38,7 @@ import org.opensearch.plugin.wlm.rest.RestGetWorkloadGroupAction; import org.opensearch.plugin.wlm.rest.RestUpdateWorkloadGroupAction; import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType; +import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureValueValidator; import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.Plugin; @@ -51,6 +52,7 @@ import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.rule.spi.RuleFrameworkExtension; import org.opensearch.rule.storage.DefaultAttributeValueStore; +import org.opensearch.rule.storage.IndexBasedRuleDuplicateChecker; import org.opensearch.rule.storage.IndexBasedRuleQueryMapper; import org.opensearch.rule.storage.XContentRuleParser; import org.opensearch.script.ScriptService; @@ -67,6 +69,7 @@ * Plugin class for WorkloadManagement */ public class WorkloadManagementPlugin extends Plugin implements ActionPlugin, SystemIndexPlugin, RuleFrameworkExtension { + /** * The name of the index where rules are stored. */ @@ -99,15 +102,18 @@ public Collection createComponents( IndexNameExpressionResolver indexNameExpressionResolver, Supplier repositoriesServiceSupplier ) { + FeatureTypeHolder.featureType = new WorkloadGroupFeatureType(new WorkloadGroupFeatureValueValidator(clusterService)); RulePersistenceServiceHolder.rulePersistenceService = new IndexStoredRulePersistenceService( INDEX_NAME, client, + clusterService, MAX_RULES_PER_PAGE, - new XContentRuleParser(WorkloadGroupFeatureType.INSTANCE), - new IndexBasedRuleQueryMapper() + new XContentRuleParser(FeatureTypeHolder.featureType), + new IndexBasedRuleQueryMapper(), + new IndexBasedRuleDuplicateChecker() ); InMemoryRuleProcessingService ruleProcessingService = new InMemoryRuleProcessingService( - WorkloadGroupFeatureType.INSTANCE, + FeatureTypeHolder.featureType, DefaultAttributeValueStore::new ); autoTaggingActionFilter = new AutoTaggingActionFilter(ruleProcessingService, threadPool); @@ -169,10 +175,14 @@ public Supplier getRulePersistenceServiceSupplier() { @Override public FeatureType getFeatureType() { - return WorkloadGroupFeatureType.INSTANCE; + return FeatureTypeHolder.featureType; } static class RulePersistenceServiceHolder { private static RulePersistenceService rulePersistenceService; } + + static class FeatureTypeHolder { + private static FeatureType featureType; + } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java index f3ccf9795e93b..bb0f4c7e90122 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPluginModule.java @@ -11,11 +11,6 @@ import org.opensearch.common.inject.AbstractModule; import org.opensearch.common.inject.Singleton; import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; -import org.opensearch.plugin.wlm.rule.service.WlmRulePersistenceService; -import org.opensearch.plugin.wlm.rule.service.WlmRuleProcessingService; -import org.opensearch.plugin.wlm.rule.service.WlmRuleResponseBuilder; -import org.opensearch.rule.service.IndexStoredRulePersistenceService; -import org.opensearch.rule.service.RulePersistenceService; /** * Guice Module to manage WorkloadManagement related objects @@ -27,14 +22,10 @@ public class WorkloadManagementPluginModule extends AbstractModule { */ public WorkloadManagementPluginModule() {} + @Override protected void configure() { // Bind WorkloadGroupPersistenceService as a singleton to ensure a single instance is used, // preventing multiple throttling key registrations in the constructor. bind(WorkloadGroupPersistenceService.class).in(Singleton.class); - bind(RulePersistenceService.class).to(WlmRulePersistenceService.class).in(Singleton.class); - bind(RuleProcessingService.class).to(WlmRuleProcessingService.class).in(Singleton.class); - bind(RuleResponseBuilder.class).to(WlmRuleResponseBuilder.class).in(Singleton.class); - bind(RuleService.class).in(Singleton.class); - bind(RulePersistenceService.class).to(IndexStoredRulePersistenceService.class); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java deleted file mode 100644 index 5357a344da407..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupAttribute.java +++ /dev/null @@ -1,61 +0,0 @@ -/* - * 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.rule; - -import org.opensearch.autotagging.Attribute; - -import java.util.HashMap; -import java.util.Map; - -/** - * Attributes specific to the query group feature. - * @opensearch.experimental - */ -public enum QueryGroupAttribute implements Attribute { - /** - * Represents the index_pattern attribute in QueryGroupAttribute - */ - INDEX_PATTERN("index_pattern"); - - private final String name; - - QueryGroupAttribute(String name) { - this.name = name; - validateAttribute(); - } - - @Override - public String getName() { - return name; - } - - /** - * Retrieves the QueryGroupAttribute from a name string - * @param name - attribute name - */ - public static QueryGroupAttribute fromName(String name) { - for (QueryGroupAttribute attr : QueryGroupAttribute.values()) { - if (attr.getName().equals(name)) { - return attr; - } - } - throw new IllegalArgumentException("Unknown QueryGroupAttribute: " + name); - } - - /** - * Converts the QueryGroupAttribute values into a map with attribute names as keys. - */ - public static Map toMap() { - Map map = new HashMap<>(); - for (QueryGroupAttribute attr : QueryGroupAttribute.values()) { - map.put(attr.getName(), attr); - } - return map; - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java deleted file mode 100644 index 42f45d3d7c5d6..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/QueryGroupFeatureType.java +++ /dev/null @@ -1,64 +0,0 @@ -/* - * 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.rule; - -import org.opensearch.autotagging.Attribute; -import org.opensearch.autotagging.AutoTaggingRegistry; -import org.opensearch.autotagging.FeatureType; - -import java.util.Map; - -/** - * Represents a feature type specific to the query group feature - * @opensearch.experimental - */ -public class QueryGroupFeatureType implements FeatureType { - /** - * The instance for QueryGroupFeatureType - */ - public static final QueryGroupFeatureType INSTANCE = new QueryGroupFeatureType(); - /** - * Name for QueryGroupFeatureType - */ - public static final String NAME = "query_group"; - private static final int MAX_ATTRIBUTE_VALUES = 10; - private static final int MAX_ATTRIBUTE_VALUE_LENGTH = 100; - private static final Map ALLOWED_ATTRIBUTES = QueryGroupAttribute.toMap(); - - private QueryGroupFeatureType() {} - - static { - INSTANCE.registerFeatureType(); - } - - @Override - public String getName() { - return NAME; - } - - @Override - public int getMaxNumberOfValuesPerAttribute() { - return MAX_ATTRIBUTE_VALUES; - } - - @Override - public int getMaxCharLengthPerAttributeValue() { - return MAX_ATTRIBUTE_VALUE_LENGTH; - } - - @Override - public Map getAllowedAttributesRegistry() { - return ALLOWED_ATTRIBUTES; - } - - @Override - public void registerFeatureType() { - AutoTaggingRegistry.registerFeatureType(INSTANCE); - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java index 1dce67ee4e72d..fc9dfa3136277 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java @@ -10,8 +10,8 @@ import org.opensearch.rule.RuleAttribute; import org.opensearch.rule.autotagging.Attribute; -import org.opensearch.rule.autotagging.AutoTaggingRegistry; import org.opensearch.rule.autotagging.FeatureType; +import org.opensearch.rule.autotagging.FeatureValueValidator; import java.util.Map; @@ -20,10 +20,6 @@ * @opensearch.experimental */ public class WorkloadGroupFeatureType implements FeatureType { - /** - * The instance for WorkloadGroupFeatureType - */ - public static final WorkloadGroupFeatureType INSTANCE = new WorkloadGroupFeatureType(); /** * Name for WorkloadGroupFeatureType */ @@ -34,8 +30,15 @@ public class WorkloadGroupFeatureType implements FeatureType { RuleAttribute.INDEX_PATTERN.getName(), RuleAttribute.INDEX_PATTERN ); + private final FeatureValueValidator featureValueValidator; - private WorkloadGroupFeatureType() {} + /** + * constructor for WorkloadGroupFeatureType + * @param featureValueValidator + */ + public WorkloadGroupFeatureType(FeatureValueValidator featureValueValidator) { + this.featureValueValidator = featureValueValidator; + } @Override public String getName() { @@ -58,7 +61,7 @@ public Map getAllowedAttributesRegistry() { } @Override - public void registerFeatureType() { - AutoTaggingRegistry.registerFeatureType(INSTANCE); + public FeatureValueValidator getFeatureValueValidator() { + return featureValueValidator; } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java new file mode 100644 index 0000000000000..0ea7621943615 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java @@ -0,0 +1,40 @@ +/* + * 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.rule; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.rule.autotagging.FeatureValueValidator; + +/** + * Validator for the workload_group feature type + * @opensearch.experimental + */ +public class WorkloadGroupFeatureValueValidator implements FeatureValueValidator { + private final ClusterService clusterService; + private final Logger logger = LogManager.getLogger(WorkloadGroupFeatureValueValidator.class); + + /** + * constructor for WorkloadGroupFeatureValueValidator + * @param clusterService + */ + public WorkloadGroupFeatureValueValidator(ClusterService clusterService) { + this.clusterService = clusterService; + } + + @Override + public void validate(String featureValue) { + if (!clusterService.state().metadata().workloadGroups().containsKey(featureValue)) { + logger.error("{} is not a valid workload group id.", featureValue); + throw new ResourceNotFoundException(featureValue + " is not a valid workload group id."); + } + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java deleted file mode 100644 index 07c6e8a0d1ec0..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/TransportCreateWlmRuleAction.java +++ /dev/null @@ -1,61 +0,0 @@ -/* - * 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.rule.action; - -import org.opensearch.ResourceNotFoundException; -import org.opensearch.action.support.ActionFilters; -import org.opensearch.action.support.HandledTransportAction; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.inject.Inject; -import org.opensearch.core.action.ActionListener; -import org.opensearch.rule.action.CreateRuleRequest; -import org.opensearch.rule.action.CreateRuleResponse; -import org.opensearch.rule.service.RulePersistenceService; -import org.opensearch.tasks.Task; -import org.opensearch.transport.TransportService; - -/** - * Transport action to create Rule in workload management - * @opensearch.experimental - */ -public class TransportCreateWlmRuleAction extends HandledTransportAction { - - private final RulePersistenceService rulePersistenceService; - private final ClusterService clusterService; - - /** - * Constructor for TransportCreateWlmRuleAction - * - * @param transportService - a {@link TransportService} object - * @param actionFilters - a {@link ActionFilters} object - * @param clusterService - a {@link ClusterService} object - * @param rulePersistenceService - a {@link RulePersistenceService} object - */ - @Inject - public TransportCreateWlmRuleAction( - TransportService transportService, - ActionFilters actionFilters, - ClusterService clusterService, - RulePersistenceService rulePersistenceService - ) { - super(CreateWlmRuleAction.NAME, transportService, actionFilters, CreateRuleRequest::new); - this.clusterService = clusterService; - this.rulePersistenceService = rulePersistenceService; - } - - @Override - protected void doExecute(Task task, CreateRuleRequest request, ActionListener listener) { - String queryGroupId = request.getRule().getFeatureValue(); - if (!clusterService.state().metadata().queryGroups().containsKey(queryGroupId)) { - listener.onFailure(new ResourceNotFoundException("Couldn't find an existing query group with id: " + queryGroupId)); - return; - } - rulePersistenceService.createRule(request, listener); - } -} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/package-info.java deleted file mode 100644 index 4fd4eb7ac8eb3..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/action/package-info.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * 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 for the rest classes related to query groups in WorkloadManagementPlugin - */ -package org.opensearch.plugin.wlm.rule.action; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/package-info.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/package-info.java deleted file mode 100644 index 68d60df54d1b4..0000000000000 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/rest/package-info.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * 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 for the action classes related to query groups in WorkloadManagementPlugin - */ -package org.opensearch.plugin.wlm.rule.rest; diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java index c2728a36e9196..f38985e37c426 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java @@ -16,7 +16,7 @@ import java.util.Map; public class WorkloadGroupFeatureTypeTests extends OpenSearchTestCase { - WorkloadGroupFeatureType featureType = WorkloadGroupFeatureType.INSTANCE; + WorkloadGroupFeatureType featureType = new WorkloadGroupFeatureType(null); public void testGetName_returnsCorrectName() { assertEquals("workload_group", featureType.getName()); From 6b0497bc3169e385492ac9aabca20a4bf45124b5 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 29 Apr 2025 01:33:03 -0700 Subject: [PATCH 14/20] modify based on comments Signed-off-by: Ruirui Zhang --- .../opensearch/rule/CreateRuleRequest.java | 8 +- .../opensearch/rule/CreateRuleResponse.java | 2 +- ...Checker.java => DuplicateRuleChecker.java} | 2 +- .../org/opensearch/rule/GetRuleResponse.java | 4 +- .../autotagging/FeatureValueValidator.java | 3 +- .../IndexStoredRulePersistenceService.java | 92 ++++++++----------- ...va => IndexBasedDuplicateRuleChecker.java} | 6 +- .../IndexBasedDuplicateRuleCheckerTests.java | 36 ++++++++ .../rule/rest/RestCreateRuleActionTests.java | 25 +++++ .../plugin/wlm/WorkloadManagementPlugin.java | 4 +- 10 files changed, 117 insertions(+), 65 deletions(-) rename modules/autotagging-commons/common/src/main/java/org/opensearch/rule/{RuleDuplicateChecker.java => DuplicateRuleChecker.java} (94%) rename modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/{IndexBasedRuleDuplicateChecker.java => IndexBasedDuplicateRuleChecker.java} (92%) create mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleCheckerTests.java create mode 100644 modules/autotagging-commons/src/test/java/org/opensearch/rule/rest/RestCreateRuleActionTests.java diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java index fea2a9eda6e27..2234b76c83f82 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java @@ -12,6 +12,7 @@ import org.opensearch.action.ActionRequestValidationException; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.rule.autotagging.FeatureValueValidator; import org.opensearch.rule.autotagging.Rule; import java.io.IOException; @@ -24,7 +25,7 @@ * { * "description": "description1", * "index_pattern": ["log*", "event*"], - * "query_group": "poOiU851RwyLYvV5lbvv5w" + * "workload_group": "poOiU851RwyLYvV5lbvv5w" * }' * @opensearch.experimental */ @@ -51,7 +52,10 @@ public CreateRuleRequest(StreamInput in) throws IOException { @Override public ActionRequestValidationException validate() { try { - rule.getFeatureType().getFeatureValueValidator().validate(rule.getFeatureValue()); + FeatureValueValidator featureValueValidator = rule.getFeatureType().getFeatureValueValidator(); + if (featureValueValidator != null) { + featureValueValidator.validate(rule.getFeatureValue()); + } return null; } catch (Exception e) { ActionRequestValidationException validationException = new ActionRequestValidationException(); diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleResponse.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleResponse.java index 9c95c30ca13c5..f040372b69335 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleResponse.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleResponse.java @@ -28,7 +28,7 @@ * "_id":"wi6VApYBoX5wstmtU_8l", * "description":"description1", * "index_pattern":["log*", "uvent*"], - * "query_group":"poOiU851RwyLYvV5lbvv5w", + * "workload_group":"poOiU851RwyLYvV5lbvv5w", * "updated_at":"2025-04-04T20:54:22.406Z" * } * @opensearch.experimental diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleDuplicateChecker.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/DuplicateRuleChecker.java similarity index 94% rename from modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleDuplicateChecker.java rename to modules/autotagging-commons/common/src/main/java/org/opensearch/rule/DuplicateRuleChecker.java index 0e5eb1742de0f..a852dcc3d4f21 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleDuplicateChecker.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/DuplicateRuleChecker.java @@ -18,7 +18,7 @@ * Interface to check for rule duplication. */ @ExperimentalApi -public interface RuleDuplicateChecker { +public interface DuplicateRuleChecker { /** * Checks if the given rule already exists in the provided rule map. * @param rule the rule to check for duplication diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/GetRuleResponse.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/GetRuleResponse.java index e3c0bb49043a7..2ce79850084db 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/GetRuleResponse.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/GetRuleResponse.java @@ -29,9 +29,9 @@ * "rules": [ * { * "_id": "z1MJApUB0zgMcDmz-UQq", - * "description": "Rule for tagging query_group_id to index123" + * "description": "Rule for tagging workload_group_id to index123" * "index_pattern": ["index123"], - * "query_group": "query_group_id", + * "workload_group": "workload_group_id", * "updated_at": "2025-02-14T01:19:22.589Z" * }, * ... diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureValueValidator.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureValueValidator.java index 40b97e9ddd843..7c21982ce9f2a 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureValueValidator.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureValueValidator.java @@ -9,7 +9,8 @@ package org.opensearch.rule.autotagging; /** - * Validates a feature value for a specific feature type + * Interface for validating a feature value against pre-defined values (such as + * values from the index, cluster state, etc.) for a specific feature type. * @opensearch.experimental */ public interface FeatureValueValidator { diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java index 3bdf8d9338411..e8822910f870e 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -32,9 +32,9 @@ import org.opensearch.index.query.QueryBuilder; import org.opensearch.rule.CreateRuleRequest; import org.opensearch.rule.CreateRuleResponse; +import org.opensearch.rule.DuplicateRuleChecker; import org.opensearch.rule.GetRuleRequest; import org.opensearch.rule.GetRuleResponse; -import org.opensearch.rule.RuleDuplicateChecker; import org.opensearch.rule.RuleEntityParser; import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.RuleQueryMapper; @@ -66,7 +66,7 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService private final int maxRulesPerPage; private final RuleEntityParser parser; private final RuleQueryMapper queryBuilder; - private final RuleDuplicateChecker ruleDuplicateChecker; + private final DuplicateRuleChecker ruleDuplicateChecker; private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class); private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); @@ -88,7 +88,7 @@ public IndexStoredRulePersistenceService( int maxRulesPerPage, RuleEntityParser parser, RuleQueryMapper queryBuilder, - RuleDuplicateChecker ruleDuplicateChecker + DuplicateRuleChecker ruleDuplicateChecker ) { this.indexName = indexName; this.client = client; @@ -107,66 +107,51 @@ public IndexStoredRulePersistenceService( */ public void createRule(CreateRuleRequest request, ActionListener listener) { try (ThreadContext.StoredContext ctx = getContext()) { - createIndexIfAbsent(new ActionListener<>() { - @Override - public void onResponse(Void unused) { - validateNoDuplicateRule(request.getRule(), new ActionListener<>() { - @Override - public void onResponse(Void unused) { - persistRule(request.getRule(), listener); - } - - @Override - public void onFailure(Exception e) { - listener.onFailure(e); - } - }); - } - - @Override - public void onFailure(Exception e) { - listener.onFailure(e); - } - }); + if (!clusterService.state().metadata().hasIndex(indexName)) { + createIndex(listener, request.getRule()); + } else { + validateAndPersist(request.getRule(), listener); + } } } /** - * Creates the system index if it doesn't exist + * Creates the system index, then validates and persists the given rule. * @param listener - ActionListener for CreateRuleResponse + * @param rule - the rule to validate and persist */ - private void createIndexIfAbsent(ActionListener listener) { - try (ThreadContext.StoredContext ctx = getContext()) { - if (clusterService.state().metadata().hasIndex(indexName)) { - listener.onResponse(null); - return; + private void createIndex(ActionListener listener, Rule rule) { + final CreateIndexRequest request = new CreateIndexRequest(indexName).settings(indexSettings); + client.admin().indices().create(request, new ActionListener<>() { + @Override + public void onResponse(CreateIndexResponse response) { + if (!response.isAcknowledged()) { + logger.error("Index creation not acknowledged: {}", indexName); + listener.onFailure(new IllegalStateException(indexName + " index creation failed and rule cannot be persisted")); + } else { + validateAndPersist(rule, listener); + } } - final CreateIndexRequest request = new CreateIndexRequest(indexName).settings(indexSettings); - client.admin().indices().create(request, new ActionListener<>() { - @Override - public void onResponse(CreateIndexResponse response) { - if (!response.isAcknowledged()) { - logger.error("Index creation not acknowledged: {}", indexName); - listener.onFailure(new IllegalStateException(indexName + " index creation failed and rule cannot be persisted")); - } else { - logger.info("Index {} created", indexName); - listener.onResponse(null); - } + @Override + public void onFailure(Exception e) { + if (e instanceof ResourceAlreadyExistsException) { + validateAndPersist(rule, listener); + } else { + logger.error("Failed to create index {}: {}", indexName, e.getMessage()); + listener.onFailure(e); } + } + }); + } - @Override - public void onFailure(Exception e) { - if (e instanceof ResourceAlreadyExistsException) { - logger.trace("Index {} already exists", indexName); - listener.onResponse(null); - } else { - logger.error("Failed to create index {}: {}", indexName, e.getMessage()); - listener.onFailure(e); - } - } - }); - } + /** + * Validates that the rule does not already exist, then persists it if validation succeeds + * @param rule - The rule to validate and persist + * @param listener - ActionListener for CreateRuleResponse + */ + private void validateAndPersist(Rule rule, ActionListener listener) { + validateNoDuplicateRule(rule, ActionListener.wrap(unused -> persistRule(rule, listener), listener::onFailure)); } /** @@ -175,6 +160,7 @@ public void onFailure(Exception e) { * @param rule - the rule we check duplicate against * @param listener - listener for validateNoDuplicateRule response */ + private void validateNoDuplicateRule(Rule rule, ActionListener listener) { try (ThreadContext.StoredContext ctx = getContext()) { QueryBuilder query = queryBuilder.from(new GetRuleRequest(null, rule.getAttributeMap(), null, rule.getFeatureType())); diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleDuplicateChecker.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleChecker.java similarity index 92% rename from modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleDuplicateChecker.java rename to modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleChecker.java index dd0800884b8a4..7eb5457eab8fe 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleDuplicateChecker.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleChecker.java @@ -9,7 +9,7 @@ package org.opensearch.rule.storage; import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.rule.RuleDuplicateChecker; +import org.opensearch.rule.DuplicateRuleChecker; import org.opensearch.rule.autotagging.Attribute; import org.opensearch.rule.autotagging.Rule; @@ -21,12 +21,12 @@ * This class is used to check rule duplication for indexed based rules. */ @ExperimentalApi -public class IndexBasedRuleDuplicateChecker implements RuleDuplicateChecker { +public class IndexBasedDuplicateRuleChecker implements DuplicateRuleChecker { /** * Default constructor */ - public IndexBasedRuleDuplicateChecker() {} + public IndexBasedDuplicateRuleChecker() {} /** * Checks if a duplicate rule exists based on the attribute map. diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleCheckerTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleCheckerTests.java new file mode 100644 index 0000000000000..f7fcc9c32e32a --- /dev/null +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleCheckerTests.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.rule.storage; + +import org.opensearch.rule.autotagging.Rule; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.Map; +import java.util.Optional; + +import static org.opensearch.rule.action.GetRuleRequestTests._ID_ONE; +import static org.opensearch.rule.action.GetRuleRequestTests.ruleOne; + +public class IndexBasedDuplicateRuleCheckerTests extends OpenSearchTestCase { + IndexBasedDuplicateRuleChecker checker = new IndexBasedDuplicateRuleChecker(); + + public void testNoDuplicate() { + Map existingRules = Collections.emptyMap(); + Optional result = checker.getDuplicateRuleId(ruleOne, existingRules); + assertFalse(result.isPresent()); + } + + public void testDuplicateExists() { + Map existingRules = Map.of(_ID_ONE, ruleOne); + Optional result = checker.getDuplicateRuleId(ruleOne, existingRules); + assertTrue(result.isPresent()); + assertEquals(result.get(), _ID_ONE); + } +} diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/rest/RestCreateRuleActionTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/rest/RestCreateRuleActionTests.java new file mode 100644 index 0000000000000..cf8fb4068f2fc --- /dev/null +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/rest/RestCreateRuleActionTests.java @@ -0,0 +1,25 @@ +/* + * 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.rule.rest; + +import org.opensearch.test.OpenSearchTestCase; + +public class RestCreateRuleActionTests extends OpenSearchTestCase { + RestCreateRuleAction action = new RestCreateRuleAction();; + + public void testGetName() { + assertEquals("create_rule", action.getName()); + } + + public void testRoutes() { + var routes = action.routes(); + assertEquals(2, routes.size()); + assertTrue(routes.stream().anyMatch(r -> r.getPath().equals("_rules/{featureType}/"))); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index 5342b378079b4..b4c92d9e37ac5 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -52,7 +52,7 @@ import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.rule.spi.RuleFrameworkExtension; import org.opensearch.rule.storage.DefaultAttributeValueStore; -import org.opensearch.rule.storage.IndexBasedRuleDuplicateChecker; +import org.opensearch.rule.storage.IndexBasedDuplicateRuleChecker; import org.opensearch.rule.storage.IndexBasedRuleQueryMapper; import org.opensearch.rule.storage.XContentRuleParser; import org.opensearch.script.ScriptService; @@ -110,7 +110,7 @@ public Collection createComponents( MAX_RULES_PER_PAGE, new XContentRuleParser(FeatureTypeHolder.featureType), new IndexBasedRuleQueryMapper(), - new IndexBasedRuleDuplicateChecker() + new IndexBasedDuplicateRuleChecker() ); InMemoryRuleProcessingService ruleProcessingService = new InMemoryRuleProcessingService( FeatureTypeHolder.featureType, From 8dbb6583465fce6ecb51d6d19dc5e336ecff4950 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 30 Apr 2025 21:29:33 -0700 Subject: [PATCH 15/20] address comments Signed-off-by: Ruirui Zhang --- .../java/org/opensearch/rule/CreateRuleRequest.java | 7 ++++--- .../service/IndexStoredRulePersistenceService.java | 6 ++++-- .../org/opensearch/rule/action/GetRuleAction.java | 2 +- .../opensearch/rule/rest/RestCreateRuleAction.java | 2 +- .../plugin/wlm/WorkloadManagementPluginTests.java | 13 +++++++++++++ 5 files changed, 23 insertions(+), 7 deletions(-) diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java index 2234b76c83f82..aedd091fa829e 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java @@ -53,13 +53,14 @@ public CreateRuleRequest(StreamInput in) throws IOException { public ActionRequestValidationException validate() { try { FeatureValueValidator featureValueValidator = rule.getFeatureType().getFeatureValueValidator(); - if (featureValueValidator != null) { - featureValueValidator.validate(rule.getFeatureValue()); + if (featureValueValidator == null) { + throw new IllegalStateException("FeatureValueValidator is not defined for feature type " + rule.getFeatureType().getName()); } + featureValueValidator.validate(rule.getFeatureValue()); return null; } catch (Exception e) { ActionRequestValidationException validationException = new ActionRequestValidationException(); - validationException.addValidationError(e.getMessage()); + validationException.addValidationError("Validation failed: " + e.getMessage()); return validationException; } } diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java index e8822910f870e..f2908eea1207b 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -10,6 +10,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.ExceptionsHelper; import org.opensearch.ResourceAlreadyExistsException; import org.opensearch.ResourceNotFoundException; import org.opensearch.action.DocWriteResponse; @@ -135,10 +136,11 @@ public void onResponse(CreateIndexResponse response) { @Override public void onFailure(Exception e) { - if (e instanceof ResourceAlreadyExistsException) { + Throwable cause = ExceptionsHelper.unwrapCause(e); + if (cause instanceof ResourceAlreadyExistsException) { validateAndPersist(rule, listener); } else { - logger.error("Failed to create index {}: {}", indexName, e.getMessage()); + logger.error("Failed to create index {}: {}", indexName, e.getMessage(), e); listener.onFailure(e); } } diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleAction.java index e59eabf682510..aa795d3aa4f72 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleAction.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/GetRuleAction.java @@ -18,7 +18,7 @@ public class GetRuleAction extends ActionType { /** - * An instance of GetWlmRuleAction + * An instance of GetRuleAction */ public static final GetRuleAction INSTANCE = new GetRuleAction(); diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java index 4aca4c9510b1c..2c680820e87f4 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java @@ -39,7 +39,7 @@ */ public class RestCreateRuleAction extends BaseRestHandler { /** - * constructor for RestUpdateRuleAction + * constructor for RestCreateRuleAction */ public RestCreateRuleAction() {} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java index 6838102ac3bf1..b79ba4107617b 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java @@ -98,6 +98,19 @@ public void testGetSystemIndexDescriptorsReturnsCorrectDescriptor() { } public void testGetFeatureTypeReturnsWorkloadGroupFeatureType() { + plugin.createComponents( + mock(Client.class), + mock(ClusterService.class), + mock(ThreadPool.class), + mock(ResourceWatcherService.class), + mock(ScriptService.class), + mock(NamedXContentRegistry.class), + mock(Environment.class), + null, + mock(NamedWriteableRegistry.class), + mock(IndexNameExpressionResolver.class), + () -> mock(RepositoriesService.class) + ); FeatureType featureType = plugin.getFeatureType(); assertEquals("workload_group", featureType.getName()); } From 1af4bb78d24ae514142edc79fc727568b24d0bcb Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 6 May 2025 06:23:41 -0700 Subject: [PATCH 16/20] revise based on comments Signed-off-by: Ruirui Zhang --- .../opensearch/rule/CreateRuleRequest.java | 6 +- .../opensearch/rule/DuplicateRuleChecker.java | 28 -------- .../java/org/opensearch/rule/RuleUtils.java | 67 +++++++++++++++++++ .../rule/autotagging/AutoTaggingRegistry.java | 3 + .../rule/autotagging/FeatureType.java | 7 +- .../IndexStoredRulePersistenceService.java | 12 ++-- .../IndexBasedDuplicateRuleChecker.java | 59 ---------------- .../org/opensearch/rule/RuleUtilsTests.java | 60 +++++++++++++++++ .../rule/autotagging/RuleTests.java | 5 -- .../IndexBasedDuplicateRuleCheckerTests.java | 36 ---------- .../opensearch/rule/utils/RuleTestUtils.java | 5 -- .../rule/rest/RestCreateRuleAction.java | 2 +- .../InMemoryRuleProcessingServiceTests.java | 8 +-- .../rule/rest/RestCreateRuleActionTests.java | 4 +- .../plugin/wlm/WorkloadManagementPlugin.java | 18 ++--- .../wlm/rule/WorkloadGroupFeatureType.java | 16 ++++- .../WorkloadGroupFeatureValueValidator.java | 14 +++- .../rule/WorkloadGroupFeatureTypeTests.java | 14 +++- 18 files changed, 196 insertions(+), 168 deletions(-) delete mode 100644 modules/autotagging-commons/common/src/main/java/org/opensearch/rule/DuplicateRuleChecker.java create mode 100644 modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleUtils.java delete mode 100644 modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleChecker.java create mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleUtilsTests.java delete mode 100644 modules/autotagging-commons/common/src/test/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleCheckerTests.java diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java index aedd091fa829e..2a1aa0c980906 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java @@ -52,11 +52,7 @@ public CreateRuleRequest(StreamInput in) throws IOException { @Override public ActionRequestValidationException validate() { try { - FeatureValueValidator featureValueValidator = rule.getFeatureType().getFeatureValueValidator(); - if (featureValueValidator == null) { - throw new IllegalStateException("FeatureValueValidator is not defined for feature type " + rule.getFeatureType().getName()); - } - featureValueValidator.validate(rule.getFeatureValue()); + rule.getFeatureType().getFeatureValueValidator().validate(rule.getFeatureValue()); return null; } catch (Exception e) { ActionRequestValidationException validationException = new ActionRequestValidationException(); diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/DuplicateRuleChecker.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/DuplicateRuleChecker.java deleted file mode 100644 index a852dcc3d4f21..0000000000000 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/DuplicateRuleChecker.java +++ /dev/null @@ -1,28 +0,0 @@ -/* - * 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.rule; - -import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.rule.autotagging.Rule; - -import java.util.Map; -import java.util.Optional; - -/** - * Interface to check for rule duplication. - */ -@ExperimentalApi -public interface DuplicateRuleChecker { - /** - * Checks if the given rule already exists in the provided rule map. - * @param rule the rule to check for duplication - * @param ruleMap a map of existing rules, keyed by rule ID - */ - Optional getDuplicateRuleId(Rule rule, Map ruleMap); -} diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleUtils.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleUtils.java new file mode 100644 index 0000000000000..7c66eac988f9b --- /dev/null +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleUtils.java @@ -0,0 +1,67 @@ +/* + * 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.rule; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.rule.autotagging.Attribute; +import org.opensearch.rule.autotagging.Rule; + +import java.util.Collections; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** + * Utility class for operations related to {@link Rule} objects. + * @opensearch.experimental + */ +@ExperimentalApi +public class RuleUtils { + + /** + * constructor for RuleUtils + */ + public RuleUtils() {} + + /** + * Checks if a duplicate rule exists and returns its id. + * Two rules are considered to be duplicate when meeting all the criteria below + * 1. They have the same feature type + * 2. They have the exact same attributes + * 3. For each attribute, the sets of values must intersect — i.e., at least one common value must exist + * between the current rule and the one being checked. + * + * @param rule The rule to be validated against ruleMap. + * @param ruleMap This map contains existing rules to be checked + */ + public static Optional getDuplicateRuleId(Rule rule, Map ruleMap) { + Map> targetAttributeMap = rule.getAttributeMap(); + for (Map.Entry entry : ruleMap.entrySet()) { + Rule currRule = entry.getValue(); + Map> existingAttributeMap = currRule.getAttributeMap(); + + if (rule.getFeatureType() != currRule.getFeatureType() || targetAttributeMap.size() != existingAttributeMap.size()) { + continue; + } + boolean allAttributesIntersect = true; + for (Attribute attribute : targetAttributeMap.keySet()) { + Set targetAttributeValues = targetAttributeMap.get(attribute); + Set existingAttributeValues = existingAttributeMap.get(attribute); + if (existingAttributeValues == null || Collections.disjoint(targetAttributeValues, existingAttributeValues)) { + allAttributesIntersect = false; + break; + } + } + if (allAttributesIntersect) { + return Optional.of(entry.getKey()); + } + } + return Optional.empty(); + } +} diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/AutoTaggingRegistry.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/AutoTaggingRegistry.java index 9cfc5ccb1e342..be817e66cbd7a 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/AutoTaggingRegistry.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/AutoTaggingRegistry.java @@ -59,6 +59,9 @@ private static void validateFeatureType(FeatureType featureType) { "Feature type name " + name + " should not be null, empty or have more than " + MAX_FEATURE_TYPE_NAME_LENGTH + "characters" ); } + if (featureType.getFeatureValueValidator() == null) { + throw new IllegalStateException("FeatureValueValidator is not defined for feature type " + name); + } } /** diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureType.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureType.java index f56d6d3da5589..9fc2cf62b462d 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureType.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureType.java @@ -50,7 +50,12 @@ public interface FeatureType extends Writeable { /** * returns the validator for feature value */ - FeatureValueValidator getFeatureValueValidator(); + default FeatureValueValidator getFeatureValueValidator() { + return new FeatureValueValidator() { + @Override + public void validate(String featureValue) {} + }; + } /** * returns max attribute values diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java index f2908eea1207b..4e1f02f9c01ba 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -33,12 +33,12 @@ import org.opensearch.index.query.QueryBuilder; import org.opensearch.rule.CreateRuleRequest; import org.opensearch.rule.CreateRuleResponse; -import org.opensearch.rule.DuplicateRuleChecker; import org.opensearch.rule.GetRuleRequest; import org.opensearch.rule.GetRuleResponse; import org.opensearch.rule.RuleEntityParser; import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.RuleQueryMapper; +import org.opensearch.rule.RuleUtils; import org.opensearch.rule.autotagging.Rule; import org.opensearch.search.SearchHit; import org.opensearch.search.sort.SortOrder; @@ -67,7 +67,6 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService private final int maxRulesPerPage; private final RuleEntityParser parser; private final RuleQueryMapper queryBuilder; - private final DuplicateRuleChecker ruleDuplicateChecker; private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class); private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); @@ -80,7 +79,6 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService * @param maxRulesPerPage - The maximum number of rules that can be returned in a single get request. * @param parser * @param queryBuilder - * @param ruleDuplicateChecker */ public IndexStoredRulePersistenceService( String indexName, @@ -88,8 +86,7 @@ public IndexStoredRulePersistenceService( ClusterService clusterService, int maxRulesPerPage, RuleEntityParser parser, - RuleQueryMapper queryBuilder, - DuplicateRuleChecker ruleDuplicateChecker + RuleQueryMapper queryBuilder ) { this.indexName = indexName; this.client = client; @@ -97,7 +94,6 @@ public IndexStoredRulePersistenceService( this.maxRulesPerPage = maxRulesPerPage; this.parser = parser; this.queryBuilder = queryBuilder; - this.ruleDuplicateChecker = ruleDuplicateChecker; } /** @@ -140,7 +136,7 @@ public void onFailure(Exception e) { if (cause instanceof ResourceAlreadyExistsException) { validateAndPersist(rule, listener); } else { - logger.error("Failed to create index {}: {}", indexName, e.getMessage(), e); + logger.error("Failed to create index {}: {}", indexName, e.getMessage()); listener.onFailure(e); } } @@ -169,7 +165,7 @@ private void validateNoDuplicateRule(Rule rule, ActionListener listener) { getRuleFromIndex(null, query, null, new ActionListener<>() { @Override public void onResponse(GetRuleResponse getRuleResponse) { - Optional duplicateRuleId = ruleDuplicateChecker.getDuplicateRuleId(rule, getRuleResponse.getRules()); + Optional duplicateRuleId = RuleUtils.getDuplicateRuleId(rule, getRuleResponse.getRules()); duplicateRuleId.ifPresentOrElse( id -> listener.onFailure(new IllegalArgumentException("Duplicate rule exists under id " + id)), () -> listener.onResponse(null) diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleChecker.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleChecker.java deleted file mode 100644 index 7eb5457eab8fe..0000000000000 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleChecker.java +++ /dev/null @@ -1,59 +0,0 @@ -/* - * 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.rule.storage; - -import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.rule.DuplicateRuleChecker; -import org.opensearch.rule.autotagging.Attribute; -import org.opensearch.rule.autotagging.Rule; - -import java.util.Map; -import java.util.Optional; -import java.util.Set; - -/** - * This class is used to check rule duplication for indexed based rules. - */ -@ExperimentalApi -public class IndexBasedDuplicateRuleChecker implements DuplicateRuleChecker { - - /** - * Default constructor - */ - public IndexBasedDuplicateRuleChecker() {} - - /** - * Checks if a duplicate rule exists based on the attribute map. - * A rule is considered a duplicate when the attribute value already exists in the index, and the number of - * attributes in the new rule is equal to the number of attributes in an existing rule. - * - * For example, if an existing rule has: - * attribute1 = ['a'] and attribute2 = ['c'] - * And we are creating a new rule with: - * attribute1 = ['a'] - * Then it's not a duplicate because the existing rule has attribute2 and is more granular - * - * @param rule The rule to be validated against ruleMap. - * @param ruleMap This map entries are Rules that contain the attribute values from rule, meaning they - * have a partial or complete overlap with the new rule being created. - */ - @Override - public Optional getDuplicateRuleId(Rule rule, Map ruleMap) { - Map> attributeMapToValidate = rule.getAttributeMap(); - for (Map.Entry entry : ruleMap.entrySet()) { - String ruleId = entry.getKey(); - Rule currRule = entry.getValue(); - // Compare the size of the attribute maps to ensure we only check for duplicates with the same number of attributes. - if (attributeMapToValidate.size() == currRule.getAttributeMap().size()) { - return Optional.of(ruleId); - } - } - return Optional.empty(); - } -} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleUtilsTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleUtilsTests.java new file mode 100644 index 0000000000000..e2157797a9b6c --- /dev/null +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleUtilsTests.java @@ -0,0 +1,60 @@ +/* + * 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.rule; + +import org.opensearch.rule.autotagging.Rule; +import org.opensearch.rule.utils.RuleTestUtils; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static org.opensearch.rule.action.GetRuleRequestTests.ATTRIBUTE_VALUE_ONE; +import static org.opensearch.rule.action.GetRuleRequestTests.ATTRIBUTE_VALUE_TWO; +import static org.opensearch.rule.action.GetRuleRequestTests.DESCRIPTION_ONE; +import static org.opensearch.rule.action.GetRuleRequestTests.FEATURE_VALUE_ONE; +import static org.opensearch.rule.action.GetRuleRequestTests.TIMESTAMP_ONE; +import static org.opensearch.rule.action.GetRuleRequestTests._ID_ONE; +import static org.opensearch.rule.action.GetRuleRequestTests._ID_TWO; +import static org.opensearch.rule.action.GetRuleRequestTests.ruleTwo; +import static org.opensearch.rule.action.GetRuleResponseTests.ruleOne; + +public class RuleUtilsTests extends OpenSearchTestCase { + + public void testDuplicateRuleFound() { + Optional result = RuleUtils.getDuplicateRuleId(ruleOne, Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo)); + assertTrue(result.isPresent()); + assertEquals(_ID_ONE, result.get()); + } + + public void testNoDuplicate_NoAttributeIntersection() { + Optional result = RuleUtils.getDuplicateRuleId(ruleOne, Map.of(_ID_TWO, ruleTwo)); + assertTrue(result.isEmpty()); + } + + public void testNoDuplicate_AttributeSizeMismatch() { + Rule testRule = Rule.builder() + .description(DESCRIPTION_ONE) + .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_ONE) + .attributeMap( + Map.of( + RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, + Set.of(ATTRIBUTE_VALUE_ONE), + RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, + Set.of(ATTRIBUTE_VALUE_TWO) + ) + ) + .updatedAt(TIMESTAMP_ONE) + .build(); + Optional result = RuleUtils.getDuplicateRuleId(ruleOne, Map.of(_ID_TWO, testRule)); + assertTrue(result.isEmpty()); + } +} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java index 52687990f235b..c504e1ab6b53e 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java @@ -114,11 +114,6 @@ public int getMaxCharLengthPerAttributeValue() { public Map getAllowedAttributesRegistry() { return ALLOWED_ATTRIBUTES; } - - @Override - public FeatureValueValidator getFeatureValueValidator() { - return null; - } } static Rule buildRule( diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleCheckerTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleCheckerTests.java deleted file mode 100644 index f7fcc9c32e32a..0000000000000 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/storage/IndexBasedDuplicateRuleCheckerTests.java +++ /dev/null @@ -1,36 +0,0 @@ -/* - * 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.rule.storage; - -import org.opensearch.rule.autotagging.Rule; -import org.opensearch.test.OpenSearchTestCase; - -import java.util.Collections; -import java.util.Map; -import java.util.Optional; - -import static org.opensearch.rule.action.GetRuleRequestTests._ID_ONE; -import static org.opensearch.rule.action.GetRuleRequestTests.ruleOne; - -public class IndexBasedDuplicateRuleCheckerTests extends OpenSearchTestCase { - IndexBasedDuplicateRuleChecker checker = new IndexBasedDuplicateRuleChecker(); - - public void testNoDuplicate() { - Map existingRules = Collections.emptyMap(); - Optional result = checker.getDuplicateRuleId(ruleOne, existingRules); - assertFalse(result.isPresent()); - } - - public void testDuplicateExists() { - Map existingRules = Map.of(_ID_ONE, ruleOne); - Optional result = checker.getDuplicateRuleId(ruleOne, existingRules); - assertTrue(result.isPresent()); - assertEquals(result.get(), _ID_ONE); - } -} diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java index feefa61cb8c8f..87ee679031916 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java @@ -54,11 +54,6 @@ public Map getAllowedAttributesRegistry() { MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO ); } - - @Override - public FeatureValueValidator getFeatureValueValidator() { - return null; - } } public enum MockRuleAttributes implements Attribute { diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java index 2c680820e87f4..7068879e7ba54 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java @@ -50,7 +50,7 @@ public String getName() { @Override public List routes() { - return List.of(new RestHandler.Route(POST, "_rules/{featureType}/"), new RestHandler.Route(PUT, "_rules/{featureType}/")); + return List.of(new RestHandler.Route(PUT, "_rules/{featureType}")); } @Override diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java index 79eccc09ccdc1..4e024f3600c35 100644 --- a/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java @@ -138,10 +138,10 @@ public Map getAllowedAttributesRegistry() { return Map.of("test_attribute", TestAttribute.TEST_ATTRIBUTE); } - @Override - public FeatureValueValidator getFeatureValueValidator() { - return null; - } +// @Override +// public FeatureValueValidator getFeatureValueValidator() { +// return null; +// } } public enum TestAttribute implements Attribute { diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/rest/RestCreateRuleActionTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/rest/RestCreateRuleActionTests.java index cf8fb4068f2fc..b6020c7b08e1c 100644 --- a/modules/autotagging-commons/src/test/java/org/opensearch/rule/rest/RestCreateRuleActionTests.java +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/rest/RestCreateRuleActionTests.java @@ -19,7 +19,7 @@ public void testGetName() { public void testRoutes() { var routes = action.routes(); - assertEquals(2, routes.size()); - assertTrue(routes.stream().anyMatch(r -> r.getPath().equals("_rules/{featureType}/"))); + assertEquals(1, routes.size()); + assertTrue(routes.stream().anyMatch(r -> r.getPath().equals("_rules/{featureType}"))); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index b4c92d9e37ac5..eb40e464f1da6 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -49,10 +49,10 @@ import org.opensearch.rule.InMemoryRuleProcessingService; import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.autotagging.FeatureType; +import org.opensearch.rule.autotagging.FeatureValueValidator; import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.rule.spi.RuleFrameworkExtension; import org.opensearch.rule.storage.DefaultAttributeValueStore; -import org.opensearch.rule.storage.IndexBasedDuplicateRuleChecker; import org.opensearch.rule.storage.IndexBasedRuleQueryMapper; import org.opensearch.rule.storage.XContentRuleParser; import org.opensearch.script.ScriptService; @@ -102,18 +102,18 @@ public Collection createComponents( IndexNameExpressionResolver indexNameExpressionResolver, Supplier repositoriesServiceSupplier ) { - FeatureTypeHolder.featureType = new WorkloadGroupFeatureType(new WorkloadGroupFeatureValueValidator(clusterService)); + FeatureValueValidator validator = WorkloadGroupFeatureValueValidator.getInstance(clusterService); + WorkloadGroupFeatureType.initializeFeatureValueValidator(validator); RulePersistenceServiceHolder.rulePersistenceService = new IndexStoredRulePersistenceService( INDEX_NAME, client, clusterService, MAX_RULES_PER_PAGE, - new XContentRuleParser(FeatureTypeHolder.featureType), - new IndexBasedRuleQueryMapper(), - new IndexBasedDuplicateRuleChecker() + new XContentRuleParser(WorkloadGroupFeatureType.getInstance()), + new IndexBasedRuleQueryMapper() ); InMemoryRuleProcessingService ruleProcessingService = new InMemoryRuleProcessingService( - FeatureTypeHolder.featureType, + WorkloadGroupFeatureType.getInstance(), DefaultAttributeValueStore::new ); autoTaggingActionFilter = new AutoTaggingActionFilter(ruleProcessingService, threadPool); @@ -175,14 +175,10 @@ public Supplier getRulePersistenceServiceSupplier() { @Override public FeatureType getFeatureType() { - return FeatureTypeHolder.featureType; + return WorkloadGroupFeatureType.getInstance(); } static class RulePersistenceServiceHolder { private static RulePersistenceService rulePersistenceService; } - - static class FeatureTypeHolder { - private static FeatureType featureType; - } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java index fc9dfa3136277..f934ae15cda27 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java @@ -31,15 +31,29 @@ public class WorkloadGroupFeatureType implements FeatureType { RuleAttribute.INDEX_PATTERN ); private final FeatureValueValidator featureValueValidator; + private static WorkloadGroupFeatureType instance; /** * constructor for WorkloadGroupFeatureType * @param featureValueValidator */ - public WorkloadGroupFeatureType(FeatureValueValidator featureValueValidator) { + private WorkloadGroupFeatureType(FeatureValueValidator featureValueValidator) { this.featureValueValidator = featureValueValidator; } + public static void initializeFeatureValueValidator(FeatureValueValidator validator) { + if (instance == null) { + instance = new WorkloadGroupFeatureType(validator); + } + } + + public static WorkloadGroupFeatureType getInstance() { + if (instance == null) { + throw new IllegalStateException("FeatureValueValidator is not initialized. Call initializeFeatureValueValidator() first."); + } + return instance; + } + @Override public String getName() { return NAME; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java index 0ea7621943615..1b673b656f9bb 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java @@ -20,16 +20,28 @@ */ public class WorkloadGroupFeatureValueValidator implements FeatureValueValidator { private final ClusterService clusterService; + private static volatile WorkloadGroupFeatureValueValidator instance; private final Logger logger = LogManager.getLogger(WorkloadGroupFeatureValueValidator.class); /** * constructor for WorkloadGroupFeatureValueValidator * @param clusterService */ - public WorkloadGroupFeatureValueValidator(ClusterService clusterService) { + private WorkloadGroupFeatureValueValidator(ClusterService clusterService) { this.clusterService = clusterService; } + public static WorkloadGroupFeatureValueValidator getInstance(ClusterService clusterService) { + if (instance == null) { + synchronized (WorkloadGroupFeatureValueValidator.class) { + if (instance == null) { + instance = new WorkloadGroupFeatureValueValidator(clusterService); + } + } + } + return instance; + } + @Override public void validate(String featureValue) { if (!clusterService.state().metadata().workloadGroups().containsKey(featureValue)) { diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java index f38985e37c426..98e28ee53dab3 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java @@ -8,15 +8,27 @@ package org.opensearch.plugin.wlm.rule; +import org.junit.Before; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.rule.RuleAttribute; import org.opensearch.rule.autotagging.Attribute; import org.opensearch.rule.autotagging.AutoTaggingRegistry; +import org.opensearch.rule.autotagging.FeatureValueValidator; import org.opensearch.test.OpenSearchTestCase; import java.util.Map; +import static org.mockito.Mockito.mock; + public class WorkloadGroupFeatureTypeTests extends OpenSearchTestCase { - WorkloadGroupFeatureType featureType = new WorkloadGroupFeatureType(null); + WorkloadGroupFeatureType featureType; + + @Before + public void setUpFeatureType() { + FeatureValueValidator validator = WorkloadGroupFeatureValueValidator.getInstance(mock(ClusterService.class)); + WorkloadGroupFeatureType.initializeFeatureValueValidator(validator); + featureType = WorkloadGroupFeatureType.getInstance(); + } public void testGetName_returnsCorrectName() { assertEquals("workload_group", featureType.getName()); From 8dfab4c9592d3c5730d82cea76c984d565812f6d Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Mon, 12 May 2025 01:22:21 -0700 Subject: [PATCH 17/20] addrss comments Signed-off-by: Ruirui Zhang --- .../java/org/opensearch/rule/CreateRuleRequest.java | 1 - .../service/IndexStoredRulePersistenceService.java | 13 ++++++++++--- .../rule/autotagging/AutoTaggingRegistryTests.java | 12 ++++++------ .../org/opensearch/rule/utils/RuleTestUtils.java | 1 - .../opensearch/rule/rest/RestCreateRuleAction.java | 1 - .../rule/InMemoryRuleProcessingServiceTests.java | 9 ++++----- .../plugin/wlm/rule/WorkloadGroupFeatureType.java | 8 ++++++++ .../rule/WorkloadGroupFeatureValueValidator.java | 5 +++++ .../wlm/rule/WorkloadGroupFeatureTypeTests.java | 2 +- 9 files changed, 34 insertions(+), 18 deletions(-) diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java index 2a1aa0c980906..69434947f5660 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java @@ -12,7 +12,6 @@ import org.opensearch.action.ActionRequestValidationException; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; -import org.opensearch.rule.autotagging.FeatureValueValidator; import org.opensearch.rule.autotagging.Rule; import java.io.IOException; diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java index 4e1f02f9c01ba..332c91aa276f2 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -67,6 +67,7 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService private final int maxRulesPerPage; private final RuleEntityParser parser; private final RuleQueryMapper queryBuilder; + private final Object lock = new Object(); private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class); private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); @@ -107,7 +108,9 @@ public void createRule(CreateRuleRequest request, ActionListener AutoTaggingRegistry.registerFeatureType(null)); FeatureType featureType = mock(FeatureType.class); - when(featureType.getName()).thenReturn(TEST_FEATURE_TYPE); + when(featureType.getName()).thenReturn(FEATURE_TYPE_NAME); assertThrows(IllegalStateException.class, () -> AutoTaggingRegistry.registerFeatureType(featureType)); when(featureType.getName()).thenReturn(randomAlphaOfLength(MAX_FEATURE_TYPE_NAME_LENGTH + 1)); assertThrows(IllegalStateException.class, () -> AutoTaggingRegistry.registerFeatureType(featureType)); diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java index 87ee679031916..caa4aafdfca51 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java @@ -11,7 +11,6 @@ import org.opensearch.rule.autotagging.Attribute; import org.opensearch.rule.autotagging.AutoTaggingRegistry; import org.opensearch.rule.autotagging.FeatureType; -import org.opensearch.rule.autotagging.FeatureValueValidator; import java.util.Map; import java.util.Set; diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java index 7068879e7ba54..4a8254924db2e 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java @@ -29,7 +29,6 @@ import java.io.IOException; import java.util.List; -import static org.opensearch.rest.RestRequest.Method.POST; import static org.opensearch.rest.RestRequest.Method.PUT; import static org.opensearch.rule.rest.RestGetRuleAction.FEATURE_TYPE; diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java index 4e024f3600c35..6b78a5847d9dc 100644 --- a/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java @@ -12,7 +12,6 @@ import org.opensearch.rule.autotagging.Attribute; import org.opensearch.rule.autotagging.AutoTaggingRegistry; import org.opensearch.rule.autotagging.FeatureType; -import org.opensearch.rule.autotagging.FeatureValueValidator; import org.opensearch.rule.autotagging.Rule; import org.opensearch.rule.storage.DefaultAttributeValueStore; import org.opensearch.test.OpenSearchTestCase; @@ -138,10 +137,10 @@ public Map getAllowedAttributesRegistry() { return Map.of("test_attribute", TestAttribute.TEST_ATTRIBUTE); } -// @Override -// public FeatureValueValidator getFeatureValueValidator() { -// return null; -// } + // @Override + // public FeatureValueValidator getFeatureValueValidator() { + // return null; + // } } public enum TestAttribute implements Attribute { diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java index f934ae15cda27..059e01a5d5e37 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java @@ -41,12 +41,20 @@ private WorkloadGroupFeatureType(FeatureValueValidator featureValueValidator) { this.featureValueValidator = featureValueValidator; } + /** + * Initializes the singleton instance of WorkloadGroupFeatureType. + * This method should be called once before calling {@link #getInstance()}. + * @param validator the FeatureValueValidator to be used for initialization + */ public static void initializeFeatureValueValidator(FeatureValueValidator validator) { if (instance == null) { instance = new WorkloadGroupFeatureType(validator); } } + /** + * Returns the singleton instance of {@link WorkloadGroupFeatureType}. + */ public static WorkloadGroupFeatureType getInstance() { if (instance == null) { throw new IllegalStateException("FeatureValueValidator is not initialized. Call initializeFeatureValueValidator() first."); diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java index 1b673b656f9bb..cdd8a5e432438 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java @@ -31,6 +31,11 @@ private WorkloadGroupFeatureValueValidator(ClusterService clusterService) { this.clusterService = clusterService; } + /** + * Returns the singleton instance of {@code WorkloadGroupFeatureValueValidator}, initializing it if necessary. + * Uses double-checked locking to ensure thread-safe lazy initialization. + * @param clusterService the {@link ClusterService} used to construct the validator if not already initialized + */ public static WorkloadGroupFeatureValueValidator getInstance(ClusterService clusterService) { if (instance == null) { synchronized (WorkloadGroupFeatureValueValidator.class) { diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java index 98e28ee53dab3..da7a5d907d5ef 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java @@ -8,13 +8,13 @@ package org.opensearch.plugin.wlm.rule; -import org.junit.Before; import org.opensearch.cluster.service.ClusterService; import org.opensearch.rule.RuleAttribute; import org.opensearch.rule.autotagging.Attribute; import org.opensearch.rule.autotagging.AutoTaggingRegistry; import org.opensearch.rule.autotagging.FeatureValueValidator; import org.opensearch.test.OpenSearchTestCase; +import org.junit.Before; import java.util.Map; From 74beb60d0787ab799c6283912e72fd11230df7b5 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Wed, 14 May 2025 01:29:48 -0700 Subject: [PATCH 18/20] change create rule api to run only on master node to avoid race condition Signed-off-by: Ruirui Zhang --- .../opensearch/rule/CreateRuleRequest.java | 4 +- .../autotagging/AutoTaggingRegistryTests.java | 2 +- .../action/TransportCreateRuleAction.java | 50 ++++++++++++++-- .../InMemoryRuleProcessingServiceTests.java | 5 -- .../rule/action/CreateRuleActionTests.java | 22 +++++++ .../TransportCreateRuleActionTests.java | 58 +++++++++++++++++++ .../plugin/wlm/WorkloadManagementPlugin.java | 12 ++-- .../wlm/rule/WorkloadGroupFeatureType.java | 24 +------- .../WorkloadGroupFeatureValueValidator.java | 19 +----- .../rule/WorkloadGroupFeatureTypeTests.java | 14 +---- 10 files changed, 138 insertions(+), 72 deletions(-) create mode 100644 modules/autotagging-commons/src/test/java/org/opensearch/rule/action/CreateRuleActionTests.java create mode 100644 modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java index 69434947f5660..d6db4e22d6bd9 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java @@ -8,8 +8,8 @@ package org.opensearch.rule; -import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.rule.autotagging.Rule; @@ -28,7 +28,7 @@ * }' * @opensearch.experimental */ -public class CreateRuleRequest extends ActionRequest { +public class CreateRuleRequest extends ClusterManagerNodeRequest { private final Rule rule; /** diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/AutoTaggingRegistryTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/AutoTaggingRegistryTests.java index bec29c7f8e4d4..56adf333e1a1c 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/AutoTaggingRegistryTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/AutoTaggingRegistryTests.java @@ -15,9 +15,9 @@ import static org.opensearch.rule.autotagging.AutoTaggingRegistry.MAX_FEATURE_TYPE_NAME_LENGTH; import static org.opensearch.rule.autotagging.RuleTests.INVALID_FEATURE; +import static org.opensearch.rule.utils.RuleTestUtils.FEATURE_TYPE_NAME; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import static org.opensearch.rule.utils.RuleTestUtils.FEATURE_TYPE_NAME; public class AutoTaggingRegistryTests extends OpenSearchTestCase { diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java index 28dd658d107aa..8fc51176bf173 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java @@ -9,46 +9,86 @@ package org.opensearch.rule.action; import org.opensearch.action.support.ActionFilters; -import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction; +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.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.rule.CreateRuleRequest; import org.opensearch.rule.CreateRuleResponse; import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.RulePersistenceServiceRegistry; -import org.opensearch.tasks.Task; +import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; +import java.io.IOException; + +import static org.opensearch.threadpool.ThreadPool.Names.SAME; + /** * Transport action to create Rules * @opensearch.experimental */ -public class TransportCreateRuleAction extends HandledTransportAction { +public class TransportCreateRuleAction extends TransportClusterManagerNodeAction { private final RulePersistenceServiceRegistry rulePersistenceServiceRegistry; /** * Constructor for TransportCreateRuleAction * + * @param threadPool - {@link ThreadPool} object * @param transportService - a {@link TransportService} object + * @param clusterService - a {@link ClusterService} object * @param actionFilters - a {@link ActionFilters} object + * @param indexNameExpressionResolver - {@link IndexNameExpressionResolver} object * @param rulePersistenceServiceRegistry - a {@link RulePersistenceServiceRegistry} object */ @Inject public TransportCreateRuleAction( + ThreadPool threadPool, TransportService transportService, + ClusterService clusterService, ActionFilters actionFilters, + IndexNameExpressionResolver indexNameExpressionResolver, RulePersistenceServiceRegistry rulePersistenceServiceRegistry ) { - super(CreateRuleAction.NAME, transportService, actionFilters, CreateRuleRequest::new); + super( + CreateRuleAction.NAME, + transportService, + clusterService, + threadPool, + actionFilters, + CreateRuleRequest::new, + indexNameExpressionResolver + ); this.rulePersistenceServiceRegistry = rulePersistenceServiceRegistry; } @Override - protected void doExecute(Task task, CreateRuleRequest request, ActionListener listener) { + protected String executor() { + return SAME; + } + + @Override + protected CreateRuleResponse read(StreamInput in) throws IOException { + return new CreateRuleResponse(in); + } + + @Override + protected void clusterManagerOperation(CreateRuleRequest request, ClusterState state, ActionListener listener) + throws Exception { final RulePersistenceService rulePersistenceService = rulePersistenceServiceRegistry.getRulePersistenceService( request.getRule().getFeatureType() ); rulePersistenceService.createRule(request, listener); } + + @Override + protected ClusterBlockException checkBlock(CreateRuleRequest request, ClusterState state) { + return state.blocks().globalBlockedException(ClusterBlockLevel.WRITE); + } } diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java index 6b78a5847d9dc..dfe347007959a 100644 --- a/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/InMemoryRuleProcessingServiceTests.java @@ -136,11 +136,6 @@ public String getName() { public Map getAllowedAttributesRegistry() { return Map.of("test_attribute", TestAttribute.TEST_ATTRIBUTE); } - - // @Override - // public FeatureValueValidator getFeatureValueValidator() { - // return null; - // } } public enum TestAttribute implements Attribute { diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/CreateRuleActionTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/CreateRuleActionTests.java new file mode 100644 index 0000000000000..5941cd3e9a6a7 --- /dev/null +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/CreateRuleActionTests.java @@ -0,0 +1,22 @@ +/* + * 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.rule.action; + +import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.test.OpenSearchTestCase; + +public class CreateRuleActionTests extends OpenSearchTestCase { + public void testGetName() { + assertEquals("cluster:admin/opensearch/rule/_create", CreateRuleAction.NAME); + } + + public void testCreateResponseReader() { + assertTrue(CreateRuleAction.INSTANCE.getResponseReader() instanceof Writeable.Reader); + } +} diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java new file mode 100644 index 0000000000000..1574882825ffb --- /dev/null +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java @@ -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.rule.action; + +import org.opensearch.action.support.ActionFilters; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.rule.CreateRuleRequest; +import org.opensearch.rule.RulePersistenceService; +import org.opensearch.rule.RulePersistenceServiceRegistry; +import org.opensearch.rule.autotagging.Rule; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class TransportCreateRuleActionTests extends OpenSearchTestCase { + TransportCreateRuleAction sut; + + public void testExecute() throws Exception { + RulePersistenceServiceRegistry rulePersistenceServiceRegistry = mock(RulePersistenceServiceRegistry.class); + TransportService transportService = mock(TransportService.class); + ActionFilters actionFilters = mock(ActionFilters.class); + RulePersistenceService rulePersistenceService = mock(RulePersistenceService.class); + CreateRuleRequest createRuleRequest = mock(CreateRuleRequest.class); + Rule rule = mock(Rule.class); + when(createRuleRequest.getRule()).thenReturn(rule); + when(createRuleRequest.getRule().getFeatureType()).thenReturn(null); + ThreadPool threadPool = mock(ThreadPool.class); + ClusterService clusterService = mock(ClusterService.class); + IndexNameExpressionResolver indexNameExpressionResolver = mock(IndexNameExpressionResolver.class); + + when(rulePersistenceServiceRegistry.getRulePersistenceService(any())).thenReturn(rulePersistenceService); + doNothing().when(rulePersistenceService).getRule(any(), any()); + sut = new TransportCreateRuleAction( + threadPool, + transportService, + clusterService, + actionFilters, + indexNameExpressionResolver, + rulePersistenceServiceRegistry + ); + sut.clusterManagerOperation(createRuleRequest, null, null); + verify(rulePersistenceService, times(1)).createRule(any(), any()); + } +} diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index eb40e464f1da6..0cadc579ba9a4 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -49,7 +49,6 @@ import org.opensearch.rule.InMemoryRuleProcessingService; import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.autotagging.FeatureType; -import org.opensearch.rule.autotagging.FeatureValueValidator; import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.rule.spi.RuleFrameworkExtension; import org.opensearch.rule.storage.DefaultAttributeValueStore; @@ -102,14 +101,13 @@ public Collection createComponents( IndexNameExpressionResolver indexNameExpressionResolver, Supplier repositoriesServiceSupplier ) { - FeatureValueValidator validator = WorkloadGroupFeatureValueValidator.getInstance(clusterService); - WorkloadGroupFeatureType.initializeFeatureValueValidator(validator); + FeatureTypeHolder.featureType = new WorkloadGroupFeatureType(new WorkloadGroupFeatureValueValidator(clusterService)); RulePersistenceServiceHolder.rulePersistenceService = new IndexStoredRulePersistenceService( INDEX_NAME, client, clusterService, MAX_RULES_PER_PAGE, - new XContentRuleParser(WorkloadGroupFeatureType.getInstance()), + new XContentRuleParser(FeatureTypeHolder.featureType), new IndexBasedRuleQueryMapper() ); InMemoryRuleProcessingService ruleProcessingService = new InMemoryRuleProcessingService( @@ -175,10 +173,14 @@ public Supplier getRulePersistenceServiceSupplier() { @Override public FeatureType getFeatureType() { - return WorkloadGroupFeatureType.getInstance(); + return FeatureTypeHolder.featureType; } static class RulePersistenceServiceHolder { private static RulePersistenceService rulePersistenceService; } + + static class FeatureTypeHolder { + private static FeatureType featureType; + } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java index 059e01a5d5e37..fc9dfa3136277 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureType.java @@ -31,37 +31,15 @@ public class WorkloadGroupFeatureType implements FeatureType { RuleAttribute.INDEX_PATTERN ); private final FeatureValueValidator featureValueValidator; - private static WorkloadGroupFeatureType instance; /** * constructor for WorkloadGroupFeatureType * @param featureValueValidator */ - private WorkloadGroupFeatureType(FeatureValueValidator featureValueValidator) { + public WorkloadGroupFeatureType(FeatureValueValidator featureValueValidator) { this.featureValueValidator = featureValueValidator; } - /** - * Initializes the singleton instance of WorkloadGroupFeatureType. - * This method should be called once before calling {@link #getInstance()}. - * @param validator the FeatureValueValidator to be used for initialization - */ - public static void initializeFeatureValueValidator(FeatureValueValidator validator) { - if (instance == null) { - instance = new WorkloadGroupFeatureType(validator); - } - } - - /** - * Returns the singleton instance of {@link WorkloadGroupFeatureType}. - */ - public static WorkloadGroupFeatureType getInstance() { - if (instance == null) { - throw new IllegalStateException("FeatureValueValidator is not initialized. Call initializeFeatureValueValidator() first."); - } - return instance; - } - @Override public String getName() { return NAME; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java index cdd8a5e432438..0ea7621943615 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureValueValidator.java @@ -20,33 +20,16 @@ */ public class WorkloadGroupFeatureValueValidator implements FeatureValueValidator { private final ClusterService clusterService; - private static volatile WorkloadGroupFeatureValueValidator instance; private final Logger logger = LogManager.getLogger(WorkloadGroupFeatureValueValidator.class); /** * constructor for WorkloadGroupFeatureValueValidator * @param clusterService */ - private WorkloadGroupFeatureValueValidator(ClusterService clusterService) { + public WorkloadGroupFeatureValueValidator(ClusterService clusterService) { this.clusterService = clusterService; } - /** - * Returns the singleton instance of {@code WorkloadGroupFeatureValueValidator}, initializing it if necessary. - * Uses double-checked locking to ensure thread-safe lazy initialization. - * @param clusterService the {@link ClusterService} used to construct the validator if not already initialized - */ - public static WorkloadGroupFeatureValueValidator getInstance(ClusterService clusterService) { - if (instance == null) { - synchronized (WorkloadGroupFeatureValueValidator.class) { - if (instance == null) { - instance = new WorkloadGroupFeatureValueValidator(clusterService); - } - } - } - return instance; - } - @Override public void validate(String featureValue) { if (!clusterService.state().metadata().workloadGroups().containsKey(featureValue)) { diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java index da7a5d907d5ef..201dad40cc439 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java @@ -8,27 +8,15 @@ package org.opensearch.plugin.wlm.rule; -import org.opensearch.cluster.service.ClusterService; import org.opensearch.rule.RuleAttribute; import org.opensearch.rule.autotagging.Attribute; import org.opensearch.rule.autotagging.AutoTaggingRegistry; -import org.opensearch.rule.autotagging.FeatureValueValidator; import org.opensearch.test.OpenSearchTestCase; -import org.junit.Before; import java.util.Map; -import static org.mockito.Mockito.mock; - public class WorkloadGroupFeatureTypeTests extends OpenSearchTestCase { - WorkloadGroupFeatureType featureType; - - @Before - public void setUpFeatureType() { - FeatureValueValidator validator = WorkloadGroupFeatureValueValidator.getInstance(mock(ClusterService.class)); - WorkloadGroupFeatureType.initializeFeatureValueValidator(validator); - featureType = WorkloadGroupFeatureType.getInstance(); - } + WorkloadGroupFeatureType featureType = new WorkloadGroupFeatureType(new WorkloadGroupFeatureValueValidator(null)); public void testGetName_returnsCorrectName() { assertEquals("workload_group", featureType.getName()); From 656e01cd64ffe203152f3586869bbf0264484f2a Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Fri, 23 May 2025 14:02:39 -0700 Subject: [PATCH 19/20] customize transport layer to route to primary shard's node Signed-off-by: Ruirui Zhang --- .../opensearch/rule/CreateRuleRequest.java | 11 +- .../org/opensearch/rule/GetRuleRequest.java | 7 +- .../IndexStoredRulePersistenceService.java | 83 ++------ ...ndexStoredRulePersistenceServiceTests.java | 122 +++++++++++ .../rule/spi/RuleFrameworkExtension.java | 7 +- .../opensearch/rule/RuleFrameworkPlugin.java | 38 +++- .../action/TransportCreateRuleAction.java | 191 +++++++++++++++--- .../rule/RuleFrameworkPluginTests.java | 2 +- .../RulePersistenceServiceRegistryTests.java | 4 +- .../TransportCreateRuleActionTests.java | 82 +++++--- .../plugin/wlm/WorkloadManagementPlugin.java | 24 +-- .../wlm/WorkloadManagementPluginTests.java | 2 +- 12 files changed, 401 insertions(+), 172 deletions(-) diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java index d6db4e22d6bd9..7963357b893d9 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/CreateRuleRequest.java @@ -8,8 +8,8 @@ package org.opensearch.rule; +import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionRequestValidationException; -import org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.rule.autotagging.Rule; @@ -19,16 +19,15 @@ /** * A request for create Rule * Example request: - * Note that the endpoint below is for wlm rules specifically and serves only as an example - * curl -XPUT "localhost:9200/_wlm/rule/" -H 'Content-Type: application/json' -d ' + * curl -X PUT "localhost:9200/_rules/{featureType}/" -H 'Content-Type: application/json' -d ' * { * "description": "description1", - * "index_pattern": ["log*", "event*"], - * "workload_group": "poOiU851RwyLYvV5lbvv5w" + * "attribute_name": ["log*", "event*"], + * "feature_type": "poOiU851RwyLYvV5lbvv5w" * }' * @opensearch.experimental */ -public class CreateRuleRequest extends ClusterManagerNodeRequest { +public class CreateRuleRequest extends ActionRequest { private final Rule rule; /** diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/GetRuleRequest.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/GetRuleRequest.java index 9cdebb782edc9..630a329688b2b 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/GetRuleRequest.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/GetRuleRequest.java @@ -26,10 +26,9 @@ /** * A request for get Rule * Example Request: - * The endpoint "localhost:9200/_wlm/rule" is specific to the Workload Management feature to manage rules - * curl -X GET "localhost:9200/_wlm/rule" - get all rules - * curl -X GET "localhost:9200/_wlm/rule/{_id}" - get single rule by id - * curl -X GET "localhost:9200/_wlm/rule?index_pattern=a,b" - get all rules containing attribute index_pattern as a or b + * curl -X GET "localhost:9200/_rules/{featureType}/" - get all rules for {featureType} + * curl -X GET "localhost:9200/_rules/{featureType}/{_id}" - get single rule by id + * curl -X GET "localhost:9200/_rules/{featureType}?index_pattern=a,b" - get all rules containing attribute index_pattern as a or b for {featureType} * @opensearch.experimental */ @ExperimentalApi diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java index 332c91aa276f2..2b1e523c2b198 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -10,29 +10,22 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.opensearch.ExceptionsHelper; -import org.opensearch.ResourceAlreadyExistsException; import org.opensearch.ResourceNotFoundException; import org.opensearch.action.DocWriteResponse; import org.opensearch.action.delete.DeleteRequest; -import org.opensearch.action.search.SearchRequestBuilder; -import org.opensearch.action.support.clustermanager.AcknowledgedResponse; -import org.opensearch.action.admin.indices.create.CreateIndexRequest; -import org.opensearch.action.admin.indices.create.CreateIndexResponse; import org.opensearch.action.index.IndexRequest; import org.opensearch.action.search.SearchRequestBuilder; +import org.opensearch.action.support.clustermanager.AcknowledgedResponse; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.action.ActionListener; -import org.opensearch.index.engine.DocumentMissingException; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.index.query.QueryBuilders; -import org.opensearch.rule.DeleteRuleRequest; import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.index.engine.DocumentMissingException; import org.opensearch.index.query.QueryBuilder; import org.opensearch.rule.CreateRuleRequest; import org.opensearch.rule.CreateRuleResponse; +import org.opensearch.rule.DeleteRuleRequest; import org.opensearch.rule.GetRuleRequest; import org.opensearch.rule.GetRuleResponse; import org.opensearch.rule.RuleEntityParser; @@ -67,9 +60,7 @@ public class IndexStoredRulePersistenceService implements RulePersistenceService private final int maxRulesPerPage; private final RuleEntityParser parser; private final RuleQueryMapper queryBuilder; - private final Object lock = new Object(); private static final Logger logger = LogManager.getLogger(IndexStoredRulePersistenceService.class); - private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); /** * Constructs an instance of {@link IndexStoredRulePersistenceService} with the specified parameters. @@ -106,66 +97,21 @@ public IndexStoredRulePersistenceService( public void createRule(CreateRuleRequest request, ActionListener listener) { try (ThreadContext.StoredContext ctx = getContext()) { if (!clusterService.state().metadata().hasIndex(indexName)) { - createIndex(listener, request.getRule()); + logger.error("Index {} does not exist", indexName); + throw new IllegalStateException("Index" + indexName + " does not exist"); } else { - synchronized (lock) { - validateAndPersist(request.getRule(), listener); - } + Rule rule = request.getRule(); + validateNoDuplicateRule(rule, ActionListener.wrap(unused -> persistRule(rule, listener), listener::onFailure)); } } } - /** - * Creates the system index, then validates and persists the given rule. - * @param listener - ActionListener for CreateRuleResponse - * @param rule - the rule to validate and persist - */ - private void createIndex(ActionListener listener, Rule rule) { - final CreateIndexRequest request = new CreateIndexRequest(indexName).settings(indexSettings); - client.admin().indices().create(request, new ActionListener<>() { - @Override - public void onResponse(CreateIndexResponse response) { - if (!response.isAcknowledged()) { - logger.error("Index creation not acknowledged: {}", indexName); - listener.onFailure(new IllegalStateException(indexName + " index creation failed and rule cannot be persisted")); - } else { - synchronized (lock) { - validateAndPersist(rule, listener); - } - } - } - - @Override - public void onFailure(Exception e) { - Throwable cause = ExceptionsHelper.unwrapCause(e); - if (cause instanceof ResourceAlreadyExistsException) { - synchronized (lock) { - validateAndPersist(rule, listener); - } - } else { - logger.error("Failed to create index {}: {}", indexName, e.getMessage()); - listener.onFailure(e); - } - } - }); - } - - /** - * Validates that the rule does not already exist, then persists it if validation succeeds - * @param rule - The rule to validate and persist - * @param listener - ActionListener for CreateRuleResponse - */ - private void validateAndPersist(Rule rule, ActionListener listener) { - validateNoDuplicateRule(rule, ActionListener.wrap(unused -> persistRule(rule, listener), listener::onFailure)); - } - /** * Validates that no duplicate rule exists with the same attribute map. * If a conflict is found, fails the listener * @param rule - the rule we check duplicate against * @param listener - listener for validateNoDuplicateRule response */ - private void validateNoDuplicateRule(Rule rule, ActionListener listener) { try (ThreadContext.StoredContext ctx = getContext()) { QueryBuilder query = queryBuilder.from(new GetRuleRequest(null, rule.getAttributeMap(), null, rule.getFeatureType())); @@ -266,10 +212,6 @@ void handleGetRuleResponse(List hits, ActionListener listener.onResponse(new GetRuleResponse(ruleMap, nextSearchAfter)); } - private ThreadContext.StoredContext getContext() { - return client.threadPool().getThreadContext().stashContext(); - } - @Override public void deleteRule(DeleteRuleRequest request, ActionListener listener) { try (ThreadContext.StoredContext context = getContext()) { @@ -291,4 +233,15 @@ public void deleteRule(DeleteRuleRequest request, ActionListener mockRuleQueryMapper = mock(RuleQueryMapper.class); + RuleEntityParser mockRuleEntityParser = mock(RuleEntityParser.class); + ClusterService clusterService = mock(ClusterService.class); + ClusterState clusterState = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + QueryBuilder queryBuilder = mock(QueryBuilder.class); + + when(clusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + when(metadata.hasIndex(TEST_INDEX_NAME)).thenReturn(true); + when(mockRuleQueryMapper.from(any(GetRuleRequest.class))).thenReturn(queryBuilder); + when(queryBuilder.filter(any())).thenReturn(queryBuilder); + + SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + Client client = setUpMockClient(searchRequestBuilder); + RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( + TEST_INDEX_NAME, + client, + clusterService, + MAX_VALUES_PER_PAGE, + mockRuleEntityParser, + mockRuleQueryMapper + ); + ActionListener listener = mock(ActionListener.class); + when(mockRule.toXContent(any(), any())).thenAnswer(invocation -> invocation.getArgument(0)); + + SearchResponse searchResponse = mock(SearchResponse.class); + when(searchResponse.getHits()).thenReturn(new SearchHits(new SearchHit[] {}, new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1.0f)); + doAnswer((invocation) -> { + ActionListener actionListener = invocation.getArgument(0); + actionListener.onResponse(searchResponse); + return null; + }).when(searchRequestBuilder).execute(any(ActionListener.class)); + + IndexResponse indexResponse = mock(IndexResponse.class); + when(indexResponse.getId()).thenReturn(_ID_ONE); + doAnswer(invocation -> { + ActionListener actionListener = invocation.getArgument(1); + actionListener.onResponse(indexResponse); + return null; + }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); + + rulePersistenceService.createRule(createRuleRequest, listener); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(CreateRuleResponse.class); + verify(listener).onResponse(responseCaptor.capture()); + CreateRuleResponse response = responseCaptor.getValue(); + assertNotNull(response.getRule()); + } + + public void testCreateDuplicateRule() { + CreateRuleRequest createRuleRequest = mock(CreateRuleRequest.class); + Rule mockRule = mock(Rule.class); + when(createRuleRequest.getRule()).thenReturn(mockRule); + + RuleQueryMapper mockRuleQueryMapper = mock(RuleQueryMapper.class); + RuleEntityParser mockRuleEntityParser = mock(RuleEntityParser.class); + ClusterService clusterService = mock(ClusterService.class); + ClusterState clusterState = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + QueryBuilder queryBuilder = mock(QueryBuilder.class); + + when(clusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + when(metadata.hasIndex(TEST_INDEX_NAME)).thenReturn(true); + when(mockRuleQueryMapper.from(any(GetRuleRequest.class))).thenReturn(queryBuilder); + when(queryBuilder.filter(any())).thenReturn(queryBuilder); + + when(mockRule.getAttributeMap()).thenReturn(Map.of(MOCK_RULE_ATTRIBUTE_ONE, Set.of(ATTRIBUTE_VALUE_ONE))); + when(mockRule.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); + + SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); + Client client = setUpMockClient(searchRequestBuilder); + + RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( + TEST_INDEX_NAME, + client, + clusterService, + MAX_VALUES_PER_PAGE, + mockRuleEntityParser, + mockRuleQueryMapper + ); + + SearchResponse searchResponse = mock(SearchResponse.class); + SearchHit hit = new SearchHit(1); + hit.sourceRef(new BytesArray(VALID_JSON)); + SearchHits searchHits = new SearchHits(new SearchHit[] { hit }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1.0f); + when(searchResponse.getHits()).thenReturn(searchHits); + + doAnswer((invocation) -> { + ActionListener actionListener = invocation.getArgument(0); + actionListener.onResponse(searchResponse); + return null; + }).when(searchRequestBuilder).execute(any(ActionListener.class)); + + ActionListener listener = mock(ActionListener.class); + when(mockRuleEntityParser.parse(any(String.class))).thenReturn(mockRule); + rulePersistenceService.createRule(createRuleRequest, listener); + ArgumentCaptor failureCaptor = ArgumentCaptor.forClass(Exception.class); + verify(listener).onFailure(failureCaptor.capture()); + } + public void testGetRuleByIdSuccess() { GetRuleRequest getRuleRequest = mock(GetRuleRequest.class); when(getRuleRequest.getId()).thenReturn(_ID_ONE); @@ -76,10 +190,12 @@ public void testGetRuleByIdSuccess() { SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); Client client = setUpMockClient(searchRequestBuilder); + ClusterService clusterService = mock(ClusterService.class); RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( TEST_INDEX_NAME, client, + clusterService, MAX_VALUES_PER_PAGE, mockRuleEntityParser, mockRuleQueryMapper @@ -122,10 +238,12 @@ public void testGetRuleByIdNotFound() { SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); Client client = setUpMockClient(searchRequestBuilder); + ClusterService clusterService = mock(ClusterService.class); RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( TEST_INDEX_NAME, client, + clusterService, MAX_VALUES_PER_PAGE, mockRuleEntityParser, mockRuleQueryMapper @@ -176,6 +294,7 @@ public void testDeleteRule_successful() { DeleteRuleRequest request = new DeleteRuleRequest(ruleId, RuleTestUtils.MockRuleFeatureType.INSTANCE); Client client = mock(Client.class); + ClusterService clusterService = mock(ClusterService.class); ThreadPool threadPool = mock(ThreadPool.class); when(client.threadPool()).thenReturn(threadPool); when(threadPool.getThreadContext()).thenReturn(new ThreadContext(Settings.EMPTY)); @@ -183,6 +302,7 @@ public void testDeleteRule_successful() { RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( TEST_INDEX_NAME, client, + clusterService, MAX_VALUES_PER_PAGE, mock(RuleEntityParser.class), mock(RuleQueryMapper.class) @@ -214,6 +334,7 @@ public void testDeleteRule_notFound() { DeleteRuleRequest request = new DeleteRuleRequest(ruleId, RuleTestUtils.MockRuleFeatureType.INSTANCE); Client client = mock(Client.class); + ClusterService clusterService = mock(ClusterService.class); ThreadPool threadPool = mock(ThreadPool.class); when(client.threadPool()).thenReturn(threadPool); when(threadPool.getThreadContext()).thenReturn(new ThreadContext(Settings.EMPTY)); @@ -221,6 +342,7 @@ public void testDeleteRule_notFound() { RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( TEST_INDEX_NAME, client, + clusterService, MAX_VALUES_PER_PAGE, mock(RuleEntityParser.class), mock(RuleQueryMapper.class) diff --git a/modules/autotagging-commons/spi/src/main/java/org/opensearch/rule/spi/RuleFrameworkExtension.java b/modules/autotagging-commons/spi/src/main/java/org/opensearch/rule/spi/RuleFrameworkExtension.java index 6197fa9e03093..07443c31236a7 100644 --- a/modules/autotagging-commons/spi/src/main/java/org/opensearch/rule/spi/RuleFrameworkExtension.java +++ b/modules/autotagging-commons/spi/src/main/java/org/opensearch/rule/spi/RuleFrameworkExtension.java @@ -25,9 +25,8 @@ public interface RuleFrameworkExtension { Supplier getRulePersistenceServiceSupplier(); /** - * It tells the framework its FeatureType which can be used by Transport classes to handle the - * consumer specific persistence - * @return + * Flow implementation from consumer plugins into framework plugin + * @return the specific implementation of FeatureType */ - FeatureType getFeatureType(); + Supplier getFeatureTypeSupplier(); } diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java index e97d13b35308c..fbbe6e2759693 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java @@ -22,20 +22,20 @@ import org.opensearch.plugins.Plugin; import org.opensearch.rest.RestController; import org.opensearch.rest.RestHandler; -import org.opensearch.rule.action.DeleteRuleAction; -import org.opensearch.rule.action.GetRuleAction; -import org.opensearch.rule.action.TransportDeleteRuleAction; -import org.opensearch.rule.action.TransportGetRuleAction; -import org.opensearch.rule.autotagging.AutoTaggingRegistry; -import org.opensearch.rule.rest.RestDeleteRuleAction; import org.opensearch.rule.action.CreateRuleAction; +import org.opensearch.rule.action.DeleteRuleAction; import org.opensearch.rule.action.GetRuleAction; import org.opensearch.rule.action.TransportCreateRuleAction; +import org.opensearch.rule.action.TransportDeleteRuleAction; import org.opensearch.rule.action.TransportGetRuleAction; import org.opensearch.rule.autotagging.AutoTaggingRegistry; +import org.opensearch.rule.autotagging.FeatureType; import org.opensearch.rule.rest.RestCreateRuleAction; +import org.opensearch.rule.rest.RestDeleteRuleAction; import org.opensearch.rule.rest.RestGetRuleAction; import org.opensearch.rule.spi.RuleFrameworkExtension; +import org.opensearch.threadpool.ExecutorBuilder; +import org.opensearch.threadpool.FixedExecutorBuilder; import java.util.ArrayList; import java.util.Collection; @@ -49,6 +49,19 @@ */ public class RuleFrameworkPlugin extends Plugin implements ExtensiblePlugin, ActionPlugin { + /** + * The name of the thread pool dedicated to rule execution. + */ + public static final String RULE_THREAD_POOL_NAME = "rule_serial_executor"; + /** + * The number of threads allocated in the rule execution thread pool. This is set to 1 to ensure serial execution. + */ + public static final int RULE_THREAD_COUNT = 1; + /** + * The maximum number of tasks that can be queued in the rule execution thread pool. + */ + public static final int RULE_QUEUE_SIZE = 100; + /** * constructor for RuleFrameworkPlugin */ @@ -81,6 +94,11 @@ public List getRestHandlers( return List.of(new RestGetRuleAction(), new RestDeleteRuleAction(), new RestCreateRuleAction()); } + @Override + public List> getExecutorBuilders(Settings settings) { + return List.of(new FixedExecutorBuilder(settings, RULE_THREAD_POOL_NAME, RULE_THREAD_COUNT, RULE_QUEUE_SIZE, "rule-threadpool")); + } + @Override public Collection createGuiceModules() { return List.of(b -> { b.bind(RulePersistenceServiceRegistry.class).toInstance(rulePersistenceServiceRegistry); }); @@ -92,10 +110,8 @@ public void loadExtensions(ExtensionLoader loader) { } private void consumeFrameworkExtension(RuleFrameworkExtension ruleFrameworkExtension) { - AutoTaggingRegistry.registerFeatureType(ruleFrameworkExtension.getFeatureType()); - rulePersistenceServiceRegistry.register( - ruleFrameworkExtension.getFeatureType(), - ruleFrameworkExtension.getRulePersistenceServiceSupplier().get() - ); + FeatureType featureType = ruleFrameworkExtension.getFeatureTypeSupplier().get(); + AutoTaggingRegistry.registerFeatureType(featureType); + rulePersistenceServiceRegistry.register(featureType, ruleFrameworkExtension.getRulePersistenceServiceSupplier().get()); } } diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java index 8fc51176bf173..61c0a15d7ed96 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java @@ -8,87 +8,212 @@ package org.opensearch.rule.action; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.ExceptionsHelper; +import org.opensearch.ResourceAlreadyExistsException; +import org.opensearch.action.ActionListenerResponseHandler; +import org.opensearch.action.admin.indices.create.CreateIndexRequest; +import org.opensearch.action.admin.indices.create.CreateIndexResponse; import org.opensearch.action.support.ActionFilters; -import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction; +import org.opensearch.action.support.TransportAction; 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.node.DiscoveryNode; +import org.opensearch.cluster.routing.IndexRoutingTable; +import org.opensearch.cluster.routing.ShardRouting; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; -import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.rule.CreateRuleRequest; import org.opensearch.rule.CreateRuleResponse; import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.RulePersistenceServiceRegistry; +import org.opensearch.rule.service.IndexStoredRulePersistenceService; +import org.opensearch.tasks.Task; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportChannel; +import org.opensearch.transport.TransportException; +import org.opensearch.transport.TransportRequestHandler; import org.opensearch.transport.TransportService; +import org.opensearch.transport.client.Client; import java.io.IOException; +import java.util.Map; -import static org.opensearch.threadpool.ThreadPool.Names.SAME; +import static org.opensearch.rule.RuleFrameworkPlugin.RULE_THREAD_POOL_NAME; /** * Transport action to create Rules * @opensearch.experimental */ -public class TransportCreateRuleAction extends TransportClusterManagerNodeAction { - +public class TransportCreateRuleAction extends TransportAction { + private final TransportService transportService; + private final ClusterService clusterService; + private final ThreadPool threadPool; private final RulePersistenceServiceRegistry rulePersistenceServiceRegistry; + private final Client client; + private static final Logger logger = LogManager.getLogger(TransportCreateRuleAction.class); + private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); /** * Constructor for TransportCreateRuleAction * * @param threadPool - {@link ThreadPool} object + * @param client - {@link Client} object * @param transportService - a {@link TransportService} object * @param clusterService - a {@link ClusterService} object * @param actionFilters - a {@link ActionFilters} object - * @param indexNameExpressionResolver - {@link IndexNameExpressionResolver} object * @param rulePersistenceServiceRegistry - a {@link RulePersistenceServiceRegistry} object */ @Inject public TransportCreateRuleAction( - ThreadPool threadPool, + Client client, TransportService transportService, ClusterService clusterService, + ThreadPool threadPool, ActionFilters actionFilters, - IndexNameExpressionResolver indexNameExpressionResolver, RulePersistenceServiceRegistry rulePersistenceServiceRegistry ) { - super( + super(CreateRuleAction.NAME, actionFilters, transportService.getTaskManager()); + this.client = client; + this.transportService = transportService; + this.clusterService = clusterService; + this.threadPool = threadPool; + this.rulePersistenceServiceRegistry = rulePersistenceServiceRegistry; + + transportService.registerRequestHandler( CreateRuleAction.NAME, - transportService, - clusterService, - threadPool, - actionFilters, + ThreadPool.Names.WRITE, CreateRuleRequest::new, - indexNameExpressionResolver + new TransportRequestHandler() { + @Override + public void messageReceived(CreateRuleRequest request, TransportChannel channel, Task task) { + executeLocally(request, ActionListener.wrap(response -> { + try { + channel.sendResponse(response); + } catch (IOException e) { + logger.error("Failed to send CreateRuleResponse to transport channel", e); + throw new TransportException("Fail to send", e); + } + }, exception -> { + try { + channel.sendResponse(exception); + } catch (IOException e) { + logger.error("Failed to send exception response to transport channel", e); + throw new TransportException("Fail to send", e); + } + })); + } + } ); - this.rulePersistenceServiceRegistry = rulePersistenceServiceRegistry; } @Override - protected String executor() { - return SAME; + protected void doExecute(Task task, CreateRuleRequest request, ActionListener listener) { + RulePersistenceService persistenceService = rulePersistenceServiceRegistry.getRulePersistenceService( + request.getRule().getFeatureType() + ); + if (!(persistenceService instanceof IndexStoredRulePersistenceService indexStoredRulePersistenceService)) { + executeLocally(request, listener); + return; + } + + String indexName = indexStoredRulePersistenceService.getIndexName(); + try (ThreadContext.StoredContext ctx = client.threadPool().getThreadContext().stashContext()) { + if (!clusterService.state().metadata().hasIndex(indexName)) { + createIndex(() -> routeRequest(request, listener, indexName), indexName, listener); + return; + } + } + + routeRequest(request, listener, indexName); } - @Override - protected CreateRuleResponse read(StreamInput in) throws IOException { - return new CreateRuleResponse(in); + /** + * Creates the backing index if it does not exist, then runs the given success callback. + * @param onSuccess callback to run after successful index creation + * @param indexName the name of the index to create + * @param listener listener to handle failures + */ + private void createIndex(Runnable onSuccess, String indexName, ActionListener listener) { + final CreateIndexRequest request = new CreateIndexRequest(indexName).settings(indexSettings); + client.admin().indices().create(request, new ActionListener<>() { + @Override + public void onResponse(CreateIndexResponse response) { + if (!response.isAcknowledged()) { + listener.onFailure(new IllegalStateException(indexName + " index creation not acknowledged")); + } else { + onSuccess.run(); + } + } + + @Override + public void onFailure(Exception e) { + Throwable cause = ExceptionsHelper.unwrapCause(e); + if (cause instanceof ResourceAlreadyExistsException) { + onSuccess.run(); + } else { + listener.onFailure(e); + } + } + }); } - @Override - protected void clusterManagerOperation(CreateRuleRequest request, ClusterState state, ActionListener listener) - throws Exception { - final RulePersistenceService rulePersistenceService = rulePersistenceServiceRegistry.getRulePersistenceService( - request.getRule().getFeatureType() - ); - rulePersistenceService.createRule(request, listener); + /** + * Routes the CreateRuleRequest to the primary shard node for the given index. + * Executes locally if the current node is the primary. + * @param request the CreateRuleRequest + * @param listener listener to handle response or failure + * @param indexName the index name used to find the primary shard node + */ + private void routeRequest(CreateRuleRequest request, ActionListener listener, String indexName) { + DiscoveryNode primaryNode = getPrimaryShardNode(indexName); + if (primaryNode == null) { + listener.onFailure(new IllegalStateException("Primary node for index [" + indexName + "] not found")); + return; + } + + if (transportService.getLocalNode().equals(primaryNode)) { + executeLocally(request, listener); + } else { + transportService.sendRequest( + primaryNode, + CreateRuleAction.NAME, + request, + new ActionListenerResponseHandler<>(listener, CreateRuleResponse::new) + ); + } } - @Override - protected ClusterBlockException checkBlock(CreateRuleRequest request, ClusterState state) { - return state.blocks().globalBlockedException(ClusterBlockLevel.WRITE); + /** + * Retrieves the discovery node that holds the primary shard for the given index. + * @param indexName the index name + */ + private DiscoveryNode getPrimaryShardNode(String indexName) { + ClusterState state = clusterService.state(); + IndexRoutingTable indexRoutingTable = state.getRoutingTable().index(indexName); + if (indexRoutingTable == null) { + return null; + } + ShardRouting primaryShard = indexRoutingTable.shard(0).primaryShard(); + if (primaryShard == null || !primaryShard.assignedToNode()) { + return null; + } + return state.nodes().get(primaryShard.currentNodeId()); + } + + /** + * Executes the create rule operation locally on the dedicated rule thread pool. + * @param request the CreateRuleRequest + * @param listener listener to handle response or failure + */ + private void executeLocally(CreateRuleRequest request, ActionListener listener) { + threadPool.executor(RULE_THREAD_POOL_NAME).execute(() -> { + RulePersistenceService persistenceService = rulePersistenceServiceRegistry.getRulePersistenceService( + request.getRule().getFeatureType() + ); + persistenceService.createRule(request, listener); + }); } } diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/RuleFrameworkPluginTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/RuleFrameworkPluginTests.java index e9f0a68d35dea..4e11d12f9facb 100644 --- a/modules/autotagging-commons/src/test/java/org/opensearch/rule/RuleFrameworkPluginTests.java +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/RuleFrameworkPluginTests.java @@ -28,7 +28,7 @@ public class RuleFrameworkPluginTests extends OpenSearchTestCase { public void testGetActions() { List> handlers = plugin.getActions(); - assertEquals(2, handlers.size()); + assertEquals(3, handlers.size()); assertEquals(GetRuleAction.INSTANCE.name(), handlers.get(0).getAction().name()); } diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/RulePersistenceServiceRegistryTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/RulePersistenceServiceRegistryTests.java index 116ad7a730c58..aecdef198d39f 100644 --- a/modules/autotagging-commons/src/test/java/org/opensearch/rule/RulePersistenceServiceRegistryTests.java +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/RulePersistenceServiceRegistryTests.java @@ -15,8 +15,8 @@ import static org.mockito.Mockito.when; public class RulePersistenceServiceRegistryTests extends OpenSearchTestCase { - RulePersistenceServiceRegistry registry = new RulePersistenceServiceRegistry();; - FeatureType mockFeatureType = mock(FeatureType.class);; + RulePersistenceServiceRegistry registry = new RulePersistenceServiceRegistry(); + FeatureType mockFeatureType = mock(FeatureType.class); RulePersistenceService mockService = mock(RulePersistenceService.class); public void testRegisterAndGetService() { diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java index 1574882825ffb..8b063cb0afbe3 100644 --- a/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java @@ -9,50 +9,74 @@ package org.opensearch.rule.action; import org.opensearch.action.support.ActionFilters; -import org.opensearch.cluster.metadata.IndexNameExpressionResolver; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.core.action.ActionListener; import org.opensearch.rule.CreateRuleRequest; +import org.opensearch.rule.CreateRuleResponse; import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.RulePersistenceServiceRegistry; +import org.opensearch.rule.autotagging.FeatureType; import org.opensearch.rule.autotagging.Rule; +import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; +import org.opensearch.transport.client.Client; +import java.util.concurrent.ExecutorService; + +import static org.opensearch.rule.RuleFrameworkPlugin.RULE_THREAD_POOL_NAME; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +@SuppressWarnings("unchecked") public class TransportCreateRuleActionTests extends OpenSearchTestCase { - TransportCreateRuleAction sut; - - public void testExecute() throws Exception { - RulePersistenceServiceRegistry rulePersistenceServiceRegistry = mock(RulePersistenceServiceRegistry.class); - TransportService transportService = mock(TransportService.class); - ActionFilters actionFilters = mock(ActionFilters.class); - RulePersistenceService rulePersistenceService = mock(RulePersistenceService.class); - CreateRuleRequest createRuleRequest = mock(CreateRuleRequest.class); + private TransportService transportService; + private ClusterService clusterService; + private ThreadPool threadPool; + private RulePersistenceServiceRegistry registry; + private Client client; + private ActionFilters actionFilters; + private TransportCreateRuleAction action; + private FeatureType mockFeatureType; + + private final String testIndexName = "test-index"; + + public void setUp() throws Exception { + super.setUp(); + transportService = mock(TransportService.class); + clusterService = mock(ClusterService.class); + threadPool = mock(ThreadPool.class); + registry = mock(RulePersistenceServiceRegistry.class); + client = mock(Client.class); + actionFilters = mock(ActionFilters.class); + mockFeatureType = mock(FeatureType.class); + RulePersistenceServiceRegistry registry = new RulePersistenceServiceRegistry(); + when(mockFeatureType.getName()).thenReturn("test_feature"); + RulePersistenceService mockService = mock(RulePersistenceService.class); + registry.register(mockFeatureType, mockService); + + ExecutorService executorService = mock(ExecutorService.class); + doAnswer(invocation -> { + Runnable runnable = invocation.getArgument(0); + runnable.run(); + return null; + }).when(executorService).execute(any()); + when(threadPool.executor(any())).thenReturn(executorService); + action = new TransportCreateRuleAction(client, transportService, clusterService, threadPool, actionFilters, registry); + } + + public void testExecution() { + IndexStoredRulePersistenceService persistenceService = mock(IndexStoredRulePersistenceService.class); + when(registry.getRulePersistenceService(mockFeatureType)).thenReturn(persistenceService); Rule rule = mock(Rule.class); - when(createRuleRequest.getRule()).thenReturn(rule); - when(createRuleRequest.getRule().getFeatureType()).thenReturn(null); - ThreadPool threadPool = mock(ThreadPool.class); - ClusterService clusterService = mock(ClusterService.class); - IndexNameExpressionResolver indexNameExpressionResolver = mock(IndexNameExpressionResolver.class); - - when(rulePersistenceServiceRegistry.getRulePersistenceService(any())).thenReturn(rulePersistenceService); - doNothing().when(rulePersistenceService).getRule(any(), any()); - sut = new TransportCreateRuleAction( - threadPool, - transportService, - clusterService, - actionFilters, - indexNameExpressionResolver, - rulePersistenceServiceRegistry - ); - sut.clusterManagerOperation(createRuleRequest, null, null); - verify(rulePersistenceService, times(1)).createRule(any(), any()); + when(rule.getFeatureType()).thenReturn(mockFeatureType); + CreateRuleRequest request = new CreateRuleRequest(rule); + ActionListener listener = mock(ActionListener.class); + action.doExecute(null, request, listener); + verify(threadPool).executor(RULE_THREAD_POOL_NAME); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index 0cadc579ba9a4..e967bc1138083 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -77,8 +77,8 @@ public class WorkloadManagementPlugin extends Plugin implements ActionPlugin, Sy * The maximum number of rules allowed per GET request. */ public static final int MAX_RULES_PER_PAGE = 50; - - private final RulePersistenceServiceHolder rulePersistenceServiceHolder = new RulePersistenceServiceHolder(); + private static FeatureType featureType; + private static RulePersistenceService rulePersistenceService; private AutoTaggingActionFilter autoTaggingActionFilter; @@ -101,13 +101,13 @@ public Collection createComponents( IndexNameExpressionResolver indexNameExpressionResolver, Supplier repositoriesServiceSupplier ) { - FeatureTypeHolder.featureType = new WorkloadGroupFeatureType(new WorkloadGroupFeatureValueValidator(clusterService)); - RulePersistenceServiceHolder.rulePersistenceService = new IndexStoredRulePersistenceService( + featureType = new WorkloadGroupFeatureType(new WorkloadGroupFeatureValueValidator(clusterService)); + rulePersistenceService = new IndexStoredRulePersistenceService( INDEX_NAME, client, clusterService, MAX_RULES_PER_PAGE, - new XContentRuleParser(FeatureTypeHolder.featureType), + new XContentRuleParser(featureType), new IndexBasedRuleQueryMapper() ); InMemoryRuleProcessingService ruleProcessingService = new InMemoryRuleProcessingService( @@ -168,19 +168,11 @@ public Collection createGuiceModules() { @Override public Supplier getRulePersistenceServiceSupplier() { - return () -> RulePersistenceServiceHolder.rulePersistenceService; + return () -> rulePersistenceService; } @Override - public FeatureType getFeatureType() { - return FeatureTypeHolder.featureType; - } - - static class RulePersistenceServiceHolder { - private static RulePersistenceService rulePersistenceService; - } - - static class FeatureTypeHolder { - private static FeatureType featureType; + public Supplier getFeatureTypeSupplier() { + return () -> featureType; } } diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java index b79ba4107617b..458707fb1aee1 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java @@ -111,7 +111,7 @@ public void testGetFeatureTypeReturnsWorkloadGroupFeatureType() { mock(IndexNameExpressionResolver.class), () -> mock(RepositoriesService.class) ); - FeatureType featureType = plugin.getFeatureType(); + FeatureType featureType = plugin.getFeatureTypeSupplier().get(); assertEquals("workload_group", featureType.getName()); } From 778645c2d33581e597a9be7fd725521933d59347 Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Thu, 29 May 2025 15:18:43 -0700 Subject: [PATCH 20/20] adding plugin-specific routing layer Signed-off-by: Ruirui Zhang --- .../opensearch/rule/RuleRoutingService.java | 25 ++ .../IndexStoredRulePersistenceService.java | 86 +++-- .../org/opensearch/rule/RuleUtilsTests.java | 48 ++- .../rule/action/CreateRuleRequestTests.java | 4 +- .../rule/action/CreateRuleResponseTests.java | 2 +- .../rule/action/DeleteRuleRequestTests.java | 2 +- .../rule/action/GetRuleRequestTests.java | 67 +--- .../rule/action/GetRuleResponseTests.java | 10 +- .../rule/autotagging/RuleTests.java | 1 - ...ndexStoredRulePersistenceServiceTests.java | 301 +++++++++--------- .../opensearch/rule/utils/RuleTestUtils.java | 57 +++- .../rule/spi/RuleFrameworkExtension.java | 7 + .../opensearch/rule/RuleFrameworkPlugin.java | 8 +- .../rule/RulePersistenceServiceRegistry.java | 2 +- .../rule/RuleRoutingServiceRegistry.java | 48 +++ .../action/TransportCreateRuleAction.java | 135 +------- .../rule/rest/RestCreateRuleAction.java | 1 + .../RulePersistenceServiceRegistryTests.java | 2 +- .../TransportCreateRuleActionTests.java | 30 +- .../plugin/wlm/WorkloadManagementPlugin.java | 26 +- .../rule/WorkloadGroupRuleRoutingService.java | 143 +++++++++ .../wlm/AutoTaggingActionFilterTests.java | 3 - .../rule/WorkloadGroupFeatureTypeTests.java | 5 +- .../main/java/org/opensearch/node/Node.java | 1 + 24 files changed, 579 insertions(+), 435 deletions(-) create mode 100644 modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleRoutingService.java create mode 100644 modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleRoutingServiceRegistry.java create mode 100644 plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupRuleRoutingService.java diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleRoutingService.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleRoutingService.java new file mode 100644 index 0000000000000..e0d08f371a2aa --- /dev/null +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleRoutingService.java @@ -0,0 +1,25 @@ +/* + * 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.rule; + +import org.opensearch.core.action.ActionListener; + +/** + * Interface that handles rule routing logic + * @opensearch.experimental + */ +public interface RuleRoutingService { + + /** + * Handles a create rule request by routing it to the appropriate node. + * @param request the create rule request + * @param listener listener to handle the final response + */ + void handleCreateRuleRequest(CreateRuleRequest request, ActionListener listener); +} diff --git a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java index 2b1e523c2b198..b0d31a829b2b6 100644 --- a/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java +++ b/modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java @@ -14,7 +14,9 @@ import org.opensearch.action.DocWriteResponse; import org.opensearch.action.delete.DeleteRequest; import org.opensearch.action.index.IndexRequest; +import org.opensearch.action.index.IndexResponse; import org.opensearch.action.search.SearchRequestBuilder; +import org.opensearch.action.search.SearchResponse; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.util.concurrent.ThreadContext; @@ -37,7 +39,6 @@ import org.opensearch.search.sort.SortOrder; import org.opensearch.transport.client.Client; -import java.io.IOException; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -95,10 +96,10 @@ public IndexStoredRulePersistenceService( * @param listener ActionListener for CreateRuleResponse */ public void createRule(CreateRuleRequest request, ActionListener listener) { - try (ThreadContext.StoredContext ctx = getContext()) { + try (ThreadContext.StoredContext ctx = stashContext()) { if (!clusterService.state().metadata().hasIndex(indexName)) { logger.error("Index {} does not exist", indexName); - throw new IllegalStateException("Index" + indexName + " does not exist"); + listener.onFailure(new IllegalStateException("Index" + indexName + " does not exist")); } else { Rule rule = request.getRule(); validateNoDuplicateRule(rule, ActionListener.wrap(unused -> persistRule(rule, listener), listener::onFailure)); @@ -107,30 +108,28 @@ public void createRule(CreateRuleRequest request, ActionListener listener) { - try (ThreadContext.StoredContext ctx = getContext()) { - QueryBuilder query = queryBuilder.from(new GetRuleRequest(null, rule.getAttributeMap(), null, rule.getFeatureType())); - getRuleFromIndex(null, query, null, new ActionListener<>() { - @Override - public void onResponse(GetRuleResponse getRuleResponse) { - Optional duplicateRuleId = RuleUtils.getDuplicateRuleId(rule, getRuleResponse.getRules()); - duplicateRuleId.ifPresentOrElse( - id -> listener.onFailure(new IllegalArgumentException("Duplicate rule exists under id " + id)), - () -> listener.onResponse(null) - ); - } + QueryBuilder query = queryBuilder.from(new GetRuleRequest(null, rule.getAttributeMap(), null, rule.getFeatureType())); + getRuleFromIndex(null, query, null, new ActionListener<>() { + @Override + public void onResponse(GetRuleResponse getRuleResponse) { + Optional duplicateRuleId = RuleUtils.getDuplicateRuleId(rule, getRuleResponse.getRules()); + duplicateRuleId.ifPresentOrElse( + id -> listener.onFailure(new IllegalArgumentException("Duplicate rule exists under id " + id)), + () -> listener.onResponse(null) + ); + } - @Override - public void onFailure(Exception e) { - listener.onFailure(e); - } - }); - } + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + }); } /** @@ -139,17 +138,13 @@ public void onFailure(Exception e) { * @param listener - ActionListener for CreateRuleResponse */ private void persistRule(Rule rule, ActionListener listener) { - try (ThreadContext.StoredContext ctx = getContext()) { + try { IndexRequest indexRequest = new IndexRequest(indexName).source( rule.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS) ); - client.index(indexRequest, ActionListener.wrap(indexResponse -> { - listener.onResponse(new CreateRuleResponse(indexResponse.getId(), rule)); - }, e -> { - logger.warn("Failed to save Rule object due to error: {}", e.getMessage()); - listener.onFailure(e); - })); - } catch (IOException e) { + IndexResponse indexResponse = client.index(indexRequest).get(); + listener.onResponse(new CreateRuleResponse(indexResponse.getId(), rule)); + } catch (Exception e) { logger.error("Error saving rule to index: {}", indexName); listener.onFailure(new RuntimeException("Failed to save rule to index.")); } @@ -161,8 +156,10 @@ private void persistRule(Rule rule, ActionListener listener) * @param listener the listener for GetRuleResponse. */ public void getRule(GetRuleRequest getRuleRequest, ActionListener listener) { - final QueryBuilder getQueryBuilder = queryBuilder.from(getRuleRequest); - getRuleFromIndex(getRuleRequest.getId(), getQueryBuilder, getRuleRequest.getSearchAfter(), listener); + try (ThreadContext.StoredContext context = stashContext()) { + final QueryBuilder getQueryBuilder = queryBuilder.from(getRuleRequest); + getRuleFromIndex(getRuleRequest.getId(), getQueryBuilder, getRuleRequest.getSearchAfter(), listener); + } } /** @@ -173,22 +170,19 @@ public void getRule(GetRuleRequest getRuleRequest, ActionListener listener) { - // Stash the current thread context when interacting with system index to perform - // operations as the system itself, bypassing authorization checks. This ensures that - // actions within this block are trusted and executed with system-level privileges. - try (ThreadContext.StoredContext context = getContext()) { + try { SearchRequestBuilder searchRequest = client.prepareSearch(indexName).setQuery(queryBuilder).setSize(maxRulesPerPage); if (searchAfter != null) { searchRequest.addSort(_ID_STRING, SortOrder.ASC).searchAfter(new Object[] { searchAfter }); } - searchRequest.execute(ActionListener.wrap(searchResponse -> { - List hits = Arrays.asList(searchResponse.getHits().getHits()); - if (hasNoResults(id, listener, hits)) return; - handleGetRuleResponse(hits, listener); - }, e -> { - logger.error("Failed to fetch all rules: {}", e.getMessage()); - listener.onFailure(e); - })); + + SearchResponse searchResponse = searchRequest.get(); + List hits = Arrays.asList(searchResponse.getHits().getHits()); + if (hasNoResults(id, listener, hits)) return; + handleGetRuleResponse(hits, listener); + } catch (Exception e) { + logger.error("Failed to fetch all rules: {}", e.getMessage()); + listener.onFailure(e); } } @@ -214,7 +208,7 @@ void handleGetRuleResponse(List hits, ActionListener @Override public void deleteRule(DeleteRuleRequest request, ActionListener listener) { - try (ThreadContext.StoredContext context = getContext()) { + try (ThreadContext.StoredContext context = stashContext()) { DeleteRequest deleteRequest = new DeleteRequest(indexName).id(request.getRuleId()); client.delete(deleteRequest, ActionListener.wrap(deleteResponse -> { boolean acknowledged = deleteResponse.getResult() == DocWriteResponse.Result.DELETED; @@ -241,7 +235,7 @@ public String getIndexName() { return indexName; } - private ThreadContext.StoredContext getContext() { + private ThreadContext.StoredContext stashContext() { return client.threadPool().getThreadContext().stashContext(); } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleUtilsTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleUtilsTests.java index e2157797a9b6c..2780f329925c9 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleUtilsTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleUtilsTests.java @@ -9,6 +9,7 @@ package org.opensearch.rule; import org.opensearch.rule.autotagging.Rule; +import org.opensearch.rule.autotagging.RuleTests; import org.opensearch.rule.utils.RuleTestUtils; import org.opensearch.test.OpenSearchTestCase; @@ -16,15 +17,15 @@ import java.util.Optional; import java.util.Set; -import static org.opensearch.rule.action.GetRuleRequestTests.ATTRIBUTE_VALUE_ONE; -import static org.opensearch.rule.action.GetRuleRequestTests.ATTRIBUTE_VALUE_TWO; -import static org.opensearch.rule.action.GetRuleRequestTests.DESCRIPTION_ONE; -import static org.opensearch.rule.action.GetRuleRequestTests.FEATURE_VALUE_ONE; -import static org.opensearch.rule.action.GetRuleRequestTests.TIMESTAMP_ONE; -import static org.opensearch.rule.action.GetRuleRequestTests._ID_ONE; -import static org.opensearch.rule.action.GetRuleRequestTests._ID_TWO; -import static org.opensearch.rule.action.GetRuleRequestTests.ruleTwo; import static org.opensearch.rule.action.GetRuleResponseTests.ruleOne; +import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_VALUE_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_VALUE_TWO; +import static org.opensearch.rule.utils.RuleTestUtils.DESCRIPTION_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.FEATURE_VALUE_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.TIMESTAMP_ONE; +import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.utils.RuleTestUtils._ID_TWO; +import static org.opensearch.rule.utils.RuleTestUtils.ruleTwo; public class RuleUtilsTests extends OpenSearchTestCase { @@ -34,12 +35,12 @@ public void testDuplicateRuleFound() { assertEquals(_ID_ONE, result.get()); } - public void testNoDuplicate_NoAttributeIntersection() { + public void testNoAttributeIntersection() { Optional result = RuleUtils.getDuplicateRuleId(ruleOne, Map.of(_ID_TWO, ruleTwo)); assertTrue(result.isEmpty()); } - public void testNoDuplicate_AttributeSizeMismatch() { + public void testAttributeSizeMismatch() { Rule testRule = Rule.builder() .description(DESCRIPTION_ONE) .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) @@ -57,4 +58,31 @@ public void testNoDuplicate_AttributeSizeMismatch() { Optional result = RuleUtils.getDuplicateRuleId(ruleOne, Map.of(_ID_TWO, testRule)); assertTrue(result.isEmpty()); } + + public void testPartialAttributeValueIntersection() { + Rule ruleWithPartialOverlap = Rule.builder() + .description(DESCRIPTION_ONE) + .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_ONE) + .attributeMap(Map.of(RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, Set.of(ATTRIBUTE_VALUE_ONE, "extra_value"))) + .updatedAt(TIMESTAMP_ONE) + .build(); + + Optional result = RuleUtils.getDuplicateRuleId(ruleWithPartialOverlap, Map.of(_ID_ONE, ruleOne)); + assertTrue(result.isPresent()); + assertEquals(_ID_ONE, result.get()); + } + + public void testDifferentFeatureTypes() { + Rule differentFeatureTypeRule = Rule.builder() + .description(DESCRIPTION_ONE) + .featureType(RuleTests.TestFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_ONE) + .attributeMap(RuleTests.ATTRIBUTE_MAP) + .updatedAt(TIMESTAMP_ONE) + .build(); + + Optional result = RuleUtils.getDuplicateRuleId(differentFeatureTypeRule, Map.of(_ID_ONE, ruleOne)); + assertTrue(result.isEmpty()); + } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java index 4f6626df9af3c..7804714c37dcc 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleRequestTests.java @@ -15,8 +15,8 @@ import java.io.IOException; -import static org.opensearch.rule.action.GetRuleRequestTests.assertEqualRule; -import static org.opensearch.rule.action.GetRuleRequestTests.ruleOne; +import static org.opensearch.rule.utils.RuleTestUtils.assertEqualRule; +import static org.opensearch.rule.utils.RuleTestUtils.ruleOne; public class CreateRuleRequestTests extends OpenSearchTestCase { diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java index e4bdf5fd94c53..dc445dad2e82c 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/CreateRuleResponseTests.java @@ -20,9 +20,9 @@ import java.io.IOException; import java.util.Map; -import static org.opensearch.rule.action.GetRuleRequestTests.assertEqualRules; import static org.opensearch.rule.action.GetRuleResponseTests.ruleOne; import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.assertEqualRules; import static org.mockito.Mockito.mock; public class CreateRuleResponseTests extends OpenSearchTestCase { diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/DeleteRuleRequestTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/DeleteRuleRequestTests.java index 315f94c0e0437..55213a245b5ad 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/DeleteRuleRequestTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/DeleteRuleRequestTests.java @@ -16,7 +16,7 @@ import java.io.IOException; -import static org.opensearch.rule.action.GetRuleRequestTests._ID_ONE; +import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; public class DeleteRuleRequestTests extends OpenSearchTestCase { diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java index a451a58356606..c904373588f17 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleRequestTests.java @@ -11,15 +11,15 @@ import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.rule.GetRuleRequest; -import org.opensearch.rule.autotagging.Attribute; -import org.opensearch.rule.autotagging.Rule; import org.opensearch.rule.utils.RuleTestUtils; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; import java.util.HashMap; -import java.util.Map; -import java.util.Set; + +import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_MAP; +import static org.opensearch.rule.utils.RuleTestUtils.SEARCH_AFTER; +import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; public class GetRuleRequestTests extends OpenSearchTestCase { /** @@ -64,63 +64,4 @@ public void testValidate() { request = new GetRuleRequest(_ID_ONE, ATTRIBUTE_MAP, "", RuleTestUtils.MockRuleFeatureType.INSTANCE); assertThrows(IllegalArgumentException.class, request::validate); } - - public static final String _ID_ONE = "id_1"; - public static final String SEARCH_AFTER = "search_after"; - public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; - public static final String FEATURE_VALUE_ONE = "feature_value_one"; - public static final String FEATURE_VALUE_TWO = "feature_value_two"; - public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; - public static final String ATTRIBUTE_VALUE_TWO = "mock_attribute_two"; - public static final String DESCRIPTION_ONE = "description_1"; - public static final String DESCRIPTION_TWO = "description_2"; - public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; - public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; - public static final Map> ATTRIBUTE_MAP = Map.of( - RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, - Set.of(ATTRIBUTE_VALUE_ONE) - ); - - public static final Rule ruleOne = Rule.builder() - .description(DESCRIPTION_ONE) - .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_ONE) - .attributeMap(ATTRIBUTE_MAP) - .updatedAt(TIMESTAMP_ONE) - .build(); - - public static final Rule ruleTwo = Rule.builder() - .description(DESCRIPTION_TWO) - .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) - .featureValue(FEATURE_VALUE_TWO) - .attributeMap(Map.of(RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(ATTRIBUTE_VALUE_TWO))) - .updatedAt(TIMESTAMP_TWO) - .build(); - - public static Map ruleMap() { - return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); - } - - public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { - assertEquals(mapOne.size(), mapTwo.size()); - for (Map.Entry entry : mapOne.entrySet()) { - String id = entry.getKey(); - assertTrue(mapTwo.containsKey(id)); - Rule one = mapOne.get(id); - Rule two = mapTwo.get(id); - assertEqualRule(one, two, ruleUpdated); - } - } - - public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { - if (ruleUpdated) { - assertEquals(one.getDescription(), two.getDescription()); - assertEquals(one.getFeatureType(), two.getFeatureType()); - assertEquals(one.getFeatureValue(), two.getFeatureValue()); - assertEquals(one.getAttributeMap(), two.getAttributeMap()); - assertEquals(one.getAttributeMap(), two.getAttributeMap()); - } else { - assertEquals(one, two); - } - } } diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java index f01bb94fab276..4a8fc3ef6d7bc 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/action/GetRuleResponseTests.java @@ -24,10 +24,10 @@ import java.util.Map; import java.util.Set; -import static org.opensearch.rule.action.GetRuleRequestTests.SEARCH_AFTER; -import static org.opensearch.rule.action.GetRuleRequestTests._ID_ONE; -import static org.opensearch.rule.action.GetRuleRequestTests.assertEqualRules; -import static org.opensearch.rule.action.GetRuleRequestTests.ruleMap; +import static org.opensearch.rule.utils.RuleTestUtils.SEARCH_AFTER; +import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.assertEqualRules; +import static org.opensearch.rule.utils.RuleTestUtils.ruleMap; import static org.mockito.Mockito.mock; public class GetRuleResponseTests extends OpenSearchTestCase { @@ -109,7 +109,7 @@ public void testToXContentGetSingleRule() throws IOException { String expected = "{\n" + " \"rules\" : [\n" + " {\n" - + " \"_id\" : \"id_1\",\n" + + " \"_id\" : \"AgfUO5Ja9yfvhdONlYi3TQ==\",\n" + " \"description\" : \"description_1\",\n" + " \"mock_attribute_one\" : [\n" + " \"mock_attribute_one\"\n" diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java index c504e1ab6b53e..5f20640a74f24 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java @@ -39,7 +39,6 @@ public class RuleTests extends AbstractSerializingTestCase { Set.of("value2") ); public static final String UPDATED_AT = "2025-02-24T07:42:10.123456Z"; - public static final String INVALID_CLASS = "invalid_class"; public static final String INVALID_ATTRIBUTE = "invalid_attribute"; public static final String INVALID_FEATURE = "invalid_feature"; diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java index f3fb4b65961c3..eb054ea8124e5 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java @@ -20,6 +20,7 @@ import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.action.ActionFuture; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; @@ -43,16 +44,26 @@ import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.Client; -import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import org.mockito.ArgumentCaptor; import static org.opensearch.rule.XContentRuleParserTests.VALID_JSON; +import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_MAP; import static org.opensearch.rule.utils.RuleTestUtils.ATTRIBUTE_VALUE_ONE; import static org.opensearch.rule.utils.RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE; +import static org.opensearch.rule.utils.RuleTestUtils.MockRuleFeatureType; import static org.opensearch.rule.utils.RuleTestUtils.TEST_INDEX_NAME; import static org.opensearch.rule.utils.RuleTestUtils._ID_ONE; import static org.mockito.ArgumentMatchers.any; @@ -60,7 +71,6 @@ import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.any; import static org.mockito.Mockito.anyInt; -import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -68,108 +78,171 @@ @SuppressWarnings("unchecked") public class IndexStoredRulePersistenceServiceTests extends OpenSearchTestCase { - public static final int MAX_VALUES_PER_PAGE = 50; + private static final int MAX_VALUES_PER_PAGE = 50; - public void testCreateRuleOnExistingIndex() throws IOException { - CreateRuleRequest createRuleRequest = mock(CreateRuleRequest.class); - Rule mockRule = mock(Rule.class); - when(createRuleRequest.getRule()).thenReturn(mockRule); - RuleQueryMapper mockRuleQueryMapper = mock(RuleQueryMapper.class); - RuleEntityParser mockRuleEntityParser = mock(RuleEntityParser.class); - ClusterService clusterService = mock(ClusterService.class); + private Client client; + private ClusterService clusterService; + private RuleQueryMapper ruleQueryMapper; + private RuleEntityParser ruleEntityParser; + private SearchRequestBuilder searchRequestBuilder; + private RulePersistenceService rulePersistenceService; + private QueryBuilder queryBuilder; + private Rule rule; + + public void setUp() throws Exception { + super.setUp(); + searchRequestBuilder = mock(SearchRequestBuilder.class); + client = setUpMockClient(searchRequestBuilder); + + rule = mock(Rule.class); + + clusterService = mock(ClusterService.class); ClusterState clusterState = mock(ClusterState.class); Metadata metadata = mock(Metadata.class); - QueryBuilder queryBuilder = mock(QueryBuilder.class); - when(clusterService.state()).thenReturn(clusterState); when(clusterState.metadata()).thenReturn(metadata); when(metadata.hasIndex(TEST_INDEX_NAME)).thenReturn(true); - when(mockRuleQueryMapper.from(any(GetRuleRequest.class))).thenReturn(queryBuilder); + + ruleQueryMapper = mock(RuleQueryMapper.class); + ruleEntityParser = mock(RuleEntityParser.class); + queryBuilder = mock(QueryBuilder.class); when(queryBuilder.filter(any())).thenReturn(queryBuilder); + when(ruleQueryMapper.from(any(GetRuleRequest.class))).thenReturn(queryBuilder); + when(ruleEntityParser.parse(anyString())).thenReturn(rule); - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - Client client = setUpMockClient(searchRequestBuilder); - RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( + rulePersistenceService = new IndexStoredRulePersistenceService( TEST_INDEX_NAME, client, clusterService, MAX_VALUES_PER_PAGE, - mockRuleEntityParser, - mockRuleQueryMapper + ruleEntityParser, + ruleQueryMapper ); - ActionListener listener = mock(ActionListener.class); - when(mockRule.toXContent(any(), any())).thenAnswer(invocation -> invocation.getArgument(0)); + } + + public void testCreateRuleOnExistingIndex() throws Exception { + CreateRuleRequest createRuleRequest = mock(CreateRuleRequest.class); + when(createRuleRequest.getRule()).thenReturn(rule); + when(rule.toXContent(any(), any())).thenAnswer(invocation -> invocation.getArgument(0)); SearchResponse searchResponse = mock(SearchResponse.class); when(searchResponse.getHits()).thenReturn(new SearchHits(new SearchHit[] {}, new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1.0f)); - doAnswer((invocation) -> { - ActionListener actionListener = invocation.getArgument(0); - actionListener.onResponse(searchResponse); - return null; - }).when(searchRequestBuilder).execute(any(ActionListener.class)); + when(searchRequestBuilder.get()).thenReturn(searchResponse); IndexResponse indexResponse = mock(IndexResponse.class); when(indexResponse.getId()).thenReturn(_ID_ONE); - doAnswer(invocation -> { - ActionListener actionListener = invocation.getArgument(1); - actionListener.onResponse(indexResponse); - return null; - }).when(client).index(any(IndexRequest.class), any(ActionListener.class)); + ActionFuture future = mock(ActionFuture.class); + when(future.get()).thenReturn(indexResponse); + when(client.index(any(IndexRequest.class))).thenReturn(future); + ActionListener listener = mock(ActionListener.class); rulePersistenceService.createRule(createRuleRequest, listener); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(CreateRuleResponse.class); verify(listener).onResponse(responseCaptor.capture()); - CreateRuleResponse response = responseCaptor.getValue(); - assertNotNull(response.getRule()); + assertNotNull(responseCaptor.getValue().getRule()); } - public void testCreateDuplicateRule() { - CreateRuleRequest createRuleRequest = mock(CreateRuleRequest.class); - Rule mockRule = mock(Rule.class); - when(createRuleRequest.getRule()).thenReturn(mockRule); - - RuleQueryMapper mockRuleQueryMapper = mock(RuleQueryMapper.class); - RuleEntityParser mockRuleEntityParser = mock(RuleEntityParser.class); - ClusterService clusterService = mock(ClusterService.class); - ClusterState clusterState = mock(ClusterState.class); - Metadata metadata = mock(Metadata.class); - QueryBuilder queryBuilder = mock(QueryBuilder.class); - - when(clusterService.state()).thenReturn(clusterState); - when(clusterState.metadata()).thenReturn(metadata); - when(metadata.hasIndex(TEST_INDEX_NAME)).thenReturn(true); - when(mockRuleQueryMapper.from(any(GetRuleRequest.class))).thenReturn(queryBuilder); - when(queryBuilder.filter(any())).thenReturn(queryBuilder); - - when(mockRule.getAttributeMap()).thenReturn(Map.of(MOCK_RULE_ATTRIBUTE_ONE, Set.of(ATTRIBUTE_VALUE_ONE))); - when(mockRule.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); + public void testConcurrentCreateDuplicateRules() throws InterruptedException { + ExecutorService singleThreadExecutor = Executors.newSingleThreadExecutor(); + int threadCount = 10; + CountDownLatch latch = new CountDownLatch(threadCount); + Set storedAttributeMaps = ConcurrentHashMap.newKeySet(); - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - Client client = setUpMockClient(searchRequestBuilder); + CreateRuleRequest createRuleRequest = mock(CreateRuleRequest.class); + when(rule.getAttributeMap()).thenReturn(ATTRIBUTE_MAP); + when(rule.getFeatureType()).thenReturn(MockRuleFeatureType.INSTANCE); + when(createRuleRequest.getRule()).thenReturn(rule); RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( TEST_INDEX_NAME, client, clusterService, MAX_VALUES_PER_PAGE, - mockRuleEntityParser, - mockRuleQueryMapper - ); + ruleEntityParser, + ruleQueryMapper + ) { + @Override + public void createRule(CreateRuleRequest request, ActionListener listener) { + singleThreadExecutor.execute(() -> { + Rule rule = request.getRule(); + validateNoDuplicateRule(rule, new ActionListener() { + @Override + public void onResponse(Void unused) { + synchronized (storedAttributeMaps) { + storedAttributeMaps.add(MOCK_RULE_ATTRIBUTE_ONE.getName()); + } + listener.onResponse(new CreateRuleResponse("fake-id", rule)); + latch.countDown(); + } + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + latch.countDown(); + } + }); + }); + } + + public void validateNoDuplicateRule(Rule rule, ActionListener listener) { + synchronized (storedAttributeMaps) { + if (storedAttributeMaps.contains(MOCK_RULE_ATTRIBUTE_ONE.getName())) { + listener.onFailure(new IllegalArgumentException("Duplicate rule exists with attribute map")); + } else { + listener.onResponse(null); + } + } + } + }; + + class TestListener implements ActionListener { + final AtomicInteger successCount = new AtomicInteger(); + final AtomicInteger failureCount = new AtomicInteger(); + final List failures = Collections.synchronizedList(new ArrayList<>()); + + @Override + public void onResponse(CreateRuleResponse response) { + successCount.incrementAndGet(); + } + + @Override + public void onFailure(Exception e) { + failureCount.incrementAndGet(); + failures.add(e); + } + } + TestListener testListener = new TestListener(); + + for (int i = 0; i < threadCount; i++) { + new Thread(() -> rulePersistenceService.createRule(createRuleRequest, testListener)).start(); + } + boolean completed = latch.await(10, TimeUnit.SECONDS); + singleThreadExecutor.shutdown(); + assertTrue("All create calls should complete", completed); + assertEquals(1, testListener.successCount.get()); + assertEquals(threadCount - 1, testListener.failureCount.get()); + for (Exception e : testListener.failures) { + assertTrue(e instanceof IllegalArgumentException); + assertTrue(e.getMessage().contains("Duplicate rule")); + } + } + + public void testCreateDuplicateRule() { + CreateRuleRequest createRuleRequest = mock(CreateRuleRequest.class); + when(createRuleRequest.getRule()).thenReturn(rule); + when(rule.getAttributeMap()).thenReturn(Map.of(MOCK_RULE_ATTRIBUTE_ONE, Set.of(ATTRIBUTE_VALUE_ONE))); + when(rule.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); SearchResponse searchResponse = mock(SearchResponse.class); SearchHit hit = new SearchHit(1); hit.sourceRef(new BytesArray(VALID_JSON)); SearchHits searchHits = new SearchHits(new SearchHit[] { hit }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1.0f); when(searchResponse.getHits()).thenReturn(searchHits); - - doAnswer((invocation) -> { - ActionListener actionListener = invocation.getArgument(0); - actionListener.onResponse(searchResponse); - return null; - }).when(searchRequestBuilder).execute(any(ActionListener.class)); + when(searchRequestBuilder.get()).thenReturn(searchResponse); ActionListener listener = mock(ActionListener.class); - when(mockRuleEntityParser.parse(any(String.class))).thenReturn(mockRule); + when(ruleEntityParser.parse(any(String.class))).thenReturn(rule); rulePersistenceService.createRule(createRuleRequest, listener); ArgumentCaptor failureCaptor = ArgumentCaptor.forClass(Exception.class); verify(listener).onFailure(failureCaptor.capture()); @@ -179,90 +252,34 @@ public void testGetRuleByIdSuccess() { GetRuleRequest getRuleRequest = mock(GetRuleRequest.class); when(getRuleRequest.getId()).thenReturn(_ID_ONE); when(getRuleRequest.getAttributeFilters()).thenReturn(new HashMap<>()); - QueryBuilder queryBuilder = mock(QueryBuilder.class); - RuleQueryMapper mockRuleQueryMapper = mock(RuleQueryMapper.class); - RuleEntityParser mockRuleEntityParser = mock(RuleEntityParser.class); - Rule mockRule = mock(Rule.class); - - when(mockRuleEntityParser.parse(anyString())).thenReturn(mockRule); - when(mockRuleQueryMapper.from(getRuleRequest)).thenReturn(queryBuilder); - when(queryBuilder.filter(any())).thenReturn(queryBuilder); - - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - Client client = setUpMockClient(searchRequestBuilder); - ClusterService clusterService = mock(ClusterService.class); - - RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( - TEST_INDEX_NAME, - client, - clusterService, - MAX_VALUES_PER_PAGE, - mockRuleEntityParser, - mockRuleQueryMapper - ); + when(getRuleRequest.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); SearchResponse searchResponse = mock(SearchResponse.class); - SearchHits searchHits = new SearchHits(new SearchHit[] { new SearchHit(1) }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1.0f); + SearchHit searchHit = new SearchHit(1); + searchHit.sourceRef(new BytesArray(VALID_JSON)); + SearchHits searchHits = new SearchHits(new SearchHit[] { searchHit }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1.0f); when(searchResponse.getHits()).thenReturn(searchHits); - SearchHit hit = searchHits.getHits()[0]; - hit.sourceRef(new BytesArray(VALID_JSON)); + when(searchRequestBuilder.get()).thenReturn(searchResponse); ActionListener listener = mock(ActionListener.class); - - doAnswer((invocation) -> { - ActionListener actionListener = invocation.getArgument(0); - actionListener.onResponse(searchResponse); - return null; - }).when(searchRequestBuilder).execute(any(ActionListener.class)); - - when(getRuleRequest.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); - rulePersistenceService.getRule(getRuleRequest, listener); - ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(GetRuleResponse.class); + rulePersistenceService.getRule(getRuleRequest, listener); verify(listener).onResponse(responseCaptor.capture()); GetRuleResponse response = responseCaptor.getValue(); - assertEquals(response.getRules().size(), 1); + assertEquals(1, response.getRules().size()); } public void testGetRuleByIdNotFound() { GetRuleRequest getRuleRequest = mock(GetRuleRequest.class); when(getRuleRequest.getId()).thenReturn(_ID_ONE); - QueryBuilder queryBuilder = mock(QueryBuilder.class); - RuleQueryMapper mockRuleQueryMapper = mock(RuleQueryMapper.class); - RuleEntityParser mockRuleEntityParser = mock(RuleEntityParser.class); - Rule mockRule = mock(Rule.class); - - when(mockRuleEntityParser.parse(anyString())).thenReturn(mockRule); - when(mockRuleQueryMapper.from(getRuleRequest)).thenReturn(queryBuilder); - when(queryBuilder.filter(any())).thenReturn(queryBuilder); - - SearchRequestBuilder searchRequestBuilder = mock(SearchRequestBuilder.class); - Client client = setUpMockClient(searchRequestBuilder); - ClusterService clusterService = mock(ClusterService.class); - - RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( - TEST_INDEX_NAME, - client, - clusterService, - MAX_VALUES_PER_PAGE, - mockRuleEntityParser, - mockRuleQueryMapper - ); + when(getRuleRequest.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); SearchResponse searchResponse = mock(SearchResponse.class); + when(searchRequestBuilder.get()).thenReturn(searchResponse); when(searchResponse.getHits()).thenReturn(new SearchHits(new SearchHit[] {}, new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1.0f)); - ActionListener listener = mock(ActionListener.class); - doAnswer(invocationOnMock -> { - ActionListener actionListener = invocationOnMock.getArgument(0); - actionListener.onResponse(searchResponse); - return null; - }).when(searchRequestBuilder).execute(any(ActionListener.class)); - - when(getRuleRequest.getFeatureType()).thenReturn(RuleTestUtils.MockRuleFeatureType.INSTANCE); rulePersistenceService.getRule(getRuleRequest, listener); - ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); verify(listener).onFailure(exceptionCaptor.capture()); Exception exception = exceptionCaptor.getValue(); @@ -292,28 +309,15 @@ private Client setUpMockClient(SearchRequestBuilder searchRequestBuilder) { public void testDeleteRule_successful() { String ruleId = "test-rule-id"; DeleteRuleRequest request = new DeleteRuleRequest(ruleId, RuleTestUtils.MockRuleFeatureType.INSTANCE); - - Client client = mock(Client.class); - ClusterService clusterService = mock(ClusterService.class); ThreadPool threadPool = mock(ThreadPool.class); when(client.threadPool()).thenReturn(threadPool); when(threadPool.getThreadContext()).thenReturn(new ThreadContext(Settings.EMPTY)); - RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( - TEST_INDEX_NAME, - client, - clusterService, - MAX_VALUES_PER_PAGE, - mock(RuleEntityParser.class), - mock(RuleQueryMapper.class) - ); - ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(DeleteRequest.class); ArgumentCaptor> listenerCaptor = ArgumentCaptor.forClass( ActionListener.class ); - @SuppressWarnings("unchecked") ActionListener listener = mock(ActionListener.class); rulePersistenceService.deleteRule(request, listener); @@ -332,28 +336,15 @@ public void testDeleteRule_successful() { public void testDeleteRule_notFound() { String ruleId = "missing-rule-id"; DeleteRuleRequest request = new DeleteRuleRequest(ruleId, RuleTestUtils.MockRuleFeatureType.INSTANCE); - - Client client = mock(Client.class); - ClusterService clusterService = mock(ClusterService.class); ThreadPool threadPool = mock(ThreadPool.class); when(client.threadPool()).thenReturn(threadPool); when(threadPool.getThreadContext()).thenReturn(new ThreadContext(Settings.EMPTY)); - RulePersistenceService rulePersistenceService = new IndexStoredRulePersistenceService( - TEST_INDEX_NAME, - client, - clusterService, - MAX_VALUES_PER_PAGE, - mock(RuleEntityParser.class), - mock(RuleQueryMapper.class) - ); - ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(DeleteRequest.class); ArgumentCaptor> listenerCaptor = ArgumentCaptor.forClass( ActionListener.class ); - @SuppressWarnings("unchecked") ActionListener listener = mock(ActionListener.class); rulePersistenceService.deleteRule(request, listener); diff --git a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java index caa4aafdfca51..ac45b3e784446 100644 --- a/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java +++ b/modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java @@ -11,10 +11,14 @@ import org.opensearch.rule.autotagging.Attribute; import org.opensearch.rule.autotagging.AutoTaggingRegistry; import org.opensearch.rule.autotagging.FeatureType; +import org.opensearch.rule.autotagging.Rule; import java.util.Map; import java.util.Set; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + public class RuleTestUtils { public static final String _ID_ONE = "AgfUO5Ja9yfvhdONlYi3TQ=="; public static final String ATTRIBUTE_VALUE_ONE = "mock_attribute_one"; @@ -22,12 +26,63 @@ public class RuleTestUtils { public static final String DESCRIPTION_ONE = "description_1"; public static final String FEATURE_TYPE_NAME = "mock_feature_type"; public static final String TEST_INDEX_NAME = ".test_index_for_rule"; + public static final String INVALID_ATTRIBUTE = "invalid_attribute"; + + public static final String SEARCH_AFTER = "search_after"; + public static final String _ID_TWO = "G5iIq84j7eK1qIAAAAIH53=1"; + public static final String FEATURE_VALUE_ONE = "feature_value_one"; + public static final String FEATURE_VALUE_TWO = "feature_value_two"; + public static final String DESCRIPTION_TWO = "description_2"; + public static final String TIMESTAMP_ONE = "2024-01-26T08:58:57.558Z"; + public static final String TIMESTAMP_TWO = "2023-01-26T08:58:57.558Z"; + public static final Map> ATTRIBUTE_MAP = Map.of( MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE, Set.of(ATTRIBUTE_VALUE_ONE) ); - public static final String INVALID_ATTRIBUTE = "invalid_attribute"; + public static final Rule ruleOne = Rule.builder() + .description(DESCRIPTION_ONE) + .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_ONE) + .attributeMap(ATTRIBUTE_MAP) + .updatedAt(TIMESTAMP_ONE) + .build(); + + public static final Rule ruleTwo = Rule.builder() + .description(DESCRIPTION_TWO) + .featureType(RuleTestUtils.MockRuleFeatureType.INSTANCE) + .featureValue(FEATURE_VALUE_TWO) + .attributeMap(Map.of(RuleTestUtils.MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO, Set.of(ATTRIBUTE_VALUE_TWO))) + .updatedAt(TIMESTAMP_TWO) + .build(); + + public static Map ruleMap() { + return Map.of(_ID_ONE, ruleOne, _ID_TWO, ruleTwo); + } + + public static void assertEqualRules(Map mapOne, Map mapTwo, boolean ruleUpdated) { + assertEquals(mapOne.size(), mapTwo.size()); + for (Map.Entry entry : mapOne.entrySet()) { + String id = entry.getKey(); + assertTrue(mapTwo.containsKey(id)); + Rule one = mapOne.get(id); + Rule two = mapTwo.get(id); + assertEqualRule(one, two, ruleUpdated); + } + } + + public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) { + if (ruleUpdated) { + assertEquals(one.getDescription(), two.getDescription()); + assertEquals(one.getFeatureType(), two.getFeatureType()); + assertEquals(one.getFeatureValue(), two.getFeatureValue()); + assertEquals(one.getAttributeMap(), two.getAttributeMap()); + assertEquals(one.getAttributeMap(), two.getAttributeMap()); + } else { + assertEquals(one, two); + } + } public static class MockRuleFeatureType implements FeatureType { diff --git a/modules/autotagging-commons/spi/src/main/java/org/opensearch/rule/spi/RuleFrameworkExtension.java b/modules/autotagging-commons/spi/src/main/java/org/opensearch/rule/spi/RuleFrameworkExtension.java index 07443c31236a7..5c34bc29efdda 100644 --- a/modules/autotagging-commons/spi/src/main/java/org/opensearch/rule/spi/RuleFrameworkExtension.java +++ b/modules/autotagging-commons/spi/src/main/java/org/opensearch/rule/spi/RuleFrameworkExtension.java @@ -9,6 +9,7 @@ package org.opensearch.rule.spi; import org.opensearch.rule.RulePersistenceService; +import org.opensearch.rule.RuleRoutingService; import org.opensearch.rule.autotagging.FeatureType; import java.util.function.Supplier; @@ -24,6 +25,12 @@ public interface RuleFrameworkExtension { */ Supplier getRulePersistenceServiceSupplier(); + /** + * This method is used to flow implementation from consumer plugins into framework plugin + * @return the plugin specific implementation of RuleRoutingService + */ + Supplier getRuleRoutingServiceSupplier(); + /** * Flow implementation from consumer plugins into framework plugin * @return the specific implementation of FeatureType diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java index fbbe6e2759693..cd6197cf890f7 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleFrameworkPlugin.java @@ -68,6 +68,7 @@ public class RuleFrameworkPlugin extends Plugin implements ExtensiblePlugin, Act public RuleFrameworkPlugin() {} private final RulePersistenceServiceRegistry rulePersistenceServiceRegistry = new RulePersistenceServiceRegistry(); + private final RuleRoutingServiceRegistry ruleRoutingServiceRegistry = new RuleRoutingServiceRegistry(); private final List ruleFrameworkExtensions = new ArrayList<>(); @Override @@ -101,7 +102,10 @@ public List> getExecutorBuilders(Settings settings) { @Override public Collection createGuiceModules() { - return List.of(b -> { b.bind(RulePersistenceServiceRegistry.class).toInstance(rulePersistenceServiceRegistry); }); + return List.of(b -> { + b.bind(RulePersistenceServiceRegistry.class).toInstance(rulePersistenceServiceRegistry); + b.bind(RuleRoutingServiceRegistry.class).toInstance(ruleRoutingServiceRegistry); + }); } @Override @@ -113,5 +117,7 @@ private void consumeFrameworkExtension(RuleFrameworkExtension ruleFrameworkExten FeatureType featureType = ruleFrameworkExtension.getFeatureTypeSupplier().get(); AutoTaggingRegistry.registerFeatureType(featureType); rulePersistenceServiceRegistry.register(featureType, ruleFrameworkExtension.getRulePersistenceServiceSupplier().get()); + ruleRoutingServiceRegistry.register(featureType, ruleFrameworkExtension.getRuleRoutingServiceSupplier().get()); + } } diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/RulePersistenceServiceRegistry.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/RulePersistenceServiceRegistry.java index 57f47f1a6ad0f..1479d1d3ded21 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/RulePersistenceServiceRegistry.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/RulePersistenceServiceRegistry.java @@ -41,7 +41,7 @@ public void register(FeatureType featureType, RulePersistenceService rulePersist */ public RulePersistenceService getRulePersistenceService(FeatureType featureType) { if (!rulePersistenceServices.containsKey(featureType.getName())) { - throw new IllegalArgumentException("Unknown feature type: " + featureType.getName()); + throw new IllegalArgumentException("Unknown feature type for persistence service: " + featureType.getName()); } return rulePersistenceServices.get(featureType.getName()); } diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleRoutingServiceRegistry.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleRoutingServiceRegistry.java new file mode 100644 index 0000000000000..d8bc29640a21e --- /dev/null +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/RuleRoutingServiceRegistry.java @@ -0,0 +1,48 @@ +/* + * 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.rule; + +import org.opensearch.rule.autotagging.FeatureType; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * This class manages implementations of {@link RuleRoutingService} + */ +public class RuleRoutingServiceRegistry { + private final Map ruleRoutingServices = new ConcurrentHashMap<>(); + + /** + * default constructor + */ + public RuleRoutingServiceRegistry() {} + + /** + * This method is used to register the concrete implementations of RuleRoutingService + * @param featureType + * @param ruleRoutingService + */ + public void register(FeatureType featureType, RuleRoutingService ruleRoutingService) { + if (ruleRoutingServices.put(featureType.getName(), ruleRoutingService) != null) { + throw new IllegalArgumentException("Duplicate rule routing service: " + featureType.getName()); + } + } + + /** + * It is used to get feature type specific {@link RuleRoutingService} implementation + * @param featureType - the type of feature to retrieve the routing service for + */ + public RuleRoutingService getRuleRoutingService(FeatureType featureType) { + if (!ruleRoutingServices.containsKey(featureType.getName())) { + throw new IllegalArgumentException("Unknown feature type for routing service: " + featureType.getName()); + } + return ruleRoutingServices.get(featureType.getName()); + } +} diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java index 61c0a15d7ed96..f808cf7427cbc 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/action/TransportCreateRuleAction.java @@ -8,38 +8,23 @@ package org.opensearch.rule.action; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.opensearch.ExceptionsHelper; -import org.opensearch.ResourceAlreadyExistsException; -import org.opensearch.action.ActionListenerResponseHandler; -import org.opensearch.action.admin.indices.create.CreateIndexRequest; -import org.opensearch.action.admin.indices.create.CreateIndexResponse; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.TransportAction; -import org.opensearch.cluster.ClusterState; -import org.opensearch.cluster.node.DiscoveryNode; -import org.opensearch.cluster.routing.IndexRoutingTable; -import org.opensearch.cluster.routing.ShardRouting; -import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; -import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; import org.opensearch.rule.CreateRuleRequest; import org.opensearch.rule.CreateRuleResponse; import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.RulePersistenceServiceRegistry; -import org.opensearch.rule.service.IndexStoredRulePersistenceService; +import org.opensearch.rule.RuleRoutingServiceRegistry; import org.opensearch.tasks.Task; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportChannel; import org.opensearch.transport.TransportException; import org.opensearch.transport.TransportRequestHandler; import org.opensearch.transport.TransportService; -import org.opensearch.transport.client.Client; import java.io.IOException; -import java.util.Map; import static org.opensearch.rule.RuleFrameworkPlugin.RULE_THREAD_POOL_NAME; @@ -48,43 +33,34 @@ * @opensearch.experimental */ public class TransportCreateRuleAction extends TransportAction { - private final TransportService transportService; - private final ClusterService clusterService; private final ThreadPool threadPool; + private final RuleRoutingServiceRegistry ruleRoutingServiceRegistry; private final RulePersistenceServiceRegistry rulePersistenceServiceRegistry; - private final Client client; - private static final Logger logger = LogManager.getLogger(TransportCreateRuleAction.class); - private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); /** * Constructor for TransportCreateRuleAction - * - * @param threadPool - {@link ThreadPool} object - * @param client - {@link Client} object * @param transportService - a {@link TransportService} object - * @param clusterService - a {@link ClusterService} object * @param actionFilters - a {@link ActionFilters} object + * @param threadPool - a {@link ThreadPool} object * @param rulePersistenceServiceRegistry - a {@link RulePersistenceServiceRegistry} object + * @param ruleRoutingServiceRegistry - a {@link RuleRoutingServiceRegistry} object */ @Inject public TransportCreateRuleAction( - Client client, TransportService transportService, - ClusterService clusterService, ThreadPool threadPool, ActionFilters actionFilters, - RulePersistenceServiceRegistry rulePersistenceServiceRegistry + RulePersistenceServiceRegistry rulePersistenceServiceRegistry, + RuleRoutingServiceRegistry ruleRoutingServiceRegistry ) { super(CreateRuleAction.NAME, actionFilters, transportService.getTaskManager()); - this.client = client; - this.transportService = transportService; - this.clusterService = clusterService; + this.ruleRoutingServiceRegistry = ruleRoutingServiceRegistry; this.threadPool = threadPool; this.rulePersistenceServiceRegistry = rulePersistenceServiceRegistry; transportService.registerRequestHandler( CreateRuleAction.NAME, - ThreadPool.Names.WRITE, + ThreadPool.Names.SAME, CreateRuleRequest::new, new TransportRequestHandler() { @Override @@ -111,96 +87,7 @@ public void messageReceived(CreateRuleRequest request, TransportChannel channel, @Override protected void doExecute(Task task, CreateRuleRequest request, ActionListener listener) { - RulePersistenceService persistenceService = rulePersistenceServiceRegistry.getRulePersistenceService( - request.getRule().getFeatureType() - ); - if (!(persistenceService instanceof IndexStoredRulePersistenceService indexStoredRulePersistenceService)) { - executeLocally(request, listener); - return; - } - - String indexName = indexStoredRulePersistenceService.getIndexName(); - try (ThreadContext.StoredContext ctx = client.threadPool().getThreadContext().stashContext()) { - if (!clusterService.state().metadata().hasIndex(indexName)) { - createIndex(() -> routeRequest(request, listener, indexName), indexName, listener); - return; - } - } - - routeRequest(request, listener, indexName); - } - - /** - * Creates the backing index if it does not exist, then runs the given success callback. - * @param onSuccess callback to run after successful index creation - * @param indexName the name of the index to create - * @param listener listener to handle failures - */ - private void createIndex(Runnable onSuccess, String indexName, ActionListener listener) { - final CreateIndexRequest request = new CreateIndexRequest(indexName).settings(indexSettings); - client.admin().indices().create(request, new ActionListener<>() { - @Override - public void onResponse(CreateIndexResponse response) { - if (!response.isAcknowledged()) { - listener.onFailure(new IllegalStateException(indexName + " index creation not acknowledged")); - } else { - onSuccess.run(); - } - } - - @Override - public void onFailure(Exception e) { - Throwable cause = ExceptionsHelper.unwrapCause(e); - if (cause instanceof ResourceAlreadyExistsException) { - onSuccess.run(); - } else { - listener.onFailure(e); - } - } - }); - } - - /** - * Routes the CreateRuleRequest to the primary shard node for the given index. - * Executes locally if the current node is the primary. - * @param request the CreateRuleRequest - * @param listener listener to handle response or failure - * @param indexName the index name used to find the primary shard node - */ - private void routeRequest(CreateRuleRequest request, ActionListener listener, String indexName) { - DiscoveryNode primaryNode = getPrimaryShardNode(indexName); - if (primaryNode == null) { - listener.onFailure(new IllegalStateException("Primary node for index [" + indexName + "] not found")); - return; - } - - if (transportService.getLocalNode().equals(primaryNode)) { - executeLocally(request, listener); - } else { - transportService.sendRequest( - primaryNode, - CreateRuleAction.NAME, - request, - new ActionListenerResponseHandler<>(listener, CreateRuleResponse::new) - ); - } - } - - /** - * Retrieves the discovery node that holds the primary shard for the given index. - * @param indexName the index name - */ - private DiscoveryNode getPrimaryShardNode(String indexName) { - ClusterState state = clusterService.state(); - IndexRoutingTable indexRoutingTable = state.getRoutingTable().index(indexName); - if (indexRoutingTable == null) { - return null; - } - ShardRouting primaryShard = indexRoutingTable.shard(0).primaryShard(); - if (primaryShard == null || !primaryShard.assignedToNode()) { - return null; - } - return state.nodes().get(primaryShard.currentNodeId()); + ruleRoutingServiceRegistry.getRuleRoutingService(request.getRule().getFeatureType()).handleCreateRuleRequest(request, listener); } /** @@ -210,10 +97,10 @@ private DiscoveryNode getPrimaryShardNode(String indexName) { */ private void executeLocally(CreateRuleRequest request, ActionListener listener) { threadPool.executor(RULE_THREAD_POOL_NAME).execute(() -> { - RulePersistenceService persistenceService = rulePersistenceServiceRegistry.getRulePersistenceService( + final RulePersistenceService rulePersistenceService = rulePersistenceServiceRegistry.getRulePersistenceService( request.getRule().getFeatureType() ); - persistenceService.createRule(request, listener); + rulePersistenceService.createRule(request, listener); }); } } diff --git a/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java b/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java index 4a8254924db2e..7a5f45e95a0ba 100644 --- a/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java +++ b/modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestCreateRuleAction.java @@ -58,6 +58,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli try (XContentParser parser = request.contentParser()) { Builder builder = Builder.fromXContent(parser, featureType); CreateRuleRequest createRuleRequest = new CreateRuleRequest(builder.updatedAt(Instant.now().toString()).build()); + return channel -> client.execute(CreateRuleAction.INSTANCE, createRuleRequest, createRuleResponse(channel)); } } diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/RulePersistenceServiceRegistryTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/RulePersistenceServiceRegistryTests.java index aecdef198d39f..57cc35ac4ac24 100644 --- a/modules/autotagging-commons/src/test/java/org/opensearch/rule/RulePersistenceServiceRegistryTests.java +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/RulePersistenceServiceRegistryTests.java @@ -43,6 +43,6 @@ public void testGetRulePersistenceService_UnknownFeature() { IllegalArgumentException.class, () -> registry.getRulePersistenceService(mockFeatureType) ); - assertTrue(ex.getMessage().contains("Unknown feature type: unknown_feature")); + assertTrue(ex.getMessage().contains("Unknown feature type")); } } diff --git a/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java b/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java index 8b063cb0afbe3..fa413044c5efc 100644 --- a/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java +++ b/modules/autotagging-commons/src/test/java/org/opensearch/rule/action/TransportCreateRuleActionTests.java @@ -9,55 +9,52 @@ package org.opensearch.rule.action; import org.opensearch.action.support.ActionFilters; -import org.opensearch.cluster.service.ClusterService; import org.opensearch.core.action.ActionListener; import org.opensearch.rule.CreateRuleRequest; import org.opensearch.rule.CreateRuleResponse; -import org.opensearch.rule.RulePersistenceService; import org.opensearch.rule.RulePersistenceServiceRegistry; +import org.opensearch.rule.RuleRoutingService; +import org.opensearch.rule.RuleRoutingServiceRegistry; import org.opensearch.rule.autotagging.FeatureType; import org.opensearch.rule.autotagging.Rule; import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; -import org.opensearch.transport.client.Client; import java.util.concurrent.ExecutorService; -import static org.opensearch.rule.RuleFrameworkPlugin.RULE_THREAD_POOL_NAME; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @SuppressWarnings("unchecked") public class TransportCreateRuleActionTests extends OpenSearchTestCase { private TransportService transportService; - private ClusterService clusterService; private ThreadPool threadPool; - private RulePersistenceServiceRegistry registry; - private Client client; private ActionFilters actionFilters; private TransportCreateRuleAction action; private FeatureType mockFeatureType; + private RuleRoutingServiceRegistry routingRegistry; + private RulePersistenceServiceRegistry persistenceRegistry; + private RuleRoutingService mockRoutingService; private final String testIndexName = "test-index"; public void setUp() throws Exception { super.setUp(); transportService = mock(TransportService.class); - clusterService = mock(ClusterService.class); threadPool = mock(ThreadPool.class); - registry = mock(RulePersistenceServiceRegistry.class); - client = mock(Client.class); actionFilters = mock(ActionFilters.class); mockFeatureType = mock(FeatureType.class); - RulePersistenceServiceRegistry registry = new RulePersistenceServiceRegistry(); + routingRegistry = mock(RuleRoutingServiceRegistry.class); + persistenceRegistry = mock(RulePersistenceServiceRegistry.class); when(mockFeatureType.getName()).thenReturn("test_feature"); - RulePersistenceService mockService = mock(RulePersistenceService.class); - registry.register(mockFeatureType, mockService); + mockRoutingService = mock(RuleRoutingService.class); + routingRegistry.register(mockFeatureType, mockRoutingService); ExecutorService executorService = mock(ExecutorService.class); doAnswer(invocation -> { @@ -66,17 +63,18 @@ public void setUp() throws Exception { return null; }).when(executorService).execute(any()); when(threadPool.executor(any())).thenReturn(executorService); - action = new TransportCreateRuleAction(client, transportService, clusterService, threadPool, actionFilters, registry); + action = new TransportCreateRuleAction(transportService, threadPool, actionFilters, persistenceRegistry, routingRegistry); } public void testExecution() { IndexStoredRulePersistenceService persistenceService = mock(IndexStoredRulePersistenceService.class); - when(registry.getRulePersistenceService(mockFeatureType)).thenReturn(persistenceService); + when(persistenceRegistry.getRulePersistenceService(mockFeatureType)).thenReturn(persistenceService); + when(routingRegistry.getRuleRoutingService(mockFeatureType)).thenReturn(mockRoutingService); Rule rule = mock(Rule.class); when(rule.getFeatureType()).thenReturn(mockFeatureType); CreateRuleRequest request = new CreateRuleRequest(rule); ActionListener listener = mock(ActionListener.class); action.doExecute(null, request, listener); - verify(threadPool).executor(RULE_THREAD_POOL_NAME); + verify(routingRegistry, times(1)).getRuleRoutingService(request.getRule().getFeatureType()); } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index e967bc1138083..3ea3f5548ef5c 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -14,6 +14,7 @@ import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Module; +import org.opensearch.common.network.NetworkService; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.IndexScopedSettings; import org.opensearch.common.settings.Setting; @@ -22,6 +23,7 @@ import org.opensearch.core.action.ActionResponse; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.discovery.SeedHostsProvider; import org.opensearch.env.Environment; import org.opensearch.env.NodeEnvironment; import org.opensearch.indices.SystemIndexDescriptor; @@ -39,8 +41,10 @@ import org.opensearch.plugin.wlm.rest.RestUpdateWorkloadGroupAction; import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType; import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureValueValidator; +import org.opensearch.plugin.wlm.rule.WorkloadGroupRuleRoutingService; import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; import org.opensearch.plugins.ActionPlugin; +import org.opensearch.plugins.DiscoveryPlugin; import org.opensearch.plugins.Plugin; import org.opensearch.plugins.SystemIndexPlugin; import org.opensearch.repositories.RepositoriesService; @@ -48,6 +52,7 @@ import org.opensearch.rest.RestHandler; import org.opensearch.rule.InMemoryRuleProcessingService; import org.opensearch.rule.RulePersistenceService; +import org.opensearch.rule.RuleRoutingService; import org.opensearch.rule.autotagging.FeatureType; import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.rule.spi.RuleFrameworkExtension; @@ -56,18 +61,20 @@ import org.opensearch.rule.storage.XContentRuleParser; import org.opensearch.script.ScriptService; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; import org.opensearch.transport.client.Client; import org.opensearch.watcher.ResourceWatcherService; import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.function.Supplier; /** * Plugin class for WorkloadManagement */ -public class WorkloadManagementPlugin extends Plugin implements ActionPlugin, SystemIndexPlugin, RuleFrameworkExtension { +public class WorkloadManagementPlugin extends Plugin implements ActionPlugin, SystemIndexPlugin, DiscoveryPlugin, RuleFrameworkExtension { /** * The name of the index where rules are stored. @@ -79,6 +86,7 @@ public class WorkloadManagementPlugin extends Plugin implements ActionPlugin, Sy public static final int MAX_RULES_PER_PAGE = 50; private static FeatureType featureType; private static RulePersistenceService rulePersistenceService; + private static RuleRoutingService ruleRoutingService; private AutoTaggingActionFilter autoTaggingActionFilter; @@ -110,14 +118,21 @@ public Collection createComponents( new XContentRuleParser(featureType), new IndexBasedRuleQueryMapper() ); + ruleRoutingService = new WorkloadGroupRuleRoutingService(client, clusterService); InMemoryRuleProcessingService ruleProcessingService = new InMemoryRuleProcessingService( - WorkloadGroupFeatureType.getInstance(), + featureType, DefaultAttributeValueStore::new ); autoTaggingActionFilter = new AutoTaggingActionFilter(ruleProcessingService, threadPool); return Collections.emptyList(); } + @Override + public Map> getSeedHostProviders(TransportService transportService, NetworkService networkService) { + ((WorkloadGroupRuleRoutingService) ruleRoutingService).setTransportService(transportService); + return Collections.emptyMap(); + } + @Override public List getActionFilters() { return List.of(autoTaggingActionFilter); @@ -135,7 +150,7 @@ public List getActionFilters() { @Override public Collection getSystemIndexDescriptors(Settings settings) { - return List.of(new SystemIndexDescriptor(INDEX_NAME, "System index used for storing rules")); + return List.of(new SystemIndexDescriptor(INDEX_NAME, "System index used for storing workload_group rules")); } @Override @@ -171,6 +186,11 @@ public Supplier getRulePersistenceServiceSupplier() { return () -> rulePersistenceService; } + @Override + public Supplier getRuleRoutingServiceSupplier() { + return () -> ruleRoutingService; + } + @Override public Supplier getFeatureTypeSupplier() { return () -> featureType; diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupRuleRoutingService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupRuleRoutingService.java new file mode 100644 index 0000000000000..a70482eb40186 --- /dev/null +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/WorkloadGroupRuleRoutingService.java @@ -0,0 +1,143 @@ +/* + * 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.rule; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.ExceptionsHelper; +import org.opensearch.ResourceAlreadyExistsException; +import org.opensearch.action.ActionListenerResponseHandler; +import org.opensearch.action.admin.indices.create.CreateIndexRequest; +import org.opensearch.action.admin.indices.create.CreateIndexResponse; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.routing.IndexShardRoutingTable; +import org.opensearch.cluster.routing.ShardRouting; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.action.ActionListener; +import org.opensearch.plugin.wlm.WorkloadManagementPlugin; +import org.opensearch.rule.CreateRuleRequest; +import org.opensearch.rule.CreateRuleResponse; +import org.opensearch.rule.RuleRoutingService; +import org.opensearch.rule.action.CreateRuleAction; +import org.opensearch.transport.TransportService; +import org.opensearch.transport.client.Client; + +import java.util.Map; +import java.util.Optional; + +/** + * Service responsible for routing CreateRule requests to the correct node based on primary shard ownership. + * @opensearch.experimental + */ +public class WorkloadGroupRuleRoutingService implements RuleRoutingService { + private final Client client; + private final ClusterService clusterService; + private TransportService transportService; + private static final Logger logger = LogManager.getLogger(WorkloadGroupRuleRoutingService.class); + private static final Map indexSettings = Map.of("index.number_of_shards", 1, "index.auto_expand_replicas", "0-all"); + + /** + * Constructor for WorkloadGroupRuleRoutingService + * @param client + * @param clusterService + */ + public WorkloadGroupRuleRoutingService(Client client, ClusterService clusterService) { + this.client = client; + this.clusterService = clusterService; + } + + /** + * Set {@link TransportService} for WorkloadGroupRuleRoutingService + * @param transportService + */ + public void setTransportService(TransportService transportService) { + this.transportService = transportService; + } + + @Override + public void handleCreateRuleRequest(CreateRuleRequest request, ActionListener listener) { + String indexName = WorkloadManagementPlugin.INDEX_NAME; + + try (ThreadContext.StoredContext ctx = client.threadPool().getThreadContext().stashContext()) { + if (clusterService.state().metadata().hasIndex(indexName)) { + routeRequest(request, listener, indexName); + return; + } + createIndex(indexName, new ActionListener<>() { + @Override + public void onResponse(CreateIndexResponse response) { + if (!response.isAcknowledged()) { + logger.error("Failed to create index " + indexName); + listener.onFailure(new IllegalStateException(indexName + " index creation not acknowledged")); + } else { + routeRequest(request, listener, indexName); + } + } + + @Override + public void onFailure(Exception e) { + Throwable cause = ExceptionsHelper.unwrapCause(e); + if (cause instanceof ResourceAlreadyExistsException) { + routeRequest(request, listener, indexName); + } else { + logger.error("Failed to create index {}: {}", indexName, e.getMessage()); + listener.onFailure(e); + } + } + }); + } + } + + /** + * Creates the backing index if it does not exist, then runs the given success callback. + * @param indexName the name of the index to create + * @param listener listener to handle failures + */ + private void createIndex(String indexName, ActionListener listener) { + final CreateIndexRequest createRequest = new CreateIndexRequest(indexName).settings(indexSettings); + client.admin().indices().create(createRequest, listener); + } + + /** + * Routes the CreateRuleRequest to the primary shard node for the given index. + * Executes locally if the current node is the primary. + * @param request the CreateRuleRequest + * @param listener listener to handle response or failure + * @param indexName the index name used to find the primary shard node + */ + private void routeRequest(CreateRuleRequest request, ActionListener listener, String indexName) { + Optional optionalPrimaryNode = getPrimaryShardNode(indexName); + if (optionalPrimaryNode.isEmpty()) { + listener.onFailure(new IllegalStateException("Primary node for index [" + indexName + "] not found")); + return; + } + DiscoveryNode primaryNode = optionalPrimaryNode.get(); + transportService.sendRequest( + primaryNode, + CreateRuleAction.NAME, + request, + new ActionListenerResponseHandler<>(listener, CreateRuleResponse::new) + ); + } + + /** + * Retrieves the discovery node that holds the primary shard for the given index. + * @param indexName the index name + */ + private Optional getPrimaryShardNode(String indexName) { + ClusterState state = clusterService.state(); + return Optional.ofNullable(state.getRoutingTable().index(indexName)) + .map(table -> table.shard(0)) + .map(IndexShardRoutingTable::primaryShard) + .filter(ShardRouting::assignedToNode) + .map(shard -> state.nodes().get(shard.currentNodeId())); + } +} diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/AutoTaggingActionFilterTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/AutoTaggingActionFilterTests.java index ec458bfb9b9f1..45802c28ee458 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/AutoTaggingActionFilterTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/AutoTaggingActionFilterTests.java @@ -92,9 +92,6 @@ public String getName() { public Map getAllowedAttributesRegistry() { return Map.of("test_attribute", TestAttribute.TEST_ATTRIBUTE); } - - @Override - public void registerFeatureType() {} } public enum TestAttribute implements Attribute { diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java index 201dad40cc439..a55e345fd56da 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/WorkloadGroupFeatureTypeTests.java @@ -8,6 +8,7 @@ package org.opensearch.plugin.wlm.rule; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.rule.RuleAttribute; import org.opensearch.rule.autotagging.Attribute; import org.opensearch.rule.autotagging.AutoTaggingRegistry; @@ -15,8 +16,10 @@ import java.util.Map; +import static org.mockito.Mockito.mock; + public class WorkloadGroupFeatureTypeTests extends OpenSearchTestCase { - WorkloadGroupFeatureType featureType = new WorkloadGroupFeatureType(new WorkloadGroupFeatureValueValidator(null)); + WorkloadGroupFeatureType featureType = new WorkloadGroupFeatureType(new WorkloadGroupFeatureValueValidator(mock(ClusterService.class))); public void testGetName_returnsCorrectName() { assertEquals("workload_group", featureType.getName()); diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index 25a43aa635127..22ccfb6e438df 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -1362,6 +1362,7 @@ protected Node(final Environment initialEnvironment, Collection clas clusterManagerMetrics, remoteClusterStateService ); + final SearchPipelineService searchPipelineService = new SearchPipelineService( clusterService, threadPool,