Updates exception handling with FlowFrameworkException AND adds dryrun param to Create Workflow

src/main/java/org/opensearch/flowframework/exception/FlowFrameworkException.java

@@ -9,11 +9,15 @@
 package org.opensearch.flowframework.exception;
 
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.core.xcontent.ToXContentObject;
+import org.opensearch.core.xcontent.XContentBuilder;
+
+import java.io.IOException;
 
 /**
  * Representation of Flow Framework Exceptions
  */
-public class FlowFrameworkException extends RuntimeException {
+public class FlowFrameworkException extends RuntimeException implements ToXContentObject {
 
     private static final long serialVersionUID = 1L;
 
@@ -60,4 +64,9 @@ public FlowFrameworkException(String message, Throwable cause, RestStatus restSt
     public RestStatus getRestStatus() {
         return restStatus;
     }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        return builder.startObject().field("error", "Request failed with exception: [" + this.getMessage() + "]").endObject();
+    }
 }

src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java

@@ -12,6 +12,7 @@
 import com.google.common.io.Resources;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.admin.indices.create.CreateIndexRequest;
 import org.opensearch.action.admin.indices.create.CreateIndexResponse;
 import org.opensearch.action.admin.indices.mapping.put.PutMappingRequest;
@@ -29,6 +30,7 @@
 import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.commons.authuser.User;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
@@ -148,7 +150,7 @@
                     }
                 }, e -> {
                     logger.error("Failed to create index " + indexName, e);
-                    internalListener.onFailure(e);
+                    internalListener.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
                 });
                 CreateIndexRequest request = new CreateIndexRequest(indexName).mapping(mapping).settings(indexSettings);
                 client.admin().indices().create(request, actionListener);
@@ -182,7 +184,12 @@
                                                     }
                                                 }, exception -> {
                                                     logger.error("Failed to update index setting for: " + indexName, exception);
-                                                    internalListener.onFailure(exception);
+                                                    internalListener.onFailure(
+                                                        new FlowFrameworkException(
+                                                            exception.getMessage(),
+                                                            ExceptionsHelper.status(exception)
+                                                        )
+                                                    );
                                                 }));
                                         } else {
                                             internalListener.onFailure(
@@ -191,7 +198,9 @@
                                         }
                                     }, exception -> {
                                         logger.error("Failed to update index " + indexName, exception);
-                                        internalListener.onFailure(exception);
+                                        internalListener.onFailure(
+                                            new FlowFrameworkException(exception.getMessage(), ExceptionsHelper.status(exception))
+                                        );
                                     })
                                 );
                         } else {
@@ -201,7 +210,7 @@
                         }
                     }, e -> {
                         logger.error("Failed to update index mapping", e);
-                        internalListener.onFailure(e);
+                        internalListener.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
                     }));
                 } else {
                     // No need to update index if it's already updated.
@@ -210,7 +219,7 @@
             }
         } catch (Exception e) {
             logger.error("Failed to init index " + indexName, e);
-            listener.onFailure(e);
+            listener.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
         }
     }
 
@@ -273,7 +282,7 @@
                 client.index(request, ActionListener.runBefore(listener, () -> context.restore()));
             } catch (Exception e) {
                 logger.error("Failed to index global_context index");
-                listener.onFailure(e);
+                listener.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
             }
         }, e -> {
             logger.error("Failed to create global_context index", e);
@@ -311,12 +320,12 @@
                 client.index(request, ActionListener.runBefore(listener, () -> context.restore()));
             } catch (Exception e) {
                 logger.error("Failed to put state index document", e);
-                listener.onFailure(e);
+                listener.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
             }
 
         }, e -> {
             logger.error("Failed to create global_context index", e);
-            listener.onFailure(e);
+            listener.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
         }));
     }
 
@@ -332,7 +341,7 @@
                 + documentId
                 + ", global_context index does not exist.";
             logger.error(exceptionMessage);
-            listener.onFailure(new Exception(exceptionMessage));
+            listener.onFailure(new FlowFrameworkException(exceptionMessage, RestStatus.BAD_REQUEST));
         } else {
             IndexRequest request = new IndexRequest(GLOBAL_CONTEXT_INDEX).id(documentId);
             try (
@@ -344,7 +353,7 @@
                 client.index(request, ActionListener.runBefore(listener, () -> context.restore()));
             } catch (Exception e) {
                 logger.error("Failed to update global_context entry : {}. {}", documentId, e.getMessage());
-                listener.onFailure(e);
+                listener.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
             }
         }
     }
