Added more tests and resolved PR comments

Signed-off-by: Owais <owaiskazi19@gmail.com>

src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java

@@ -33,7 +33,6 @@
 import org.opensearch.flowframework.model.State;
 import org.opensearch.flowframework.model.Template;
 import org.opensearch.flowframework.model.Workflow;
-import org.opensearch.flowframework.transport.handler.WorkflowFunction;
 import org.opensearch.flowframework.workflow.ProcessNode;
 import org.opensearch.flowframework.workflow.WorkflowProcessSorter;
 import org.opensearch.index.query.QueryBuilder;
@@ -137,15 +136,17 @@ private void resolveUserAndExecute(
         User requestedUser,
         String workflowId,
         ActionListener<WorkflowResponse> listener,
-        WorkflowFunction function
+        Runnable function
     ) {
         try {
             // Check if user has backend roles
             // When filter by is enabled, block users creating/updating workflows who do not have backend roles.
             if (filterByEnabled) {
-                String error = checkFilterByBackendRoles(requestedUser);
-                if (error != null) {
-                    listener.onFailure(new FlowFrameworkException(error, RestStatus.BAD_REQUEST));
+                try {
+                    checkFilterByBackendRoles(requestedUser);
+                } catch (FlowFrameworkException e) {
+                    logger.error(e.getMessage(), e);
+                    listener.onFailure(e);
                     return;
                 }
             }
@@ -159,7 +160,7 @@ private void resolveUserAndExecute(
                 getWorkflow(requestedUser, workflowId, filterByBackendRole, listener, function, client, clusterService, xContentRegistry);
             } else {
                 // Create Workflow. No need to get current workflow.
-                function.execute();
+                function.run();
             }
         } catch (Exception e) {
             String errorMessage = "Failed to create or update workflow";
@@ -181,7 +182,7 @@ private void resolveUserAndExecute(
      * @param user the user making the request
      * @param listener the action listener
      */
-    protected void createExecute(WorkflowRequest request, User user, ActionListener<WorkflowResponse> listener) {
+    private void createExecute(WorkflowRequest request, User user, ActionListener<WorkflowResponse> listener) {
         Instant creationTime = Instant.now();
         Template templateWithUser = new Template(
             request.getTemplate().name(),

src/main/java/org/opensearch/flowframework/transport/DeleteWorkflowTransportAction.java

@@ -109,7 +109,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Dele
      * @param listener the action listener
      * @param context the thread context
      */
-    protected void executeDeleteRequest(
+    private void executeDeleteRequest(
         WorkflowRequest request,
         ActionListener<DeleteResponse> listener,
         ThreadContext.StoredContext context

src/main/java/org/opensearch/flowframework/transport/ReprovisionWorkflowTransportAction.java

@@ -160,7 +160,7 @@ protected void doExecute(Task task, ReprovisionWorkflowRequest request, ActionLi
      * @param listener the action listener
      * @param context the thread context
      */
-    protected void executeReprovisionRequest(
+    private void executeReprovisionRequest(
         ReprovisionWorkflowRequest request,
         ActionListener<WorkflowResponse> listener,
         ThreadContext.StoredContext context

src/main/java/org/opensearch/flowframework/transport/SearchWorkflowStateTransportAction.java

@@ -10,12 +10,14 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.transport.handler.SearchHandler;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
@@ -43,6 +45,13 @@ public SearchWorkflowStateTransportAction(TransportService transportService, Act
 
     @Override
     protected void doExecute(Task task, SearchRequest request, ActionListener<SearchResponse> actionListener) {
-        searchHandler.search(request, actionListener);
+        try {
+            searchHandler.search(request, actionListener);
+        } catch (Exception e) {
+            String errorMessage = "Failed to search workflow states in global context";
+            logger.error(errorMessage, e);
+            actionListener.onFailure(new FlowFrameworkException(errorMessage, ExceptionsHelper.status(e)));
+        }
+
     }
 }

src/main/java/org/opensearch/flowframework/transport/SearchWorkflowTransportAction.java

@@ -10,12 +10,14 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.transport.handler.SearchHandler;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
@@ -43,6 +45,12 @@ public SearchWorkflowTransportAction(TransportService transportService, ActionFi
 
     @Override
     protected void doExecute(Task task, SearchRequest request, ActionListener<SearchResponse> actionListener) {
-        searchHandler.search(request, actionListener);
+        try {
+            searchHandler.search(request, actionListener);
+        } catch (Exception e) {
+            String errorMessage = "Failed to search workflows in global context";
+            logger.error(errorMessage, e);
+            actionListener.onFailure(new FlowFrameworkException(errorMessage, ExceptionsHelper.status(e)));
+        }
     }
 }

src/main/java/org/opensearch/flowframework/transport/handler/SearchHandler.java

@@ -31,7 +31,7 @@
 public class SearchHandler {
     private final Logger logger = LogManager.getLogger(SearchHandler.class);
     private final Client client;
-    private volatile Boolean filterEnabled;
+    private volatile Boolean filterByBackendRole;
 
     /**
      * Instantiates a new SearchHandler
@@ -42,8 +42,8 @@ public class SearchHandler {
      */
     public SearchHandler(Settings settings, ClusterService clusterService, Client client, Setting<Boolean> filterByBackendRoleSetting) {
         this.client = client;
-        filterEnabled = filterByBackendRoleSetting.get(settings);
-        clusterService.getClusterSettings().addSettingsUpdateConsumer(filterByBackendRoleSetting, it -> filterEnabled = it);
+        filterByBackendRole = filterByBackendRoleSetting.get(settings);
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(filterByBackendRoleSetting, it -> filterByBackendRole = it);
     }
 
     /**
@@ -78,7 +78,7 @@ public void validateRole(
         ActionListener<SearchResponse> listener,
         ThreadContext.StoredContext context
     ) {
-        if (user == null || !filterEnabled || isAdmin(user)) {
+        if (user == null || !filterByBackendRole || isAdmin(user)) {
             // Case 1: user == null when 1. Security is disabled. 2. When user is super-admin
             // Case 2: If Security is enabled and filter is disabled, proceed with search as
             // user is already authenticated to hit this API.

src/main/java/org/opensearch/flowframework/util/ParseUtils.java

@@ -35,7 +35,6 @@
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.model.Template;
 import org.opensearch.flowframework.transport.WorkflowResponse;
-import org.opensearch.flowframework.transport.handler.WorkflowFunction;
 import org.opensearch.flowframework.workflow.WorkflowData;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.NestedQueryBuilder;
@@ -53,7 +52,6 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
@@ -243,7 +241,7 @@ public static Instant parseInstant(XContentParser parser) throws IOException {
      */
     public static User getUserContext(Client client) {
         String userStr = client.threadPool().getThreadContext().getTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT);
-        logger.debug("Filtering result by " + userStr);
+        logger.debug("Filtering result by {}", userStr);
         return User.parse(userStr);
     }
 
@@ -271,8 +269,10 @@ public static SearchSourceBuilder addUserBackendRolesFilter(User user, SearchSou
         } else if (query instanceof BoolQueryBuilder) {
             ((BoolQueryBuilder) query).filter(boolQueryBuilder);
         } else {
-            // e.g., wild card query
-            throw new FlowFrameworkException("Search API does not support queries other than BoolQuery", RestStatus.BAD_REQUEST);
+            // Convert any other query to a BoolQueryBuilder
+            BoolQueryBuilder boolQuery = QueryBuilders.boolQuery().must(query);
+            boolQuery.filter(boolQueryBuilder);
+            searchSourceBuilder.query(boolQuery);
         }
         return searchSourceBuilder;
     }
@@ -291,9 +291,9 @@ public static SearchSourceBuilder addUserBackendRolesFilter(User user, SearchSou
     public static void resolveUserAndExecute(
         User requestedUser,
         String workflowId,
-        boolean filterByEnabled,
+        Boolean filterByEnabled,
         ActionListener listener,
-        WorkflowFunction function,
+        Runnable function,
         Client client,
         ClusterService clusterService,
         NamedXContentRegistry xContentRegistry
@@ -303,7 +303,7 @@ public static void resolveUserAndExecute(
                 // requestedUser == null means security is disabled or user is superadmin. In this case we don't need to
                 // check if request user have access to the workflow or not.
                 // !filterByEnabled means security is enabled and filterByEnabled is disabled
-                function.execute();
+                function.run();
             } else {
                 getWorkflow(requestedUser, workflowId, filterByEnabled, listener, function, client, clusterService, xContentRegistry);
             }
@@ -344,20 +344,21 @@ private static boolean checkUserPermissions(User requestedUser, User resourceUse
     /**
      * Check if filter by backend roles is enabled and validate the requested user
      * @param requestedUser the user to execute the request
-     * @return error message if validation fails, null otherwise
      */
-    public static String checkFilterByBackendRoles(User requestedUser) {
+    public static void checkFilterByBackendRoles(User requestedUser) {
         if (requestedUser == null) {
-            return "Filter by backend roles is enabled and User is null";
+            String errorMessage = "Filter by backend roles is enabled and User is null";
+            logger.error(errorMessage);
+            throw new FlowFrameworkException(errorMessage, RestStatus.BAD_REQUEST);
         }
         if (requestedUser.getBackendRoles().isEmpty()) {
-            return String.format(
-                Locale.ROOT,
-                "Filter by backend roles is enabled and User %s does not have backend roles configured",
-                requestedUser.getName()
-            );
+            String userErrorMessage = "Filter by backend roles is enabled, but User "
+                + requestedUser.getName()
+                + " does not have any backend roles configured";
+
+            logger.error(userErrorMessage);
+            throw new FlowFrameworkException(userErrorMessage, RestStatus.FORBIDDEN);
         }
-        return null;
     }
 
     /**
@@ -376,7 +377,7 @@ public static void getWorkflow(
         String workflowId,
         Boolean filterByEnabled,
         ActionListener listener,
-        WorkflowFunction function,
+        Runnable function,
         Client client,
         ClusterService clusterService,
         NamedXContentRegistry xContentRegistry
@@ -426,7 +427,7 @@ public static void onGetWorkflowResponse(
         String workflowId,
         Boolean filterByEnabled,
         ActionListener<WorkflowResponse> listener,
-        WorkflowFunction function,
+        Runnable function,
         NamedXContentRegistry xContentRegistry
     ) {
         if (response.isExists()) {
@@ -438,7 +439,7 @@ public static void onGetWorkflowResponse(
                 User resourceUser = template.getUser();
 
                 if (!filterByEnabled || checkUserPermissions(requestUser, resourceUser, workflowId) || isAdmin(requestUser)) {
-                    function.execute();
+                    function.run();
                 } else {
                     logger.debug("User: " + requestUser.getName() + " does not have permissions to access workflow: " + workflowId);
                     listener.onFailure(

src/test/java/org/opensearch/flowframework/rest/FlowFrameworkSecureRestApiIT.java

@@ -276,7 +276,7 @@ public void testCreateWorkflowWithNoBackendRole() throws IOException {
         Template template = TestHelpers.createTemplateFromFile("register-deploylocalsparseencodingmodel.json");
         Exception exception = expectThrows(IOException.class, () -> { createWorkflow(dogClient, template); });
         assertTrue(
-            exception.getMessage().contains("Filter by backend roles is enabled and User dog does not have backend roles configured")
+            exception.getMessage().contains("Filter by backend roles is enabled, but User dog does not have any backend roles configured")
         );
     }