CompletedFuture exceptionally if fields are not present

Signed-off-by: Owais Kazi <[email protected]>
opensearch-project · Oct 25, 2023 · 6163fc2 · 6163fc2
1 parent ae5c155
commit 6163fc2
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 1 deletion.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java
@@ -12,6 +12,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.ml.client.MachineLearningNodeClient;
 import org.opensearch.ml.common.connector.ConnectorAction;
@@ -129,6 +130,10 @@ public void onFailure(Exception e) {
                 .build();
 
             mlClient.createConnector(mlInput, actionListener);
+        } else {
+            createConnectorFuture.completeExceptionally(
+                new FlowFrameworkException("Required fields are not provided", RestStatus.BAD_REQUEST)
+            );
         }
 
         return createConnectorFuture;

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java
@@ -12,6 +12,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.ml.client.MachineLearningNodeClient;
 import org.opensearch.ml.common.transport.deploy.MLDeployModelResponse;
@@ -68,7 +69,13 @@ public void onFailure(Exception e) {
                 break;
             }
         }
-        mlClient.deploy(modelId, actionListener);
+
+        if (modelId != null) {
+            mlClient.deploy(modelId, actionListener);
+        } else {
+            deployModelFuture.completeExceptionally(new FlowFrameworkException("Model ID is not provided", RestStatus.BAD_REQUEST));
+        }
+
         return deployModelFuture;
     }
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java
@@ -12,6 +12,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.ml.client.MachineLearningNodeClient;
 import org.opensearch.ml.common.FunctionName;
@@ -137,6 +138,10 @@ public void onFailure(Exception e) {
                 .build();
 
             mlClient.register(mlInput, actionListener);
+        } else {
+            registerModelFuture.completeExceptionally(
+                new FlowFrameworkException("Required fields are not provided", RestStatus.BAD_REQUEST)
+            );
         }
 
         return registerModelFuture;