@@ -365,7 +374,7 @@
         if (!doesIndexExist(indexName)) {
             String exceptionMessage = "Failed to update document for given workflow due to missing " + indexName + " index";
             logger.error(exceptionMessage);
-            listener.onFailure(new Exception(exceptionMessage));
+            listener.onFailure(new FlowFrameworkException(exceptionMessage, RestStatus.BAD_REQUEST));
         } else {
             try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
                 UpdateRequest updateRequest = new UpdateRequest(indexName, documentId);
@@ -377,7 +386,7 @@
                 client.update(updateRequest, ActionListener.runBefore(listener, () -> context.restore()));
             } catch (Exception e) {
                 logger.error("Failed to update {} entry : {}. {}", indexName, documentId, e.getMessage());
-                listener.onFailure(e);
+                listener.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
             }
         }
     }

src/main/java/org/opensearch/flowframework/model/Template.java

@@ -177,7 +177,7 @@
             }
         }
         if (name == null) {
-            throw new IOException("An template object requires a name.");
+            throw new IOException("A template object requires a name.");
         }
 
         return new Template(name, description, useCase, templateVersion, compatibilityVersion, workflows, user);

src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java

@@ -9,13 +9,20 @@
 package org.opensearch.flowframework.rest;
 
 import com.google.common.collect.ImmutableList;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
+import org.opensearch.core.xcontent.ToXContent;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.model.Template;
 import org.opensearch.flowframework.transport.CreateWorkflowAction;
 import org.opensearch.flowframework.transport.WorkflowRequest;
 import org.opensearch.rest.BaseRestHandler;
+import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestRequest;
-import org.opensearch.rest.action.RestToXContentListener;
 
 import java.io.IOException;
 import java.util.List;
