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..0a879ec1eb997 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/Rule.java @@ -0,0 +1,370 @@ +/* + * 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.cluster.AbstractDiffable; +import org.opensearch.cluster.Diff; +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; +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.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; + +/** + * Class to define the Rule schema + * { + * "_id": "fwehf8302582mglfio349==", + * "index_pattern": ["logs123", "user*"], + * "query_group": "dev_query_group_id", + * "updated_at": "01-10-2025T21:23:21.456Z" + * } + */ +@ExperimentalApi +public class Rule extends AbstractDiffable implements ToXContentObject { + private final String _id; + private final Map> attributeMap; + private final Feature feature; + private final String label; + private final String updatedAt; + public static final Map> featureAlloedAttributesMap = Map.of( + Feature.QUERY_GROUP, + Set.of(RuleAttribute.INDEX_PATTERN) + ); + + public Rule(String _id, Map> attributeMap, String label, String updatedAt, Feature feature) { + requireNonNullOrEmpty(_id, "Rule _id can't be null or empty"); + Objects.requireNonNull(feature, "Couldn't identify which feature the rule belongs to. Rule feature name can't be null."); + requireNonNullOrEmpty(label, feature.getName() + " value can't be null or empty"); + requireNonNullOrEmpty(updatedAt, "Rule update time can't be null or empty"); + if (attributeMap == null || attributeMap.isEmpty()) { + throw new IllegalArgumentException("Rule should have at least 1 attribute requirement"); + } + if (!isValid(Instant.parse(updatedAt).getMillis())) { + throw new IllegalArgumentException("Rule update time is not a valid epoch"); + } + validatedAttributeMap(attributeMap, feature); + + this._id = _id; + this.attributeMap = attributeMap; + this.feature = feature; + this.label = label; + this.updatedAt = updatedAt; + } + + public Rule(StreamInput in) throws IOException { + this( + in.readString(), + 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) { + if (value == null || value.isEmpty()) { + throw new IllegalArgumentException(message); + } + } + + public static void validatedAttributeMap(Map> attributeMap, Feature feature) { + if (!featureAlloedAttributesMap.containsKey(feature)) { + throw new IllegalArgumentException("Couldn't find any valid attribute name under the feature: " + feature.getName()); + } + Set ValidAttributesForFeature = featureAlloedAttributesMap.get(feature); + for (Map.Entry> entry : attributeMap.entrySet()) { + RuleAttribute ruleAttribute = entry.getKey(); + Set attributeValues = entry.getValue(); + if (!ValidAttributesForFeature.contains(ruleAttribute)) { + throw new IllegalArgumentException( + ruleAttribute.getName() + " is not a valid attribute name under the feature: " + feature.getName() + ); + } + if (attributeValues.size() > 10) { + throw new IllegalArgumentException( + "Each attribute can only have a maximum of 10 values. The input attribute " + ruleAttribute + " exceeds this limit." + ); + } + for (String attributeValue : attributeValues) { + if (attributeValue.isEmpty()) { + throw new IllegalArgumentException("Attribute value should not be an empty string"); + } + if (attributeValue.length() > 100) { + throw new IllegalArgumentException( + "Attribute value can only have a maximum of 100 characters. The input " + attributeValue + " exceeds this limit." + ); + } + } + } + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeString(_id); + 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 get_id() { + return _id; + } + + public String getLabel() { + return label; + } + + public String getUpdatedAt() { + return updatedAt; + } + + public Feature getFeature() { + return feature; + } + + public Map> getAttributeMap() { + return attributeMap; + } + + /** + * This Feature enum contains the different feature names for each rule. + * For example, if we're creating a rule for WLM/QueryGroup, the rule will contain the line + * "query_group": "query_group_id", + * so the feature name would be "query_group" in this case. + */ + @ExperimentalApi + public enum Feature { + QUERY_GROUP("query_group"); + + private final String name; + + Feature(String name) { + this.name = name; + } + + public String getName() { + return name; + } + + public static boolean isValidFeature(String s) { + for (Feature feature : values()) { + if (feature.getName().equalsIgnoreCase(s)) { + return true; + } + } + return false; + } + + public static Feature fromName(String s) { + for (Feature feature : values()) { + if (feature.getName().equalsIgnoreCase(s)) return feature; + + } + throw new IllegalArgumentException("Invalid value for Feature: " + s); + } + } + + /** + * This RuleAttribute enum contains the attribute names for a rule. + */ + @ExperimentalApi + 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(); + builder.field("_id", _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", updatedAt); + builder.endObject(); + return builder; + } + + public XContentBuilder toXContentWithoutId(final XContentBuilder builder, final Params params) throws IOException { + builder.startObject(); + 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", updatedAt); + builder.endObject(); + return builder; + } + + public static Diff readDiff(final StreamInput in) throws IOException { + return readDiffFrom(Rule::new, in); + } + + @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(_id, that._id) + && 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(_id, label, feature, attributeMap, updatedAt); + } + + /** + * empty builder method for the {@link Rule} + * @return Builder object + */ + public static Builder builder() { + return new Builder(); + } + + /** + * builder method for the {@link Rule} + * @return Builder object + */ + public Builder builderFromRule() { + return new Builder()._id(_id).label(label).feature(feature.getName()).updatedAt(updatedAt).attributeMap(attributeMap); + } + + /** + * Builder class for {@link Rule} + */ + @ExperimentalApi + public static class Builder { + private String _id; + 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 (fieldName.equals("_id")) { + builder._id(parser.text()); + } else if (Feature.isValidFeature(fieldName)) { + builder.feature(fieldName); + builder.label(parser.text()); + } else if (fieldName.equals("updated_at")) { + builder.updatedAt(parser.text()); + } else { + throw new IllegalArgumentException(fieldName + " is not a valid field in Rule"); + } + } else if (token == XContentParser.Token.START_ARRAY) { + RuleAttribute ruleAttribute = RuleAttribute.fromName(fieldName); + Set indexPatternList = new HashSet<>(); + while (parser.nextToken() != XContentParser.Token.END_ARRAY) { + if (parser.currentToken() == XContentParser.Token.VALUE_STRING) { + indexPatternList.add(parser.text()); + } else { + throw new XContentParseException("Unexpected token in array: " + parser.currentToken()); + } + } + attributeMap1.put(ruleAttribute, indexPatternList); + } + } + return builder.attributeMap(attributeMap1); + } + + public Builder _id(String _id) { + this._id = _id; + return this; + } + + 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(_id, attributeMap, label, updatedAt, feature); + } + + public String getLabel() { + return label; + } + } +} 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..06ddefc0161fd --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/RuleTests.java @@ -0,0 +1,175 @@ +/* + * 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.featureAlloedAttributesMap; + +public class RuleTests extends AbstractSerializingTestCase { + public static final String _ID = "AgfUfjw039vhdONlYi3TQ=="; + public static final String LABEL = "label"; + + static Rule createRandomRule(String _id, String label) { + Feature feature = randomFeature(); + return Rule.builder() + ._id(_id) + .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<>(featureAlloedAttributesMap.get(feature)); + 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, 10); + for (int i = 0; i < numberOfValues; i++) { + res.add(randomAlphaOfLength(randomIntBetween(1, 100))); + } + 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(_ID, LABEL); + } + + static Rule buildRule(String _id, String label, String feature, Map> attributeListMap, String updatedAt) { + return Rule.builder()._id(_id).label(label).feature(feature).attributeMap(attributeListMap).updatedAt(updatedAt).build(); + } + + public void testInvalidId() { + assertThrows(IllegalArgumentException.class, () -> createRandomRule(null, LABEL)); + assertThrows(IllegalArgumentException.class, () -> createRandomRule("", LABEL)); + } + + public void testInvalidFeature() { + assertThrows( + IllegalArgumentException.class, + () -> buildRule(_ID, LABEL, null, randomAttributeMaps(null), Instant.now().toString()) + ); + assertThrows( + IllegalArgumentException.class, + () -> buildRule(_ID, LABEL, "invalid", randomAttributeMaps(null), Instant.now().toString()) + ); + } + + public void testInvalidLabel() { + assertThrows(IllegalArgumentException.class, () -> createRandomRule(_ID, null)); + assertThrows(IllegalArgumentException.class, () -> createRandomRule(_ID, "")); + } + + public void testInvalidUpdateTime() { + Feature feature = randomFeature(); + assertThrows(IllegalArgumentException.class, () -> buildRule(_ID, LABEL, feature.toString(), randomAttributeMaps(feature), null)); + } + + public void testNullOrEmptyAttributeMap() { + Feature feature = randomFeature(); + assertThrows( + IllegalArgumentException.class, + () -> buildRule(_ID, LABEL, feature.toString(), new HashMap<>(), Instant.now().toString()) + ); + assertThrows(IllegalArgumentException.class, () -> buildRule(_ID, 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(_ID, LABEL, randomFeature().getName(), map, Instant.now().toString())); + + map.put(RuleAttribute.INDEX_PATTERN, Set.of(randomAlphaOfLength(101))); + assertThrows(IllegalArgumentException.class, () -> buildRule(_ID, LABEL, randomFeature().getName(), map, Instant.now().toString())); + + map.put(RuleAttribute.INDEX_PATTERN, new HashSet<>()); + for (int i = 0; i < 11; i++) { + map.get(RuleAttribute.INDEX_PATTERN).add(String.valueOf(i)); + } + assertThrows(IllegalArgumentException.class, () -> buildRule(_ID, 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(_ID, LABEL, Feature.QUERY_GROUP.getName(), map, updatedAt); + assertNotNull(rule.get_id()); + assertEquals(_ID, rule.get_id()); + 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(_ID, LABEL, Feature.QUERY_GROUP.getName(), map, updatedAt); + + XContentBuilder builder = JsonXContent.contentBuilder(); + rule.toXContent(builder, ToXContent.EMPTY_PARAMS); + + assertEquals( + "{\"_id\":\"" + _ID + "\",\"index_pattern\":[\"log*\"],\"query_group\":\"label\",\"updated_at\":\"" + updatedAt + "\"}", + builder.toString() + ); + } +}