Add synchronous execution option to workflow provisioning (#990)

* Add synchronous execution option to workflow provisioning Signed-off-by: Junwei Dai <[email protected]> * code refactor Signed-off-by: Junwei Dai <[email protected]> * add change log Signed-off-by: Junwei Dai <[email protected]> * refactor code based on comment Signed-off-by: Junwei Dai <[email protected]> * fix spotless check Signed-off-by: Junwei Dai <[email protected]> * Limit workflow timeout to a range of 1 to 300 seconds Signed-off-by: Junwei Dai <[email protected]> * Limit workflow timeout to a range of 1 to 300 seconds Signed-off-by: Junwei Dai <[email protected]> # Conflicts: # src/main/java/org/opensearch/flowframework/util/WorkflowTimeoutUtility.java * Limit workflow timeout to non-negative Signed-off-by: Junwei Dai <[email protected]> * Add synchronous execution to reprovision Signed-off-by: Junwei Dai <[email protected]> * remove unsued common value Signed-off-by: Junwei Dai <[email protected]> * add reprovision sync execution Signed-off-by: Junwei Dai <[email protected]> # Conflicts: # src/test/java/org/opensearch/flowframework/workflow/DeleteConnectorStepTests.java * fix test for WorkflowTimeoutUtilityTests Signed-off-by: Junwei Dai <[email protected]> * fix test name for WorkflowTimeoutUtilityTests Signed-off-by: Junwei Dai <[email protected]> * Add comments to explain AtomicBoolean usage in WorkflowTimeoutUtility, update error message Signed-off-by: Junwei Dai <[email protected]> * fix spotless check Signed-off-by: Junwei Dai <[email protected]> * addressed some comments Signed-off-by: Junwei Dai <[email protected]> --------- Signed-off-by: Junwei Dai <[email protected]> Co-authored-by: Junwei Dai <[email protected]>
opensearch-project · Jan 16, 2025 · 33579a3 · 33579a3
1 parent 06b7831
commit 33579a3
Show file tree

Hide file tree

Showing 18 changed files with 1,026 additions and 39 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -18,6 +18,8 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
 
 ## [Unreleased 2.x](https://github.com/opensearch-project/flow-framework/compare/2.18...2.x)
 ### Features
+- Add synchronous execution option to workflow provisioning ([#990](https://github.com/opensearch-project/flow-framework/pull/990))
+
 ### Enhancements
 ### Bug Fixes
 - Remove useCase and defaultParams field in WorkflowRequest ([#758](https://github.com/opensearch-project/flow-framework/pull/758))

diff --git a/src/main/java/org/opensearch/flowframework/common/CommonValue.java b/src/main/java/org/opensearch/flowframework/common/CommonValue.java
@@ -74,6 +74,8 @@ private CommonValue() {}
     public static final String PROVISION_WORKFLOW = "provision";
     /** The param name for update workflow field in create API */
     public static final String UPDATE_WORKFLOW_FIELDS = "update_fields";
+    /** The param name for specifying the timeout duration in seconds to wait for workflow completion */
+    public static final String WAIT_FOR_COMPLETION_TIMEOUT = "wait_for_completion_timeout";
     /** The field name for workflow steps. This field represents the name of the workflow steps to be fetched. */
     public static final String WORKFLOW_STEP = "workflow_step";
     /** The param name for default use case, used by the create workflow API */
@@ -186,6 +188,8 @@ private CommonValue() {}
     public static final String SOURCE_INDEX = "source_index";
     /** The destination index field for reindex */
     public static final String DESTINATION_INDEX = "destination_index";
+    /** Provision Timeout field */
+    public static final String PROVISION_TIMEOUT_FIELD = "provision.timeout";
     /*
      * Constants associated with resource provisioning / state
      */

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java
@@ -12,6 +12,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.ToXContent;
@@ -43,6 +44,7 @@
 import static org.opensearch.flowframework.common.CommonValue.UPDATE_WORKFLOW_FIELDS;
 import static org.opensearch.flowframework.common.CommonValue.USE_CASE;
 import static org.opensearch.flowframework.common.CommonValue.VALIDATION;
+import static org.opensearch.flowframework.common.CommonValue.WAIT_FOR_COMPLETION_TIMEOUT;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI;
 import static org.opensearch.flowframework.common.FlowFrameworkSettings.FLOW_FRAMEWORK_ENABLED;
@@ -88,6 +90,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
         boolean reprovision = request.paramAsBoolean(REPROVISION_WORKFLOW, false);
         boolean updateFields = request.paramAsBoolean(UPDATE_WORKFLOW_FIELDS, false);
         String useCase = request.param(USE_CASE);
+        TimeValue waitForCompletionTimeout = request.paramAsTime(WAIT_FOR_COMPLETION_TIMEOUT, TimeValue.MINUS_ONE);
 
         // If provisioning, consume all other params and pass to provision transport action
         Map<String, String> params = provision
@@ -145,6 +148,15 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             );
             return processError(ffe, params, request);
         }
+        // Ensure wait_for_completion is not set unless reprovision or provision is true
+        if (waitForCompletionTimeout != TimeValue.MINUS_ONE && !(reprovision || provision)) {
+            FlowFrameworkException ffe = new FlowFrameworkException(
+                "Request parameters 'wait_for_completion_timeout' are not allowed unless the 'provision' or 'reprovision' parameter is set to true.",
+                RestStatus.BAD_REQUEST
+            );
+            return processError(ffe, params, request);
+        }
+
         try {
             Template template;
             Map<String, String> useCaseDefaultsMap = Collections.emptyMap();
@@ -219,7 +231,9 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             if (updateFields) {
                 params = Map.of(UPDATE_WORKFLOW_FIELDS, "true");
             }
-
+            if (waitForCompletionTimeout != TimeValue.MINUS_ONE) {
+                params = Map.of(WAIT_FOR_COMPLETION_TIMEOUT, waitForCompletionTimeout.toString());
+            }
             WorkflowRequest workflowRequest = new WorkflowRequest(
                 workflowId,
                 template,

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java
@@ -12,6 +12,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.ToXContent;
@@ -33,6 +34,7 @@
 import java.util.stream.Collectors;
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+import static org.opensearch.flowframework.common.CommonValue.WAIT_FOR_COMPLETION_TIMEOUT;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI;
 import static org.opensearch.flowframework.common.FlowFrameworkSettings.FLOW_FRAMEWORK_ENABLED;
@@ -73,6 +75,7 @@ public List<Route> routes() {
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
         String workflowId = request.param(WORKFLOW_ID);
+        TimeValue waitForCompletionTimeout = request.paramAsTime(WAIT_FOR_COMPLETION_TIMEOUT, TimeValue.MINUS_ONE);
         try {
             Map<String, String> params = parseParamsAndContent(request);
             if (!flowFrameworkFeatureEnabledSetting.isFlowFrameworkEnabled()) {
@@ -86,7 +89,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
                 throw new FlowFrameworkException("workflow_id cannot be null", RestStatus.BAD_REQUEST);
             }
             // Create request and provision
-            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null, params);
+            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null, params, waitForCompletionTimeout);
             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));

diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
@@ -53,6 +53,7 @@
 import static org.opensearch.flowframework.common.CommonValue.GLOBAL_CONTEXT_INDEX;
 import static org.opensearch.flowframework.common.CommonValue.PROVISIONING_PROGRESS_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.STATE_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.WAIT_FOR_COMPLETION_TIMEOUT;
 import static org.opensearch.flowframework.common.FlowFrameworkSettings.FILTER_BY_BACKEND_ROLES;
 import static org.opensearch.flowframework.util.ParseUtils.checkFilterByBackendRoles;
 import static org.opensearch.flowframework.util.ParseUtils.getUserContext;
@@ -214,6 +215,16 @@ private void createExecute(WorkflowRequest request, User user, ActionListener<Wo
             }
         }
         String workflowId = request.getWorkflowId();
+        TimeValue waitForTimeCompletion;
+        if (request.getParams().containsKey(WAIT_FOR_COMPLETION_TIMEOUT)) {
+            waitForTimeCompletion = TimeValue.parseTimeValue(
+                request.getParams().get(WAIT_FOR_COMPLETION_TIMEOUT),
+                WAIT_FOR_COMPLETION_TIMEOUT
+            );
+        } else {
+            // default to minus one indicate async execution
+            waitForTimeCompletion = TimeValue.MINUS_ONE;
+        }
         if (workflowId == null) {
             // This is a new workflow (POST)
             // Throttle incoming requests
@@ -251,7 +262,8 @@ private void createExecute(WorkflowRequest request, User user, ActionListener<Wo
                                                     WorkflowRequest workflowRequest = new WorkflowRequest(
                                                         globalContextResponse.getId(),
                                                         null,
-                                                        request.getParams()
+                                                        request.getParams(),
+                                                        waitForTimeCompletion
                                                     );
                                                     logger.info(
                                                         "Provisioning parameter is set, continuing to provision workflow {}",
@@ -261,7 +273,14 @@ private void createExecute(WorkflowRequest request, User user, ActionListener<Wo
                                                         ProvisionWorkflowAction.INSTANCE,
                                                         workflowRequest,
                                                         ActionListener.wrap(provisionResponse -> {
-                                                            listener.onResponse(new WorkflowResponse(provisionResponse.getWorkflowId()));
+                                                            listener.onResponse(
+                                                                (workflowRequest.getWaitForCompletionTimeout() == TimeValue.MINUS_ONE)
+                                                                    ? new WorkflowResponse(provisionResponse.getWorkflowId())
+                                                                    : new WorkflowResponse(
+                                                                        provisionResponse.getWorkflowId(),
+                                                                        provisionResponse.getWorkflowState()
+                                                                    )
+                                                            );
                                                         }, exception -> {
                                                             String errorMessage = "Provisioning failed.";
                                                             logger.error(errorMessage, exception);
@@ -346,19 +365,26 @@ private void createExecute(WorkflowRequest request, User user, ActionListener<Wo
                                 .build();
 
                         if (request.isReprovision()) {
-
                             // Reprovision request
                             ReprovisionWorkflowRequest reprovisionRequest = new ReprovisionWorkflowRequest(
                                 getResponse.getId(),
                                 existingTemplate,
-                                template
+                                template,
+                                waitForTimeCompletion
                             );
                             logger.info("Reprovisioning parameter is set, continuing to reprovision workflow {}", getResponse.getId());
                             client.execute(
                                 ReprovisionWorkflowAction.INSTANCE,
                                 reprovisionRequest,
                                 ActionListener.wrap(reprovisionResponse -> {
-                                    listener.onResponse(new WorkflowResponse(reprovisionResponse.getWorkflowId()));
+                                    listener.onResponse(
+                                        reprovisionRequest.getWaitForCompletionTimeout() == TimeValue.MINUS_ONE
+                                            ? new WorkflowResponse(reprovisionResponse.getWorkflowId())
+                                            : new WorkflowResponse(
+                                                reprovisionResponse.getWorkflowId(),
+                                                reprovisionResponse.getWorkflowState()
+                                            )
+                                    );
                                 }, exception -> {
                                     String errorMessage = ParameterizedMessageFactory.INSTANCE.newMessage(
                                         "Reprovisioning failed for workflow {}",