Skip to content
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

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
## [Unreleased 2.x]
### Added
- Latency and Memory allocation improvements to Multi Term Aggregation queries ([#14993](https://github.com/opensearch-project/OpenSearch/pull/14993))
- Add rule schema changes for rule based auto tagging ([#17238](https://github.com/opensearch-project/OpenSearch/pull/17238))
- Add logic in master service to optimize performance and retain detailed logging for critical cluster operations. ([#14795](https://github.com/opensearch-project/OpenSearch/pull/14795))
- Add Setting to adjust the primary constraint weights ([#16471](https://github.com/opensearch-project/OpenSearch/pull/16471))
- Introduce a setting to disable download of full cluster state from remote on term mismatch([#16798](https://github.com/opensearch-project/OpenSearch/pull/16798/))
Expand Down
377 changes: 377 additions & 0 deletions server/src/main/java/org/opensearch/autotagging/Rule.java
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==",
* "name": "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 name;
private final String label;
Copy link
Collaborator

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?

private final String updatedAt;
public static final String _ID_STRING = "_id";
public static final String NAME_STRING = "name";
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 name, Map<Attribute, Set<String>> attributeMap, String label, String updatedAt, Feature feature) {
ValidationException validationException = new ValidationException();
validateRuleInputs(name, attributeMap, label, updatedAt, feature, validationException);
if (!validationException.validationErrors().isEmpty()) {
throw new IllegalArgumentException(validationException);
}

this.name = name;
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 name,
Map<Attribute, Set<String>> attributeMap,
String label,
String updatedAt,
Feature feature,
ValidationException validationException
Copy link
Contributor

Choose a reason for hiding this comment

The 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?

public static Optional<ValidationException> validateRuleInputs() {}

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.

Copy link
Contributor

Choose a reason for hiding this comment

The 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 Rule class level above?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think creating a Validator class with all of the attributes which needs to be validated as member variables will be more suitable.

we can create something like following

static class RuleValidator {
      String description,
        Map<Attribute, Set<String>> attributeMap;
        String label;
        String updatedAt;
        Feature feature;
....

     Optional<ValidationException> validate() {
               validateAttributes();
               validateFeature();
         ....
    }
...
}

Copy link
Contributor Author

@ruai0511 ruai0511 Feb 18, 2025

Choose a reason for hiding this comment

The 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 public static Optional<ValidationException> validateRuleInputs() {} then the internal error-checking functions like validateAttributeMap can't reuse the ValidationException object

) {
requireNonNullOrEmpty(name, "Rule name 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.");

Check warning on line 91 in server/src/main/java/org/opensearch/autotagging/Rule.java

View check run for this annotation

Codecov / codecov/patch

server/src/main/java/org/opensearch/autotagging/Rule.java#L91

Added line #L91 was not covered by tests
}
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");

Check warning on line 99 in server/src/main/java/org/opensearch/autotagging/Rule.java

View check run for this annotation

Codecov / codecov/patch

server/src/main/java/org/opensearch/autotagging/Rule.java#L99

Added line #L99 was not covered by tests
}
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()

Check warning on line 122 in server/src/main/java/org/opensearch/autotagging/Rule.java

View check run for this annotation

Codecov / codecov/patch

server/src/main/java/org/opensearch/autotagging/Rule.java#L121-L122

Added lines #L121 - L122 were not covered by tests
);
}
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(name);
out.writeMap(attributeMap, Attribute::writeTo, StreamOutput::writeStringCollection);
out.writeString(label);
out.writeString(updatedAt);
out.writeString(feature.getName());
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be delegated to Feature class

}

public static Rule fromXContent(final XContentParser parser) throws IOException {
return Builder.fromXContent(parser).build();
}

public String getName() {
return name;

Check warning on line 154 in server/src/main/java/org/opensearch/autotagging/Rule.java

View check run for this annotation

Codecov / codecov/patch

server/src/main/java/org/opensearch/autotagging/Rule.java#L154

Added line #L154 was not covered by tests
}

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 {
Copy link
Contributor

Choose a reason for hiding this comment

The 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 Attribute

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");
Copy link
Contributor

@sgup432 sgup432 Feb 17, 2025

Choose a reason for hiding this comment

The 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?
Will they add another enum for their usecase? That doesn't seem right. As core should be agnostic of plugin level logic.

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);

Check warning on line 238 in server/src/main/java/org/opensearch/autotagging/Rule.java

View check run for this annotation

Codecov / codecov/patch

server/src/main/java/org/opensearch/autotagging/Rule.java#L238

Added line #L238 was not covered by tests
}
}

@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(NAME_STRING, name);
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);
Copy link
Collaborator

Choose a reason for hiding this comment

The 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(name, that.name)
&& 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(name, 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 name;
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() + "]");

Check warning on line 305 in server/src/main/java/org/opensearch/autotagging/Rule.java

View check run for this annotation

Codecov / codecov/patch

server/src/main/java/org/opensearch/autotagging/Rule.java#L305

Added line #L305 was not covered by tests
}
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(NAME_STRING)) {
builder.name(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");

Check warning on line 321 in server/src/main/java/org/opensearch/autotagging/Rule.java

View check run for this annotation

Codecov / codecov/patch

server/src/main/java/org/opensearch/autotagging/Rule.java#L321

Added line #L321 was not covered by tests
}
} 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());

Check warning on line 338 in server/src/main/java/org/opensearch/autotagging/Rule.java

View check run for this annotation

Codecov / codecov/patch

server/src/main/java/org/opensearch/autotagging/Rule.java#L338

Added line #L338 was not covered by tests
}
}
attributeMap.put(attribute, attributeValueSet);
}

public Builder name(String name) {
this.name = name;
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(name, attributeMap, label, updatedAt, feature);
}

public String getLabel() {
return label;

Check warning on line 374 in server/src/main/java/org/opensearch/autotagging/Rule.java

View check run for this annotation

Codecov / codecov/patch

server/src/main/java/org/opensearch/autotagging/Rule.java#L374

Added line #L374 was not covered by tests
}
}
}
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;
Loading
Loading