address comments

Signed-off-by: Kaushal Kumar <[email protected]>
opensearch-project · Jul 3, 2024 · 8a45648 · 8a45648
1 parent 42d3494
commit 8a45648
Show file tree

Hide file tree

Showing 4 changed files with 82 additions and 29 deletions.
diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java
@@ -11,7 +11,7 @@
 import org.opensearch.cluster.AbstractDiffable;
 import org.opensearch.cluster.Diff;
 import org.opensearch.common.UUIDs;
-import org.opensearch.common.annotation.PublicApi;
+import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.xcontent.ToXContentObject;
@@ -37,10 +37,10 @@
  *              "updatedAt": 4513232415
  * }
  */
-@PublicApi(since = "2.15")
+@ExperimentalApi
 public class QueryGroup extends AbstractDiffable<QueryGroup> implements ToXContentObject {
 
-    public static final int MAX_CHARS_ALLOWED_IN_NAME = 50;
+    private static final int MAX_CHARS_ALLOWED_IN_NAME = 50;
     private final String name;
     private final String _id;
     private final ResiliencyMode resiliencyMode;
@@ -154,7 +154,7 @@ public static QueryGroup fromXContent(final XContentParser parser) throws IOExce
         String fieldName = "";
         // Map to hold resources
         final Map<ResourceType, Object> resourceLimits = new HashMap<>();
-        while ((token = parser.nextToken()) != null) {
+        while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
             if (token == XContentParser.Token.FIELD_NAME) {
                 fieldName = parser.currentName();
             } else if (token.isValue()) {
@@ -245,7 +245,7 @@ public static Builder builder() {
      * ENFORCED - means that it will never breach the assigned limits and will cancel as soon as the limits are breached
      * MONITOR - it will not cause any cancellation but just log the eligible task cancellations
      */
-    @PublicApi(since = "2.15")
+    @ExperimentalApi
     public enum ResiliencyMode {
         SOFT("soft"),
         ENFORCED("enforced"),
@@ -274,7 +274,7 @@ public static ResiliencyMode fromName(String s) {
     /**
      * Builder class for {@link QueryGroup}
      */
-    @PublicApi(since = "2.15")
+    @ExperimentalApi
     public static class Builder {
         private String name;
         private String _id;

diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroupMetadata.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroupMetadata.java
@@ -16,7 +16,6 @@
 import org.opensearch.core.common.Strings;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
-import org.opensearch.core.xcontent.ConstructingObjectParser;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
@@ -47,22 +46,6 @@ public class QueryGroupMetadata implements Metadata.Custom {
 
     private final Map<String, QueryGroup> queryGroups;
 
-    @SuppressWarnings("unchecked")
-    static final ConstructingObjectParser<QueryGroupMetadata, Void> PARSER = new ConstructingObjectParser<>(
-        "queryGroupsParser",
-        args -> new QueryGroupMetadata((Map<String, QueryGroup>) args[0])
-    );
-
-    static {
-        PARSER.declareObjectArray(ConstructingObjectParser.constructorArg(), (p, c) -> {
-            Map<String, QueryGroup> queryGroupMap = new HashMap<>();
-            while (p.nextToken() != XContentParser.Token.END_OBJECT) {
-                queryGroupMap.put(p.currentName(), QueryGroup.fromXContent(p));
-            }
-            return queryGroupMap;
-        }, QUERY_GROUP_FIELD);
-    }
-
     public QueryGroupMetadata(Map<String, QueryGroup> queryGroups) {
         this.queryGroups = queryGroups;
     }
@@ -105,7 +88,29 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
     }
 
     public static QueryGroupMetadata fromXContent(XContentParser parser) throws IOException {
-        return PARSER.parse(parser, null);
+        Map<String, QueryGroup> queryGroupMap = new HashMap<>();
+
+        if (parser.currentToken() == null) {
+            parser.nextToken();
+        }
+
+        if (parser.currentToken() != XContentParser.Token.START_OBJECT) {
+            throw new IllegalArgumentException(
+                "QueryGroupMetadata.fromXContent was expecting a { token but found : " + parser.currentToken()
+            );
+        }
+        XContentParser.Token token = parser.currentToken();
+        String fieldName = parser.currentName();
+        while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+            if (token == XContentParser.Token.FIELD_NAME) {
+                fieldName = parser.currentName();
+            } else {
+                QueryGroup queryGroup = QueryGroup.fromXContent(parser);
+                queryGroupMap.put(fieldName, queryGroup);
+            }
+        }
+
+        return new QueryGroupMetadata(queryGroupMap);
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java
@@ -16,7 +16,7 @@
 /**
  * Enum to hold the resource type
  */
-@PublicApi(since = "1.3")
+@PublicApi(since = "2.x")
 public enum ResourceType {
     CPU("cpu"),
     MEMORY("memory");

diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java
@@ -8,15 +8,46 @@
 
 package org.opensearch.cluster.metadata;
 
+import org.opensearch.cluster.Diff;
+import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
-import org.opensearch.test.AbstractNamedWriteableTestCase;
+import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.search.ResourceType;
+import org.opensearch.test.AbstractDiffableSerializationTestCase;
 
+import java.io.IOException;
 import java.util.Collections;
 import java.util.Map;
 
 import static org.opensearch.cluster.metadata.QueryGroupTests.createRandomQueryGroup;
 
-public class QueryGroupMetadataTests extends AbstractNamedWriteableTestCase<QueryGroupMetadata> {
+public class QueryGroupMetadataTests extends AbstractDiffableSerializationTestCase<Metadata.Custom> {
+
+    public void testToXContent() throws IOException {
+        long updatedAt = 1720047207;
+        QueryGroupMetadata queryGroupMetadata = new QueryGroupMetadata(
+            Map.of(
+                "ajakgakg983r92_4242",
+                new QueryGroup(
+                    "test",
+                    "ajakgakg983r92_4242",
+                    QueryGroup.ResiliencyMode.ENFORCED,
+                    Map.of(ResourceType.MEMORY, 0.5),
+                    updatedAt
+                )
+            )
+        );
+        XContentBuilder builder = JsonXContent.contentBuilder();
+        builder.startObject();
+        queryGroupMetadata.toXContent(builder, null);
+        builder.endObject();
+        assertEquals(
+            "{\"ajakgakg983r92_4242\":{\"_id\":\"ajakgakg983r92_4242\",\"name\":\"test\",\"resiliency_mode\":\"enforced\",\"updatedAt\":1720047207,\"resourceLimits\":{\"memory\":0.5}}}",
+            builder.toString()
+        );
+    }
 
     @Override
     protected NamedWriteableRegistry getNamedWriteableRegistry() {
@@ -28,8 +59,25 @@ protected NamedWriteableRegistry getNamedWriteableRegistry() {
     }
 
     @Override
-    protected Class<QueryGroupMetadata> categoryClass() {
-        return QueryGroupMetadata.class;
+    protected Metadata.Custom makeTestChanges(Metadata.Custom testInstance) {
+        final QueryGroup queryGroup = createRandomQueryGroup("asdfakgjwrir23r25");
+        final QueryGroupMetadata queryGroupMetadata = new QueryGroupMetadata(Map.of(queryGroup.get_id(), queryGroup));
+        return queryGroupMetadata;
+    }
+
+    @Override
+    protected Writeable.Reader<Diff<Metadata.Custom>> diffReader() {
+        return QueryGroupMetadata::readDiffFrom;
+    }
+
+    @Override
+    protected Metadata.Custom doParseInstance(XContentParser parser) throws IOException {
+        return QueryGroupMetadata.fromXContent(parser);
+    }
+
+    @Override
+    protected Writeable.Reader<Metadata.Custom> instanceReader() {
+        return QueryGroupMetadata::new;
     }
 
     @Override