revise based on comments

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

Author: ruai0511

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

@@ -10,6 +10,8 @@
 
 import org.opensearch.rule.autotagging.Attribute;
 
+import java.util.TreeMap;
+
 /**
  * Generic Rule attributes that features can use out of the use by using the lib.
  * @opensearch.experimental
@@ -32,6 +34,11 @@ public String getName() {
         return name;
     }
 
+    @Override
+    public TreeMap<Integer, String> getPrioritizedSubfields() {
+        return new TreeMap<>();
+    }
+
     /**
      * Retrieves the RuleAttribute from a name string
      * @param name - attribute name

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

@@ -19,6 +19,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
 
 /**
  * Security attribute for the rules. Example:
@@ -30,20 +32,20 @@
  */
 public enum SecurityAttribute implements Attribute {
     /**
-     * Represents the index_pattern attribute in RuleAttribute
+     * Represents the principal attribute
      */
     PRINCIPAL("principal");
 
     /**
-     * Key representing the username subfield.
+     * Represents the username subfield
      */
     public static final String USERNAME = "username";
     /**
-     * Key representing the role subfield.
+     * Represents the role subfield
      */
     public static final String ROLE = "role";
+    private static final TreeMap<Integer, String> PRIORITIZED_SUBFIELDS = new TreeMap<>(Map.of(1, USERNAME, 2, ROLE));
     private final String name;
-    private static final List<String> ALLOWED_SUBFIELDS = List.of(USERNAME, ROLE);
 
     SecurityAttribute(String name) {
         this.name = name;
@@ -56,10 +58,20 @@ public String getName() {
     }
 
     @Override
-    public List<String> getPrioritizedSubfields() {
-        return ALLOWED_SUBFIELDS;
+    public TreeMap<Integer, String> getPrioritizedSubfields() {
+        return PRIORITIZED_SUBFIELDS;
     }
 
+    /**
+     * Parses the attribute values for security attribute
+     * Example:
+     * {
+     *     "username": ["alice"],
+     *     "role": ["all_access"]
+     * }
+     * will be parsed into a set with values "username|alice" and "role|all_access"
+     * @param parser the XContent parser
+     */
     @Override
     public Set<String> fromXContentParseAttributeValues(XContentParser parser) throws IOException {
         Set<String> resultSet = new HashSet<>();
@@ -70,32 +82,21 @@ public Set<String> fromXContentParseAttributeValues(XContentParser parser) throw
                 "Expected START_OBJECT token for " + getName() + " attribute but got " + parser.currentToken()
             );
         }
+        List<String> allowedSubfieldsName = PRIORITIZED_SUBFIELDS.values().stream().toList();
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
             String subFieldName = parser.currentName();
             parser.nextToken();
-            if (!ALLOWED_SUBFIELDS.contains(subFieldName)) {
-                throw new XContentParseException(
-                    parser.getTokenLocation(),
-                    "Invalid field: " + subFieldName + ". Allowed fields are: " + String.join(", ", ALLOWED_SUBFIELDS)
-                );
-            }
-            if (parser.currentToken() != XContentParser.Token.START_ARRAY) {
+            if (!allowedSubfieldsName.contains(subFieldName)) {
                 throw new XContentParseException(
                     parser.getTokenLocation(),
-                    "Expected array for field: " + subFieldName + " but got " + parser.currentToken()
+                    "Invalid field: " + subFieldName + ". Allowed fields are: " + String.join(", ", allowedSubfieldsName)
                 );
             }
-            while (parser.nextToken() != XContentParser.Token.END_ARRAY) {
-                if (parser.currentToken() == XContentParser.Token.VALUE_STRING) {
-                    // prefix each value with the subFieldName (e.g., "username_name1")
-                    resultSet.add(String.join("_", subFieldName, parser.text()));
-                } else {
-                    throw new XContentParseException(
-                        parser.getTokenLocation(),
-                        "Expected string value in array under '" + subFieldName + "', but got " + parser.currentToken()
-                    );
-                }
-            }
+            resultSet.addAll(
+                Attribute.super.fromXContentParseAttributeValues(parser).stream()
+                    .map(s -> subFieldName + '|' + s)
+                    .collect(Collectors.toSet())
+            );
         }
 
         return resultSet;
@@ -105,21 +106,18 @@ public Set<String> fromXContentParseAttributeValues(XContentParser parser) throw
     public void toXContentWriteAttributeValues(XContentBuilder builder, Set<String> values) throws IOException {
         builder.startObject(getName());
         Map<String, Set<String>> grouped = new HashMap<>();
-
-        // For each string in the values set, split it into two parts using the first underscore as delimiter:
+        // For each string in the values set, split it into two parts using the first '|' as delimiter:
         // parts[0] is the prefix (e.g., "username" or "role")
         // parts[1] is the actual value (e.g., "name1", "role1")
         for (String value : values) {
-            String[] parts = value.split("_", 2);
+            String[] parts = value.split("\\|", 2);
             if (parts.length == 2) {
                 grouped.computeIfAbsent(parts[0], k -> new HashSet<>()).add(parts[1]);
             }
         }
-
         for (Map.Entry<String, Set<String>> entry : grouped.entrySet()) {
             builder.array(entry.getKey(), entry.getValue().toArray(new String[0]));
         }
-
         builder.endObject();
     }
 }

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

@@ -20,7 +20,7 @@ public interface AttributeExtractor<V> {
      * Defines the combination style used when a request contains multiple values
      * for an attribute.
      */
-    enum CombinationStyle {
+    enum LogicalOperator {
         /**
          * Logical AND
          */
@@ -44,11 +44,11 @@ enum CombinationStyle {
     Iterable<V> extract();
 
     /**
-     * Returns the combination style used when a request contains multiple values
+     * Returns the logical operator used when a request contains multiple values
      * for an attribute.
      * For example, if the request targets both index A and B, then a rule must
-     * have both index A and B as attributes, requiring an AND combination.
-     * @return the combination style (e.g., AND, OR)
+     * have both index A and B as attributes, requiring an AND operator.
+     * @return the logical operator (e.g., AND, OR)
      */
-    CombinationStyle getCombinationStyle();
+    LogicalOperator getLogicalOperator();
 }

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

@@ -16,10 +16,9 @@
 import org.opensearch.core.xcontent.XContentParser;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Set;
+import java.util.TreeMap;
 
 /**
  * Represents an attribute within the auto-tagging feature. Attributes define characteristics that can
@@ -37,11 +36,9 @@ public interface Attribute extends Writeable {
     String getName();
 
     /**
-     * Returns the allowed subfields ordered from highest to lowest priority
+     * Returns a map of subfields ordered by priority, where 1 represents the highest priority.
      */
-    default List<String> getPrioritizedSubfields() {
-        return new ArrayList<>();
-    }
+    TreeMap<Integer, String> getPrioritizedSubfields();
 
     /**
      * Ensure that `validateAttribute` is called in the constructor of attribute implementations
@@ -60,7 +57,9 @@ default void writeTo(StreamOutput out) throws IOException {
     }
 
     /**
-     * Parses attribute values for specific attributes
+     * 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 {

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

@@ -13,8 +13,8 @@
 import org.opensearch.core.common.io.stream.Writeable;
 
 import java.io.IOException;
-import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 /**
  * Represents a feature type within the auto-tagging feature. Feature types define different categories of
@@ -43,16 +43,18 @@ public interface FeatureType extends Writeable {
     String getName();
 
     /**
-     * Returns the registry of allowed attributes for this feature type.
-     * Implementations must ensure that access to this registry is thread-safe.
+     * 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<String, Attribute> getAllowedAttributesRegistry();
+    Map<Attribute, Integer> getOrderedAttributes();
 
     /**
-     * Returns the list of attributes ordered by their processing priority.
-     * @return List of prioritized attributes.
+     * Returns the registry of allowed attributes for this feature type.
+     * Implementations must ensure that access to this registry is thread-safe.
      */
-    List<Attribute> getPrioritizedAttributesList();
+    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/storage/DefaultAttributeValueStore.java

@@ -10,7 +10,6 @@
 
 import org.apache.commons.collections4.trie.PatriciaTrie;
 
-import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -72,19 +71,17 @@ public void remove(K key, V value) {
     @Override
     public List<Set<V>> get(String key) {
         readLock.lock();
-        List<Set<V>> results = new ArrayList<>();
         try {
-            for (int len = key.length(); len >= 0; len--) {
-                String prefix = key.substring(0, len);
-                Set<V> values = trie.get(prefix);
-                if (values != null && !values.isEmpty()) {
-                    results.add(values);
-                }
-            }
+            return trie.keySet()
+                .stream()
+                .filter(k -> key.startsWith(k))
+                .sorted((a, b) -> Integer.compare(b.length(), a.length()))
+                .map(trie::get)
+                .filter(v -> v != null && !v.isEmpty())
+                .toList();
         } finally {
             readLock.unlock();
         }
-        return results;
     }
 
     @Override

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

@@ -45,20 +45,20 @@ public void testFromXContent() throws IOException {
 
         parser.nextToken();
         Set<String> result = SecurityAttribute.PRINCIPAL.fromXContentParseAttributeValues(parser);
-        assertTrue(result.contains("username_alice"));
-        assertTrue(result.contains("username_bob"));
-        assertTrue(result.contains("role_admin"));
+        assertTrue(result.contains("username|alice"));
+        assertTrue(result.contains("username|bob"));
+        assertTrue(result.contains("role|admin"));
     }
 
     public void testToXContent() throws IOException {
-        Set<String> input = Set.of("username_alice", "username_bob_lastname", "role_admin_admin");
+        Set<String> input = Set.of("username|alice", "username|bob_lastname", "role|admin|admin");
         XContentBuilder builder = XContentFactory.jsonBuilder();
         builder.startObject();
         SecurityAttribute.PRINCIPAL.toXContentWriteAttributeValues(builder, input);
         builder.endObject();
         String json = builder.toString();
         assertTrue(json.contains("alice"));
         assertTrue(json.contains("bob_lastname"));
-        assertTrue(json.contains("admin_admin"));
+        assertTrue(json.contains("admin|admin"));
     }
 }

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

@@ -18,9 +18,9 @@
 
 import java.io.IOException;
 import java.time.Instant;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 
 import static org.opensearch.rule.autotagging.RuleTests.TestAttribute.TEST_ATTRIBUTE_1;
 import static org.opensearch.rule.autotagging.RuleTests.TestAttribute.TEST_ATTRIBUTE_2;
@@ -71,6 +71,11 @@ public enum TestAttribute implements Attribute {
         public String getName() {
             return name;
         }
+
+        @Override
+        public TreeMap<Integer, String> getPrioritizedSubfields() {
+            return new TreeMap<>();
+        }
     }
 
     public static class TestFeatureType implements FeatureType {
@@ -96,6 +101,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;
@@ -110,11 +120,6 @@ public int getMaxCharLengthPerAttributeValue() {
         public Map<String, Attribute> getAllowedAttributesRegistry() {
             return ALLOWED_ATTRIBUTES;
         }
-
-        @Override
-        public List<Attribute> getPrioritizedAttributesList() {
-            return List.of(TEST_ATTRIBUTE_1, TEST_ATTRIBUTE_2);
-        }
     }
 
     static Rule buildRule(

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

@@ -26,18 +26,18 @@ public void setUp() throws Exception {
 
     public void testPut() {
         subjectUnderTest.put("foo", "bar");
-        assertEquals("bar", subjectUnderTest.get("foo").get());
+        assertEquals("bar", subjectUnderTest.get("foo").getFirst().iterator().next());
     }
 
     public void testRemove() {
         subjectUnderTest.put("foo", "bar");
-        subjectUnderTest.remove("foo");
+        subjectUnderTest.remove("foo", "bar");
         assertEquals(0, subjectUnderTest.size());
     }
 
     public void tesGet() {
         subjectUnderTest.put("foo", "bar");
-        assertEquals("bar", subjectUnderTest.get("foo").get());
+        assertEquals("bar", subjectUnderTest.get("foo").getFirst());
     }
 
     public void testGetWhenNoProperPrefixIsPresent() {