address comments

Signed-off-by: Ruirui Zhang <mariazrr@amazon.com>

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequest.java

@@ -11,9 +11,11 @@
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRequestValidationException;
 import org.opensearch.cluster.metadata.QueryGroup;
+import org.opensearch.common.UUIDs;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.xcontent.XContentParser;
+import org.joda.time.Instant;
 
 import java.io.IOException;
 
@@ -56,10 +58,8 @@ public CreateQueryGroupRequest(StreamInput in) throws IOException {
      * @param parser - A {@link XContentParser} object
      */
     public static CreateQueryGroupRequest fromXContent(XContentParser parser) throws IOException {
-        QueryGroup queryGroup1 = QueryGroup.fromXContent(parser);
-        // creating this queryGroup to ensure forceful creation of _id and updatedAt
-        QueryGroup queryGroup = new QueryGroup(queryGroup1.getName(), queryGroup1.getResiliencyMode(), queryGroup1.getResourceLimits());
-        return new CreateQueryGroupRequest(queryGroup);
+        QueryGroup.Builder builder = QueryGroup.Builder.fromXContent(parser);
+        return new CreateQueryGroupRequest(builder._id(UUIDs.randomBase64UUID()).updatedAt(Instant.now().getMillis()).build());
     }
 
     @Override

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java

@@ -138,7 +138,7 @@ public void clusterStateProcessed(String source, ClusterState oldState, ClusterS
      * @param queryGroup - the QueryGroup we're currently creating
      * @param currentClusterState - the cluster state before the update
      */
-    public ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final ClusterState currentClusterState) {
+    ClusterState saveQueryGroupInClusterState(final QueryGroup queryGroup, final ClusterState currentClusterState) {
         final Map<String, QueryGroup> existingQueryGroups = currentClusterState.metadata().queryGroups();
         String groupName = queryGroup.getName();
 

plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupRequestTests.java

@@ -22,6 +22,9 @@
 
 public class CreateQueryGroupRequestTests extends OpenSearchTestCase {
 
+    /**
+     * Test case to verify the serialization and deserialization of CreateQueryGroupRequest.
+     */
     public void testSerialization() throws IOException {
         CreateQueryGroupRequest request = new CreateQueryGroupRequest(queryGroupOne);
         BytesStreamOutput out = new BytesStreamOutput();

plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/CreateQueryGroupResponseTests.java

@@ -26,6 +26,9 @@
 
 public class CreateQueryGroupResponseTests extends OpenSearchTestCase {
 
+    /**
+     * Test case to verify the serialization and deserialization of CreateQueryGroupResponse.
+     */
     public void testSerialization() throws IOException {
         CreateQueryGroupResponse response = new CreateQueryGroupResponse(QueryGroupTestUtils.queryGroupOne, RestStatus.OK);
         BytesStreamOutput out = new BytesStreamOutput();
@@ -42,6 +45,9 @@ public void testSerialization() throws IOException {
         QueryGroupTestUtils.compareQueryGroups(listOne, listTwo);
     }
 
+    /**
+     * Test case to verify the toXContent method of CreateQueryGroupResponse.
+     */
     public void testToXContentCreateQueryGroup() throws IOException {
         XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint();
         CreateQueryGroupResponse response = new CreateQueryGroupResponse(QueryGroupTestUtils.queryGroupOne, RestStatus.OK);

plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java

@@ -43,6 +43,9 @@
 
 public class QueryGroupPersistenceServiceTests extends OpenSearchTestCase {
 
+    /**
+     * Test case to validate the creation logic of a single QueryGroup
+     */
     public void testCreateQueryGroup() {
         List<Object> setup = preparePersistenceServiceSetup(new HashMap<>());
         QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0);
@@ -58,6 +61,10 @@ public void testCreateQueryGroup() {
         compareQueryGroups(listOne, listTwo);
     }
 
+    /**
+     * Test case to validate the logic for adding a new QueryGroup to a cluster state that already contains
+     * an existing QueryGroup
+     */
     public void testCreateAnotherQueryGroup() {
         List<Object> setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne));
         QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0);
@@ -70,6 +77,9 @@ public void testCreateAnotherQueryGroup() {
         compareQueryGroups(queryGroupList(), new ArrayList<>(values));
     }
 
+    /**
+     * Test case to ensure the error is thrown when we try to create another QueryGroup with duplicate name
+     */
     public void testCreateQueryGroupDuplicateName() {
         List<Object> setup = preparePersistenceServiceSetup(Map.of(_ID_ONE, queryGroupOne));
         QueryGroupPersistenceService queryGroupPersistenceService1 = (QueryGroupPersistenceService) setup.get(0);
@@ -83,6 +93,10 @@ public void testCreateQueryGroupDuplicateName() {
         assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState));
     }
 
+    /**
+     * Test case to ensure the error is thrown when we try to create another QueryGroup that will make
+     * the total resource limits go above 1
+     */
     public void testCreateQueryGroupOverflowAllocation() {
         List<Object> setup = preparePersistenceServiceSetup(Map.of(_ID_TWO, queryGroupTwo));
         QueryGroup toCreate = builder().name(NAME_ONE)
@@ -96,6 +110,10 @@ public void testCreateQueryGroupOverflowAllocation() {
         assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState));
     }
 
+    /**
+     * Test case to ensure the error is thrown when we already have the max allowed number of QueryGroups, but
+     * we try to create another one
+     */
     public void testCreateQueryGroupOverflowCount() {
         QueryGroup toCreate = builder().name(NAME_NONE_EXISTED)
             ._id("W5iIqHyhgi4K1qIAAAAIHw==")

server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java

@@ -148,56 +148,7 @@ public XContentBuilder toXContent(final XContentBuilder builder, final Params pa
     }
 
     public static QueryGroup fromXContent(final XContentParser parser) throws IOException {
-        if (parser.currentToken() == null) { // fresh parser? move to the first token
-            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() + "]");
-        }
-
-        String fieldName = "";
-        // Map to hold resources
-        final Map<ResourceType, Double> resourceLimits = new HashMap<>();
-        while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
-            if (token == XContentParser.Token.FIELD_NAME) {
-                fieldName = parser.currentName();
-            } else if (token.isValue()) {
-                if (fieldName.equals(_ID_STRING)) {
-                    builder._id(parser.text());
-                } else if (fieldName.equals(NAME_STRING)) {
-                    builder.name(parser.text());
-                } else if (fieldName.equals(RESILIENCY_MODE_STRING)) {
-                    builder.mode(parser.text());
-                } else if (fieldName.equals(UPDATED_AT_STRING)) {
-                    builder.updatedAt(parser.longValue());
-                } else {
-                    throw new IllegalArgumentException(fieldName + " is not a valid field in QueryGroup");
-                }
-            } else if (token == XContentParser.Token.START_OBJECT) {
-
-                if (!fieldName.equals(RESOURCE_LIMITS_STRING)) {
-                    throw new IllegalArgumentException(
-                        "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token
-                    );
-                }
-
-                while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
-                    if (token == XContentParser.Token.FIELD_NAME) {
-                        fieldName = parser.currentName();
-                    } else {
-                        resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue());
-                    }
-                }
-
-            }
-        }
-        builder.resourceLimits(resourceLimits);
-        return builder.build();
+        return Builder.fromXContent(parser).build();
     }
 
     public static Diff<QueryGroup> readDiff(final StreamInput in) throws IOException {
@@ -293,6 +244,58 @@ public static class Builder {
 
         private Builder() {}
 
+        public static Builder fromXContent(XContentParser parser) throws IOException {
+            if (parser.currentToken() == null) { // fresh parser? move to the first token
+                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() + "]");
+            }
+
+            String fieldName = "";
+            // Map to hold resources
+            final Map<ResourceType, Double> resourceLimits = new HashMap<>();
+            while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+                if (token == XContentParser.Token.FIELD_NAME) {
+                    fieldName = parser.currentName();
+                } else if (token.isValue()) {
+                    if (fieldName.equals(_ID_STRING)) {
+                        builder._id(parser.text());
+                    } else if (fieldName.equals(NAME_STRING)) {
+                        builder.name(parser.text());
+                    } else if (fieldName.equals(RESILIENCY_MODE_STRING)) {
+                        builder.mode(parser.text());
+                    } else if (fieldName.equals(UPDATED_AT_STRING)) {
+                        builder.updatedAt(parser.longValue());
+                    } else {
+                        throw new IllegalArgumentException(fieldName + " is not a valid field in QueryGroup");
+                    }
+                } else if (token == XContentParser.Token.START_OBJECT) {
+
+                    if (!fieldName.equals(RESOURCE_LIMITS_STRING)) {
+                        throw new IllegalArgumentException(
+                            "QueryGroup.resourceLimits is an object and expected token was { " + " but found " + token
+                        );
+                    }
+
+                    while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+                        if (token == XContentParser.Token.FIELD_NAME) {
+                            fieldName = parser.currentName();
+                        } else {
+                            resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue());
+                        }
+                    }
+
+                }
+            }
+            return builder.resourceLimits(resourceLimits);
+        }
+
         public Builder name(String name) {
             this.name = name;
             return this;
@@ -319,11 +322,7 @@ public Builder resourceLimits(Map<ResourceType, Double> resourceLimits) {
         }
 
         public QueryGroup build() {
-            if (_id == null && updatedAt == 0L) {
-                return new QueryGroup(name, resiliencyMode, resourceLimits);
-            } else {
-                return new QueryGroup(name, _id, resiliencyMode, resourceLimits, updatedAt);
-            }
+            return new QueryGroup(name, _id, resiliencyMode, resourceLimits, updatedAt);
         }
     }
 }