@@ -29,6 +36,7 @@
  */
 public class RestCreateWorkflowAction extends BaseRestHandler {
 
+    private static final Logger logger = LogManager.getLogger(RestCreateWorkflowAction.class);
     private static final String CREATE_WORKFLOW_ACTION = "create_workflow_action";
 
     /**
@@ -53,11 +61,28 @@
 
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
-
-        String workflowId = request.param(WORKFLOW_ID);
-        Template template = Template.parse(request.content().utf8ToString());
-        WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template);
-        return channel -> client.execute(CreateWorkflowAction.INSTANCE, workflowRequest, new RestToXContentListener<>(channel));
+        try {
+            String workflowId = request.param(WORKFLOW_ID);
+            Template template = Template.parse(request.content().utf8ToString());
+            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template);
+            return channel -> client.execute(CreateWorkflowAction.INSTANCE, workflowRequest, ActionListener.wrap(response -> {
+                XContentBuilder builder = response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS);
+                channel.sendResponse(new BytesRestResponse(RestStatus.CREATED, builder));
+            }, exception -> {
+                try {
+                    FlowFrameworkException ex = (FlowFrameworkException) exception;
+                    XContentBuilder exceptionBuilder = ex.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS);
+                    channel.sendResponse(new BytesRestResponse(ex.getRestStatus(), exceptionBuilder));
+                } catch (IOException e) {
+                    logger.error("Failed to send back create workflow exception", e);
+                }
+            }));
+        } catch (Exception e) {
+            FlowFrameworkException ex = new FlowFrameworkException(e.getMessage(), RestStatus.BAD_REQUEST);
+            return channel -> channel.sendResponse(
+                new BytesRestResponse(ex.getRestStatus(), ex.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS))
+            );
+        }
     }
 
 }

src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java

@@ -9,14 +9,19 @@
 package org.opensearch.flowframework.rest;
 
 import com.google.common.collect.ImmutableList;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.core.xcontent.ToXContent;
+import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.transport.ProvisionWorkflowAction;
 import org.opensearch.flowframework.transport.WorkflowRequest;
 import org.opensearch.rest.BaseRestHandler;
+import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestRequest;
-import org.opensearch.rest.action.RestToXContentListener;
 
 import java.io.IOException;
 import java.util.List;
@@ -30,6 +35,8 @@
  */
 public class RestProvisionWorkflowAction extends BaseRestHandler {
 
+    private static final Logger logger = LogManager.getLogger(RestProvisionWorkflowAction.class);
+
     private static final String PROVISION_WORKFLOW_ACTION = "provision_workflow_action";
 
     /**
@@ -52,21 +59,35 @@
 
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
-
-        // Validate content
-        if (request.hasContent()) {
-            throw new FlowFrameworkException("Invalid request format", RestStatus.BAD_REQUEST);
-        }
-
-        // Validate params
         String workflowId = request.param(WORKFLOW_ID);
-        if (workflowId == null) {
-            throw new FlowFrameworkException("workflow_id cannot be null", RestStatus.BAD_REQUEST);
+        try {
+            // Validate content
+            if (request.hasContent()) {
+                throw new FlowFrameworkException("Invalid request format", RestStatus.BAD_REQUEST);
+            }
+            // Validate params
+            if (workflowId == null) {
+                throw new FlowFrameworkException("workflow_id cannot be null", RestStatus.BAD_REQUEST);
+            }
+            // Create request and provision
+            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null);
+            return channel -> client.execute(ProvisionWorkflowAction.INSTANCE, workflowRequest, ActionListener.wrap(response -> {
+                XContentBuilder builder = response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS);
+                channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder));
+            }, exception -> {
+                try {
+                    FlowFrameworkException ex = (FlowFrameworkException) exception;
+                    XContentBuilder exceptionBuilder = ex.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS);
+                    channel.sendResponse(new BytesRestResponse(ex.getRestStatus(), exceptionBuilder));
+                } catch (IOException e) {
+                    logger.error("Failed to send back provision workflow exception", e);
+                }
+            }));
+        } catch (FlowFrameworkException ex) {
+            return channel -> channel.sendResponse(
+                new BytesRestResponse(ex.getRestStatus(), ex.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS))
+            );
         }
-
-        // Create request and provision
-        WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null);
-        return channel -> client.execute(ProvisionWorkflowAction.INSTANCE, workflowRequest, new RestToXContentListener<>(channel));
     }
 
 }

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

@@ -11,6 +11,7 @@
 import com.google.common.collect.ImmutableMap;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
 import org.opensearch.client.Client;
@@ -83,12 +84,21 @@
                         listener.onResponse(new WorkflowResponse(globalContextResponse.getId()));
                     }, exception -> {
                         logger.error("Failed to save workflow state : {}", exception.getMessage());
-                        listener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.BAD_REQUEST));
+                        if (exception instanceof FlowFrameworkException) {
+                            listener.onFailure(exception);
+                        } else {
+                            listener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.BAD_REQUEST));
+                        }
                     })
                 );
             }, exception -> {
                 logger.error("Failed to save use case template : {}", exception.getMessage());
-                listener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.INTERNAL_SERVER_ERROR));
+                if (exception instanceof FlowFrameworkException) {
+                    listener.onFailure(exception);
+                } else {
+                    listener.onFailure(new FlowFrameworkException(exception.getMessage(), ExceptionsHelper.status(exception)));
+                }
+
             }));
         } else {
             // Update existing entry, full document replacement
@@ -105,12 +115,21 @@
                             listener.onResponse(new WorkflowResponse(request.getWorkflowId()));
                         }, exception -> {
                             logger.error("Failed to update workflow state : {}", exception.getMessage());
-                            listener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.BAD_REQUEST));
+                            if (exception instanceof FlowFrameworkException) {
+                                listener.onFailure(exception);
+                            } else {
+                                listener.onFailure(new FlowFrameworkException(exception.getMessage(), ExceptionsHelper.status(exception)));
+                            }
                         })
                     );
                 }, exception -> {
                     logger.error("Failed to updated use case template {} : {}", request.getWorkflowId(), exception.getMessage());
-                    listener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.INTERNAL_SERVER_ERROR));
+                    if (exception instanceof FlowFrameworkException) {
+                        listener.onFailure(exception);
+                    } else {
+                        listener.onFailure(new FlowFrameworkException(exception.getMessage(), ExceptionsHelper.status(exception)));
+                    }
+
                 })
             );
         }