[Rule-based auto tagging] Bug fix and improvements

CHANGELOG.md

@@ -23,6 +23,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add BooleanQuery rewrite moving constant-scoring must clauses to filter clauses ([#18510](https://github.com/opensearch-project/OpenSearch/issues/18510))
 - Add functionality for plugins to inject QueryCollectorContext during QueryPhase ([#18637](https://github.com/opensearch-project/OpenSearch/pull/18637))
 - Add support for non-timing info in profiler ([#18460](https://github.com/opensearch-project/OpenSearch/issues/18460))
+- [Rule-based auto tagging] Bug fix and improvements ([#18726](https://github.com/opensearch-project/OpenSearch/pull/18726))
 - Extend Approximation Framework to other numeric types ([#18530](https://github.com/opensearch-project/OpenSearch/issues/18530))
 - Add Semantic Version field type mapper and extensive unit tests([#18454](https://github.com/opensearch-project/OpenSearch/pull/18454))
 - Pass index settings to system ingest processor factories. ([#18708](https://github.com/opensearch-project/OpenSearch/pull/18708))

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

@@ -49,14 +49,7 @@ public CreateRuleRequest(StreamInput in) throws IOException {
 
     @Override
     public ActionRequestValidationException validate() {
-        try {
-            rule.getFeatureType().getFeatureValueValidator().validate(rule.getFeatureValue());
-            return null;
-        } catch (Exception e) {
-            ActionRequestValidationException validationException = new ActionRequestValidationException();
-            validationException.addValidationError("Validation failed: " + e.getMessage());
-            return validationException;
-        }
+        return null;
     }
 
     @Override

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

@@ -114,6 +114,11 @@ private List<String> validateFeatureType() {
         if (featureType == null) {
             return List.of("Couldn't identify which feature the rule belongs to. Rule feature can't be null.");
         }
+        try {
+            featureType.getFeatureValueValidator().validate(featureValue);
+        } catch (Exception e) {
+            return List.of(e.getMessage());
+        }
         return new ArrayList<>();
     }
 

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

@@ -16,6 +16,7 @@
 import org.opensearch.action.index.IndexRequest;
 import org.opensearch.action.search.SearchRequestBuilder;
 import org.opensearch.action.search.SearchResponse;
+import org.opensearch.action.support.WriteRequest;
 import org.opensearch.action.support.clustermanager.AcknowledgedResponse;
 import org.opensearch.action.update.UpdateRequest;
 import org.opensearch.cluster.service.ClusterService;
@@ -183,6 +184,7 @@ public void onFailure(Exception e) {
     private void persistRule(Rule rule, ActionListener<CreateRuleResponse> listener) {
         try {
             IndexRequest indexRequest = new IndexRequest(indexName).id(rule.getId())
+                .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
                 .source(rule.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS));
             client.index(indexRequest).get();
             listener.onResponse(new CreateRuleResponse(rule));
@@ -312,9 +314,8 @@ public void onFailure(Exception e) {
      */
     private void persistUpdatedRule(String ruleId, Rule updatedRule, ActionListener<UpdateRuleResponse> listener) {
         try {
-            UpdateRequest updateRequest = new UpdateRequest(indexName, ruleId).doc(
-                updatedRule.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS)
-            );
+            UpdateRequest updateRequest = new UpdateRequest(indexName, ruleId).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
+                .doc(updatedRule.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS));
             client.update(updateRequest).get();
             listener.onResponse(new UpdateRuleResponse(updatedRule));
         } catch (Exception e) {

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

@@ -23,6 +23,9 @@
 import static org.opensearch.rule.autotagging.RuleTests.TestAttribute.TEST_ATTRIBUTE_1;
 import static org.opensearch.rule.autotagging.RuleTests.UPDATED_AT;
 import static org.opensearch.rule.autotagging.RuleTests._ID;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class RuleValidatorTests extends OpenSearchTestCase {
 
@@ -118,4 +121,14 @@ public void testEqualRuleValidators() {
         RuleValidator otherValidator = new RuleValidator(_ID, DESCRIPTION, ATTRIBUTE_MAP, FEATURE_VALUE, UPDATED_AT, FEATURE_TYPE);
         assertEquals(validator, otherValidator);
     }
+
+    public void testFeatureValueValidationThrows() {
+        FeatureType mockFeatureType = mock(FeatureType.class);
+        FeatureValueValidator mockValidator = mock(FeatureValueValidator.class);
+        when(mockFeatureType.getFeatureValueValidator()).thenReturn(mockValidator);
+        doThrow(new IllegalArgumentException("Invalid feature value")).when(mockValidator).validate("bad-value");
+        RuleValidator validator = new RuleValidator(_ID, DESCRIPTION, ATTRIBUTE_MAP, "bad-value", UPDATED_AT, mockFeatureType);
+        IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, validator::validate);
+        assertTrue(ex.getMessage().contains("Invalid feature value"));
+    }
 }

modules/autotagging-commons/src/main/java/org/opensearch/rule/rest/RestGetRuleAction.java

@@ -74,10 +74,12 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
         }
 
         final FeatureType featureType = FeatureType.from(request.param(FEATURE_TYPE));
-        final Set<String> excludedKeys = Set.of(FEATURE_TYPE, ID_STRING, SEARCH_AFTER_STRING, "pretty");
-        final List<String> requestParams = request.params().keySet().stream().filter(key -> !excludedKeys.contains(key)).toList();
-
-        for (String attributeName : requestParams) {
+        final List<String> attributeParams = request.params()
+            .keySet()
+            .stream()
+            .filter(key -> featureType.getAllowedAttributesRegistry().containsKey(key))
+            .toList();
+        for (String attributeName : attributeParams) {
             Attribute attribute = featureType.getAttributeFromName(attributeName);
             if (attribute == null) {
                 throw new IllegalArgumentException(attributeName + " is not a valid attribute under feature type " + featureType.getName());

modules/autotagging-commons/src/test/java/org/opensearch/rule/rest/RestGetRuleActionTests.java

@@ -8,10 +8,26 @@
 
 package org.opensearch.rule.rest;
 
+import org.opensearch.rest.RestRequest;
+import org.opensearch.rest.RestRequest.Method;
+import org.opensearch.rule.InMemoryRuleProcessingServiceTests;
+import org.opensearch.rule.autotagging.FeatureType;
 import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.test.rest.FakeRestRequest;
+import org.junit.Before;
+
+import java.util.Map;
+
+import static org.opensearch.rule.rest.RestGetRuleAction.FEATURE_TYPE;
 
 public class RestGetRuleActionTests extends OpenSearchTestCase {
-    RestGetRuleAction action = new RestGetRuleAction();;
+
+    private RestGetRuleAction action;
+
+    @Before
+    public void setUpAction() {
+        action = new RestGetRuleAction();
+    }
 
     public void testGetName() {
         assertEquals("get_rule", action.getName());
@@ -23,4 +39,12 @@ public void testRoutes() {
         assertTrue(routes.stream().anyMatch(r -> r.getPath().equals("_rules/{featureType}/")));
         assertTrue(routes.stream().anyMatch(r -> r.getPath().equals("_rules/{featureType}/{id}")));
     }
+
+    public void testPrepareRequestFiltersAllowedAttributes() {
+        FeatureType featureType = InMemoryRuleProcessingServiceTests.WLMFeatureType.WLM;
+        String validAttrName = featureType.getAllowedAttributesRegistry().keySet().iterator().next();
+        Map<String, String> params = Map.of(FEATURE_TYPE, featureType.getName(), validAttrName, "value1", "invalidAttr", "shouldBeIgnored");
+        RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(Method.GET).withParams(params).build();
+        assertNotNull(action.prepareRequest(request, null));
+    }
 }

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/sync/RefreshBasedSyncMechanism.java

@@ -15,6 +15,7 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.plugin.wlm.WlmClusterSettingValuesProvider;
 import org.opensearch.plugin.wlm.rule.sync.detect.RuleEvent;
 import org.opensearch.plugin.wlm.rule.sync.detect.RuleEventClassifier;
@@ -121,6 +122,10 @@ public void onResponse(GetRuleResponse response) {
 
                 @Override
                 public void onFailure(Exception e) {
+                    if (e instanceof IndexNotFoundException) {
+                        logger.debug("Rule index not found, skipping rule processing.");
+                        return;
+                    }
                     logger.warn("Failed to get rules from persistence service", e);
                 }
             }

plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/rule/sync/RefreshBasedSyncMechanismTests.java

@@ -11,6 +11,7 @@
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.plugin.wlm.AutoTaggingActionFilterTests;
 import org.opensearch.plugin.wlm.WlmClusterSettingValuesProvider;
 import org.opensearch.plugin.wlm.WorkloadManagementPlugin;
@@ -247,4 +248,17 @@ public void test_doRun_RefreshesRulesAndCheckInMemoryView() {
         // Here 1 is due to add in the second run and 10 for adding 10 rules as part of first run
         verify(ruleProcessingService, times(updateEventCount + 1 + 10)).add(any(Rule.class));
     }
+
+    @SuppressWarnings("unchecked")
+    public void testDoRunIgnoresIndexNotFoundException() {
+        doAnswer(invocation -> {
+            ActionListener<GetRuleResponse> listener = invocation.getArgument(1);
+            listener.onFailure(new IndexNotFoundException("rules index not found"));
+            return null;
+        }).when(rulePersistenceService).getRule(any(GetRuleRequest.class), any(ActionListener.class));
+        sut.doRun();
+        verify(rulePersistenceService, times(1)).getRule(any(GetRuleRequest.class), any(ActionListener.class));
+        verify(ruleProcessingService, times(0)).add(any(Rule.class));
+        verify(ruleProcessingService, times(0)).remove(any(Rule.class));
+    }
 }