[Rule-based Auto-tagging] add the schema for security attributes (opensearch-project#19345)

Signed-off-by: Ruirui Zhang <[email protected]>

Author: ruai0511

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 3.x]
 ### Added
 - Expand fetch phase profiling to support inner hits and top hits aggregation phases  ([##18936](https://github.com/opensearch-project/OpenSearch/pull/18936))
+- [Rule-based Auto-tagging] add the schema for security attributes ([##19345](https://github.com/opensearch-project/OpenSearch/pull/19345))
 - Add temporal routing processors for time-based document routing ([#18920](https://github.com/opensearch-project/OpenSearch/issues/18920))
 - Implement Query Rewriting Infrastructure ([#19060](https://github.com/opensearch-project/OpenSearch/pull/19060))
 - The dynamic mapping parameter supports false_allow_templates ([#19065](https://github.com/opensearch-project/OpenSearch/pull/19065) ([#19097](https://github.com/opensearch-project/OpenSearch/pull/19097)))

modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/Attribute.java

@@ -11,8 +11,14 @@
 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.XContentBuilder;
+import org.opensearch.core.xcontent.XContentParseException;
+import org.opensearch.core.xcontent.XContentParser;
 
 import java.io.IOException;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.TreeMap;
 
 /**
  * Represents an attribute within the auto-tagging feature. Attributes define characteristics that can
@@ -29,6 +35,13 @@ public interface Attribute extends Writeable {
      */
     String getName();
 
+    /**
+     * Returns a map of subfields ordered by priority, where 1 represents the highest priority.
+     */
+    default TreeMap<Integer, String> getPrioritizedSubfields() {
+        return new TreeMap<>();
+    }
+
     /**
      * Ensure that `validateAttribute` is called in the constructor of attribute implementations
      * to prevent potential serialization issues.
@@ -45,6 +58,39 @@ default void writeTo(StreamOutput out) throws IOException {
         out.writeString(getName());
     }
 
+    /**
+     * Parses attribute values for specific attributes. This default function takes in parser
+     * and returns a set of string.
+     * For example, ["index1", "index2"] will be parsed to a set with values "index1" and "index2"
+     * @param parser the XContent parser
+     */
+    default Set<String> fromXContentParseAttributeValues(XContentParser parser) throws IOException {
+        if (parser.currentToken() != XContentParser.Token.START_ARRAY) {
+            throw new XContentParseException(
+                parser.getTokenLocation(),
+                "Expected START_ARRAY token for " + getName() + " attribute but got " + parser.currentToken()
+            );
+        }
+        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());
+            }
+        }
+        return attributeValueSet;
+    }
+
+    /**
+     * Writes a set of attribute values for a specific attribute
+     * @param builder the XContent builder
+     * @param values the set of string values to write
+     */
+    default void toXContentWriteAttributeValues(XContentBuilder builder, Set<String> values) throws IOException {
+        builder.array(getName(), values.toArray(new String[0]));
+    }
+
     /**
      * Retrieves an attribute from the given feature type based on its name.
      * Implementations of `FeatureType.getAttributeFromName` must be thread-safe as this method

modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/AutoTaggingRegistry.java

@@ -59,6 +59,11 @@ private static void validateFeatureType(FeatureType featureType) {
                 "Feature type name " + name + " should not be null, empty or have more than  " + MAX_FEATURE_TYPE_NAME_LENGTH + "characters"
             );
         }
+        if (featureType.getOrderedAttributes() == null) {
+            throw new IllegalStateException(
+                "Function getOrderedAttributes() should not return null for feature type: " + featureType.getName()
+            );
+        }
         if (featureType.getFeatureValueValidator() == null) {
             throw new IllegalStateException("FeatureValueValidator is not defined for feature type " + name);
         }

modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/FeatureType.java

@@ -14,6 +14,7 @@
 
 import java.io.IOException;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 /**
  * Represents a feature type within the auto-tagging feature. Feature types define different categories of
@@ -41,11 +42,19 @@ public interface FeatureType extends Writeable {
      */
     String getName();
 
+    /**
+     * Returns a map of top-level attributes sorted by priority, with 1 representing the highest priority.
+     * Subfields within each attribute are managed separately here {@link org.opensearch.rule.autotagging.Attribute#getPrioritizedSubfields()}.
+     */
+    Map<Attribute, Integer> getOrderedAttributes();
+
     /**
      * Returns the registry of allowed attributes for this feature type.
      * Implementations must ensure that access to this registry is thread-safe.
      */
-    Map<String, Attribute> getAllowedAttributesRegistry();
+    default Map<String, Attribute> getAllowedAttributesRegistry() {
+        return getOrderedAttributes().keySet().stream().collect(Collectors.toUnmodifiableMap(Attribute::getName, attribute -> attribute));
+    }
 
     /**
      * returns the validator for feature value

modules/autotagging-commons/common/src/main/java/org/opensearch/rule/autotagging/Rule.java

@@ -177,7 +177,7 @@ public XContentBuilder toXContent(final XContentBuilder builder, final Params pa
         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]));
+            entry.getKey().toXContentWriteAttributeValues(builder, entry.getValue());
         }
         builder.field(featureType.getName(), featureValue);
         builder.field(UPDATED_AT_STRING, updatedAt);
@@ -258,14 +258,14 @@ public static Builder fromXContent(XContentParser parser, FeatureType featureTyp
                         builder.featureType(featureType);
                         builder.featureValue(parser.text());
                     }
-                } else if (token == XContentParser.Token.START_ARRAY) {
-                    fromXContentParseArray(parser, fieldName, featureType, attributeMap1);
+                } else if (token == XContentParser.Token.START_ARRAY || token == XContentParser.Token.START_OBJECT) {
+                    fromXContentParseAttribute(parser, fieldName, featureType, attributeMap1);
                 }
             }
             return builder.attributeMap(attributeMap1);
         }
 
-        private static void fromXContentParseArray(
+        private static void fromXContentParseAttribute(
             XContentParser parser,
             String fieldName,
             FeatureType featureType,
@@ -275,15 +275,7 @@ private static void fromXContentParseArray(
             if (attribute == null) {
                 throw new XContentParseException(fieldName + " is not a valid attribute within the " + featureType.getName() + " feature.");
             }
-            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);
+            attributeMap.put(attribute, attribute.fromXContentParseAttributeValues(parser));
         }
 
         /**

modules/autotagging-commons/common/src/test/java/org/opensearch/rule/RuleAttributeTests.java

@@ -8,8 +8,14 @@
 
 package org.opensearch.rule;
 
+import org.opensearch.common.xcontent.json.JsonXContent;
+import org.opensearch.core.xcontent.XContentParseException;
+import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.test.OpenSearchTestCase;
 
+import java.io.IOException;
+import java.util.Set;
+
 public class RuleAttributeTests extends OpenSearchTestCase {
 
     public void testGetName() {
@@ -25,4 +31,40 @@ public void testFromName() {
     public void testFromName_throwsException() {
         assertThrows(IllegalArgumentException.class, () -> RuleAttribute.fromName("invalid_attribute"));
     }
+
+    public void testGetPrioritizedSubfields() {
+        assertTrue(RuleAttribute.INDEX_PATTERN.getPrioritizedSubfields().isEmpty());
+    }
+
+    public void testFromXContentParseAttributeValues_success() throws IOException {
+        String json = "{ \"index_pattern\": [\"val1\", \"val2\"] }";
+        try (XContentParser parser = JsonXContent.jsonXContent.createParser(null, null, json)) {
+            skipTokens(parser, 3);
+            Set<String> result = RuleAttribute.INDEX_PATTERN.fromXContentParseAttributeValues(parser);
+            assertTrue(result.contains("val1"));
+            assertTrue(result.contains("val2"));
+        }
+    }
+
+    public void testFromXContentParseAttributeValues_notStartArray() throws IOException {
+        String json = "{ \"index_pattern\": \"val1\" }";
+        try (XContentParser parser = JsonXContent.jsonXContent.createParser(null, null, json)) {
+            skipTokens(parser, 3);
+            assertThrows(XContentParseException.class, () -> RuleAttribute.INDEX_PATTERN.fromXContentParseAttributeValues(parser));
+        }
+    }
+
+    public void testFromXContentParseAttributeValues_unexpectedToken() throws IOException {
+        String json = "{ \"index_pattern\": [123] }";
+        try (XContentParser parser = JsonXContent.jsonXContent.createParser(null, null, json)) {
+            skipTokens(parser, 3);
+            assertThrows(XContentParseException.class, () -> RuleAttribute.INDEX_PATTERN.fromXContentParseAttributeValues(parser));
+        }
+    }
+
+    private void skipTokens(XContentParser parser, int count) throws IOException {
+        for (int i = 0; i < count; i++) {
+            parser.nextToken();
+        }
+    }
 }

modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/AutoTaggingRegistryTests.java

@@ -13,6 +13,8 @@
 import org.opensearch.test.OpenSearchTestCase;
 import org.junit.BeforeClass;
 
+import java.util.HashMap;
+
 import static org.opensearch.rule.autotagging.AutoTaggingRegistry.MAX_FEATURE_TYPE_NAME_LENGTH;
 import static org.opensearch.rule.autotagging.RuleTests.INVALID_FEATURE;
 import static org.opensearch.rule.utils.RuleTestUtils.FEATURE_TYPE_NAME;
@@ -38,10 +40,22 @@ public void testRuntimeException() {
 
     public void testIllegalStateExceptionException() {
         assertThrows(IllegalStateException.class, () -> AutoTaggingRegistry.registerFeatureType(null));
+
         FeatureType featureType = mock(FeatureType.class);
         when(featureType.getName()).thenReturn(FEATURE_TYPE_NAME);
+        when(featureType.getOrderedAttributes()).thenReturn(null);
+        when(featureType.getFeatureValueValidator()).thenReturn(new FeatureValueValidator() {
+            @Override
+            public void validate(String featureValue) {}
+        });
         assertThrows(IllegalStateException.class, () -> AutoTaggingRegistry.registerFeatureType(featureType));
+
         when(featureType.getName()).thenReturn(randomAlphaOfLength(MAX_FEATURE_TYPE_NAME_LENGTH + 1));
         assertThrows(IllegalStateException.class, () -> AutoTaggingRegistry.registerFeatureType(featureType));
+
+        when(featureType.getName()).thenReturn(FEATURE_TYPE_NAME);
+        when(featureType.getOrderedAttributes()).thenReturn(new HashMap<>());
+        when(featureType.getFeatureValueValidator()).thenReturn(null);
+        assertThrows(IllegalStateException.class, () -> AutoTaggingRegistry.registerFeatureType(featureType));
     }
 }

modules/autotagging-commons/common/src/test/java/org/opensearch/rule/autotagging/RuleTests.java

@@ -77,12 +77,6 @@ public static class TestFeatureType implements FeatureType {
         private static final String NAME = TEST_FEATURE_TYPE;
         private static final int MAX_ATTRIBUTE_VALUES = 10;
         private static final int MAX_ATTRIBUTE_VALUE_LENGTH = 100;
-        private static final Map<String, Attribute> ALLOWED_ATTRIBUTES = Map.of(
-            TEST_ATTR1_NAME,
-            TEST_ATTRIBUTE_1,
-            TEST_ATTR2_NAME,
-            TEST_ATTRIBUTE_2
-        );
 
         public TestFeatureType() {}
 
@@ -95,6 +89,11 @@ public String getName() {
             return NAME;
         }
 
+        @Override
+        public Map<Attribute, Integer> getOrderedAttributes() {
+            return Map.of(TEST_ATTRIBUTE_1, 1, TEST_ATTRIBUTE_2, 2);
+        }
+
         @Override
         public int getMaxNumberOfValuesPerAttribute() {
             return MAX_ATTRIBUTE_VALUES;
@@ -104,11 +103,6 @@ public int getMaxNumberOfValuesPerAttribute() {
         public int getMaxCharLengthPerAttributeValue() {
             return MAX_ATTRIBUTE_VALUE_LENGTH;
         }
-
-        @Override
-        public Map<String, Attribute> getAllowedAttributesRegistry() {
-            return ALLOWED_ATTRIBUTES;
-        }
     }
 
     static Rule buildRule(

modules/autotagging-commons/common/src/test/java/org/opensearch/rule/utils/RuleTestUtils.java

@@ -88,6 +88,12 @@ public static void assertEqualRule(Rule one, Rule two, boolean ruleUpdated) {
     public static class MockRuleFeatureType implements FeatureType {
 
         public static final MockRuleFeatureType INSTANCE = new MockRuleFeatureType();
+        private static final Map<Attribute, Integer> PRIORITIZED_ATTRIBUTES = Map.of(
+            MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE,
+            1,
+            MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO,
+            2
+        );
 
         private MockRuleFeatureType() {}
 
@@ -101,13 +107,8 @@ public String getName() {
         }
 
         @Override
-        public Map<String, Attribute> getAllowedAttributesRegistry() {
-            return Map.of(
-                ATTRIBUTE_VALUE_ONE,
-                MockRuleAttributes.MOCK_RULE_ATTRIBUTE_ONE,
-                ATTRIBUTE_VALUE_TWO,
-                MockRuleAttributes.MOCK_RULE_ATTRIBUTE_TWO
-            );
+        public Map<Attribute, Integer> getOrderedAttributes() {
+            return PRIORITIZED_ATTRIBUTES;
         }
     }
 

modules/autotagging-commons/spi/src/main/java/org/opensearch/rule/spi/AttributesExtension.java

@@ -0,0 +1,22 @@
+/*
+ * 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.rule.spi;
+
+import org.opensearch.rule.autotagging.Attribute;
+
+/**
+ * Represents a custom attribute extension for the rule framework.
+ * Implementations provide a single attribute to be used in auto tagging.
+ */
+public interface AttributesExtension {
+    /**
+     * Returns the attribute provided by this extension.
+     */
+    Attribute getAttribute();
+}