-
Notifications
You must be signed in to change notification settings - Fork 2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Workload Management] Add rule schema for workload management #17238
base: main
Are you sure you want to change the base?
Changes from all commits
f6a4a28
bb8d06c
b32ce69
f46caf8
e152d95
aa1f8b7
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,377 @@ | ||
/* | ||
* 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.autotagging; | ||
|
||
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==", | ||
* "description": "Assign Query Group for Index Logs123" | ||
* "index_pattern": ["logs123"], | ||
* "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<Attribute, Set<String>> attributeMap; | ||
private final Feature feature; | ||
private final String description; | ||
private final String label; | ||
private final String updatedAt; | ||
public static final String _ID_STRING = "_id"; | ||
public static final String DESCRIPTION_STRING = "description"; | ||
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(String description, Map<Attribute, Set<String>> attributeMap, String label, String updatedAt, Feature feature) { | ||
ValidationException validationException = new ValidationException(); | ||
validateRuleInputs(description, attributeMap, label, updatedAt, feature, validationException); | ||
if (!validationException.validationErrors().isEmpty()) { | ||
throw new IllegalArgumentException(validationException); | ||
} | ||
|
||
this.description = description; | ||
this.attributeMap = attributeMap; | ||
this.feature = feature; | ||
this.label = label; | ||
this.updatedAt = updatedAt; | ||
} | ||
|
||
public Rule(StreamInput in) throws IOException { | ||
this( | ||
in.readString(), | ||
in.readMap((i) -> Attribute.fromName(i.readString()), i -> new HashSet<>(i.readStringList())), | ||
in.readString(), | ||
in.readString(), | ||
Feature.fromName(in.readString()) | ||
); | ||
} | ||
|
||
public static void validateRuleInputs( | ||
String description, | ||
Map<Attribute, Set<String>> attributeMap, | ||
String label, | ||
String updatedAt, | ||
Feature feature, | ||
ValidationException validationException | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Passing validationException as a method param doesn't seem a very nice way to me. Why not do below?
Here you just pass an optional exception in case u encounter any issues. The caller can additionally check if any exceptions are present or not, and accordingly handle it. Same for below method. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In addition to this, instead of passing all these params, should we create an internal object which can be passed instead? And same can be used at There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think creating a we can create something like following
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I put it as a method param to reuse the same ValidationException object and throw all the errors together after all the checks. If we use |
||
) { | ||
requireNonNullOrEmpty(description, "Rule description can't be null or empty", 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 requireNonNullOrEmpty(String value, String message, ValidationException validationException) { | ||
if (value == null || value.isEmpty()) { | ||
validationException.addValidationError(message); | ||
} | ||
} | ||
|
||
public static void validateAttributeMap( | ||
Map<Attribute, Set<String>> attributeMap, | ||
Feature feature, | ||
ValidationException validationException | ||
) { | ||
for (Map.Entry<Attribute, Set<String>> entry : attributeMap.entrySet()) { | ||
Attribute attribute = entry.getKey(); | ||
Set<String> attributeValues = entry.getValue(); | ||
if (!feature.isValidAttribute(attribute)) { | ||
validationException.addValidationError( | ||
attribute.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 " + attribute + " 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.writeString(description); | ||
out.writeMap(attributeMap, Attribute::writeTo, StreamOutput::writeStringCollection); | ||
out.writeString(label); | ||
out.writeString(updatedAt); | ||
out.writeString(feature.getName()); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This should be delegated to |
||
} | ||
|
||
public static Rule fromXContent(final XContentParser parser) throws IOException { | ||
return Builder.fromXContent(parser).build(); | ||
} | ||
|
||
public String getDescription() { | ||
return description; | ||
} | ||
|
||
public String getLabel() { | ||
return label; | ||
} | ||
|
||
public String getUpdatedAt() { | ||
return updatedAt; | ||
} | ||
|
||
public Feature getFeature() { | ||
return feature; | ||
} | ||
|
||
public Map<Attribute, Set<String>> getAttributeMap() { | ||
return attributeMap; | ||
} | ||
|
||
/** | ||
* This enum enumerates the features that can use the Rule Based Auto-tagging | ||
* @opensearch.experimental | ||
*/ | ||
public enum Feature { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [NIT] Better to create a different file for this IMO. Same for below |
||
QUERY_GROUP("query_group", Set.of(Attribute.INDEX_PATTERN)); | ||
|
||
private final String name; | ||
private final Set<Attribute> allowedAttributes; | ||
|
||
Feature(String name, Set<Attribute> allowedAttributes) { | ||
this.name = name; | ||
this.allowedAttributes = allowedAttributes; | ||
} | ||
|
||
public String getName() { | ||
return name; | ||
} | ||
|
||
public Set<Attribute> getAllowedAttributes() { | ||
return allowedAttributes; | ||
} | ||
|
||
public boolean isValidAttribute(Attribute 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 Attribute enum contains the attribute names for a rule. | ||
* @opensearch.experimental | ||
*/ | ||
public enum Attribute { | ||
INDEX_PATTERN("index_pattern"); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. What happens if an external plugin like queryInsights or other want to use this schema and eventually this Attribute class? Same for Feature class above. |
||
|
||
private final String name; | ||
|
||
Attribute(String name) { | ||
this.name = name; | ||
} | ||
|
||
public String getName() { | ||
return name; | ||
} | ||
|
||
public static void writeTo(StreamOutput out, Attribute attribute) throws IOException { | ||
out.writeString(attribute.getName()); | ||
} | ||
|
||
public static Attribute fromName(String s) { | ||
for (Attribute attribute : values()) { | ||
if (attribute.getName().equalsIgnoreCase(s)) return attribute; | ||
|
||
} | ||
throw new IllegalArgumentException("Invalid value for Attribute: " + 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); | ||
} | ||
builder.field(DESCRIPTION_STRING, description); | ||
for (Map.Entry<Attribute, Set<String>> entry : attributeMap.entrySet()) { | ||
builder.array(entry.getKey().getName(), entry.getValue().toArray(new String[0])); | ||
} | ||
builder.field(feature.getName(), label); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. should be delegated to the feature class |
||
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(description, that.description) | ||
&& 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(description, 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 String description; | ||
private Map<Attribute, Set<String>> 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<Attribute, Set<String>> 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(DESCRIPTION_STRING)) { | ||
builder.description(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<Attribute, Set<String>> attributeMap) | ||
throws IOException { | ||
Attribute attribute = Attribute.fromName(fieldName); | ||
Set<String> 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(attribute, attributeValueSet); | ||
} | ||
|
||
public Builder description(String description) { | ||
this.description = description; | ||
return this; | ||
} | ||
|
||
public Builder label(String label) { | ||
this.label = label; | ||
return this; | ||
} | ||
|
||
public Builder attributeMap(Map<Attribute, Set<String>> 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(description, attributeMap, label, updatedAt, feature); | ||
} | ||
|
||
public String getLabel() { | ||
return label; | ||
} | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,13 @@ | ||
/* | ||
* 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 auto tagging constructs | ||
*/ | ||
|
||
package org.opensearch.autotagging; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
label is corresponding to the feature right? Does it make sense to include as part of Feature?