Addressed PR comments

Signed-off-by: Owais Kazi <[email protected]>
opensearch-project · Nov 8, 2023 · 4497b17 · 4497b17
1 parent 7a0b6d3
commit 4497b17
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 10 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java
@@ -84,7 +84,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, null, null);
+            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));

diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
@@ -111,7 +111,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
 
         if (request.getWorkflowId() == null) {
             // Throttle incoming requests
-            onSearchGlobalContext(request.getRequestTimeout(), request.getMaxWorkflows(), ActionListener.wrap(max -> {
+            checkMaxWorkflows(request.getRequestTimeout(), request.getMaxWorkflows(), ActionListener.wrap(max -> {
                 if (!max) {
                     String errorMessage = "Maximum workflows limit reached " + request.getMaxWorkflows();
                     logger.error(errorMessage);
@@ -195,7 +195,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
      *  @param maxWorkflow max workflows
      *  @param internalListener listener for search request
      */
-    protected void onSearchGlobalContext(TimeValue requestTimeOut, Integer maxWorkflow, ActionListener<Boolean> internalListener) {
+    protected void checkMaxWorkflows(TimeValue requestTimeOut, Integer maxWorkflow, ActionListener<Boolean> internalListener) {
         QueryBuilder query = QueryBuilders.matchAllQuery();
         SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder().query(query).size(0).timeout(requestTimeOut);
 

diff --git a/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java b/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java
@@ -48,6 +48,15 @@ public class WorkflowRequest extends ActionRequest {
      */
     private Integer maxWorkflows;
 
+    /**
+     * Instantiates a new WorkflowRequest, defaults dry run to false and set requestTimeout and maxWorkflows to null
+     * @param workflowId the documentId of the workflow
+     * @param template the use case template which describes the workflow
+     */
+    public WorkflowRequest(@Nullable String workflowId, @Nullable Template template) {
+        this(workflowId, template, false, null, null);
+    }
+
     /**
      * Instantiates a new WorkflowRequest and defaults dry run to false
      * @param workflowId the documentId of the workflow
@@ -76,8 +85,8 @@ public WorkflowRequest(
         @Nullable String workflowId,
         @Nullable Template template,
         boolean dryRun,
-        TimeValue requestTimeout,
-        Integer maxWorkflows
+        @Nullable TimeValue requestTimeout,
+        @Nullable Integer maxWorkflows
     ) {
         this.workflowId = workflowId;
         this.template = template;

diff --git a/...st/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java b/...st/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java
@@ -102,7 +102,7 @@ public void testProvisionWorkflow() {
         String workflowId = "1";
         @SuppressWarnings("unchecked")
         ActionListener<WorkflowResponse> listener = mock(ActionListener.class);
-        WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null, null, null);
+        WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null);
 
         doAnswer(invocation -> {
             ActionListener<GetResponse> responseListener = invocation.getArgument(1);
@@ -124,7 +124,7 @@ public void testProvisionWorkflow() {
     public void testFailedToRetrieveTemplateFromGlobalContext() {
         @SuppressWarnings("unchecked")
         ActionListener<WorkflowResponse> listener = mock(ActionListener.class);
-        WorkflowRequest request = new WorkflowRequest("1", null, null, null);
+        WorkflowRequest request = new WorkflowRequest("1", null);
         doAnswer(invocation -> {
             ActionListener<GetResponse> responseListener = invocation.getArgument(1);
             responseListener.onFailure(new Exception("Failed to retrieve template from global context."));

diff --git a/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java b/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java
@@ -55,7 +55,7 @@ public void setUp() throws Exception {
     }
 
     public void testNullIdWorkflowRequest() throws IOException {
-        WorkflowRequest nullIdRequest = new WorkflowRequest(null, template, null, null);
+        WorkflowRequest nullIdRequest = new WorkflowRequest(null, template);
         assertNull(nullIdRequest.getWorkflowId());
         assertEquals(template, nullIdRequest.getTemplate());
         assertNull(nullIdRequest.validate());
@@ -71,7 +71,7 @@ public void testNullIdWorkflowRequest() throws IOException {
     }
 
     public void testNullTemplateWorkflowRequest() throws IOException {
-        WorkflowRequest nullTemplateRequest = new WorkflowRequest("123", null, null, null);
+        WorkflowRequest nullTemplateRequest = new WorkflowRequest("123", null);
         assertNotNull(nullTemplateRequest.getWorkflowId());
         assertNull(nullTemplateRequest.getTemplate());
         assertNull(nullTemplateRequest.validate());
@@ -87,7 +87,7 @@ public void testNullTemplateWorkflowRequest() throws IOException {
     }
 
     public void testWorkflowRequest() throws IOException {
-        WorkflowRequest workflowRequest = new WorkflowRequest("123", template, null, null);
+        WorkflowRequest workflowRequest = new WorkflowRequest("123", template);
         assertNotNull(workflowRequest.getWorkflowId());
         assertEquals(template, workflowRequest.getTemplate());
         assertNull(workflowRequest.validate());