rebase from main after the schema merged

Signed-off-by: Ruirui Zhang <[email protected]>

Author: ruai0511

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/querygroup/rest/RestDeleteQueryGroupAction.java

@@ -8,8 +8,8 @@
 
 package org.opensearch.plugin.wlm.querygroup.rest;
 
-import org.opensearch.plugin.wlm.action.DeleteQueryGroupAction;
-import org.opensearch.plugin.wlm.action.DeleteQueryGroupRequest;
+import org.opensearch.plugin.wlm.querygroup.action.DeleteQueryGroupAction;
+import org.opensearch.plugin.wlm.querygroup.action.DeleteQueryGroupRequest;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.RestRequest;
 import org.opensearch.rest.action.RestToXContentListener;

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,24 +23,14 @@ public enum QueryGroupAttribute implements Attribute {
 
     QueryGroupAttribute(String name) {
         this.name = name;
-    }
-
-    static {
-        for (QueryGroupAttribute attr : QueryGroupAttribute.values()) {
-            attr.registerAttribute();
-        }
+        validateAttribute();
     }
 
     @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)) {
@@ -48,4 +39,12 @@ public static QueryGroupAttribute fromName(String name) {
         }
         throw new IllegalArgumentException("Unknown QueryGroupAttribute: " + name);
     }
+
+    public static Map<String, Attribute> toMap() {
+        Map<String, Attribute> map = new HashMap<>();
+        for (QueryGroupAttribute attr : QueryGroupAttribute.values()) {
+            map.put(attr.getName(), attr);
+        }
+        return map;
+    }
 }

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<Attribute> ALLOWED_ATTRIBUTES = Set.of(QueryGroupAttribute.values());
+    private static final Map<String, Attribute> ALLOWED_ATTRIBUTES = QueryGroupAttribute.toMap();
 
     private QueryGroupFeatureType() {}
 
@@ -47,7 +43,7 @@ public int getMaxCharLengthPerAttributeValue() {
     }
 
     @Override
-    public Set<Attribute> getAllowedAttributes() {
+    public Map<String, Attribute> getAllowedAttributesRegistry() {
         return ALLOWED_ATTRIBUTES;
     }
 

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<Attribute, Set<String>> 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<Attribute, Set<String>> attributeFilters, String searchAfter) {
-        this._id = _id;
+    public GetRuleRequest(String id, Map<Attribute, Set<String>> attributeFilters, String searchAfter) {
+        this.id = id;
         this.attributeFilters = attributeFilters;
         this.searchAfter = searchAfter;
     }
@@ -50,8 +51,8 @@ public GetRuleRequest(String _id, Map<Attribute, Set<String>> 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;
     }
 
     /**

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<String, Rule<QueryGroupFeatureType>> rules;
+    private final Map<String, Rule> 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<String, Rule<QueryGroupFeatureType>> rules, String searchAfter, RestStatus restStatus) {
+    public GetRuleResponse(final Map<String, Rule> 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<String, Rule<QueryGroupFeatureType>> entry : rules.entrySet()) {
+        for (Map.Entry<String, Rule> 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<String, Rule<QueryGroupFeatureType>> getRules() {
+    public Map<String, Rule> getRules() {
         return rules;
     }
 

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<GetRuleResponse> listener) {
-        rulePersistenceService.getRule(request.get_id(), request.getAttributeFilters(), request.getSearchAfter(), listener);
+        rulePersistenceService.getRule(request.getId(), request.getAttributeFilters(), request.getSearchAfter(), listener);
     }
 }

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;