From 4497b175bd4643f7b6d0deb1af092500661c9934 Mon Sep 17 00:00:00 2001 From: Owais Kazi Date: Wed, 8 Nov 2023 14:31:50 -0800 Subject: [PATCH] Addressed PR comments Signed-off-by: Owais Kazi --- .../rest/RestProvisionWorkflowAction.java | 2 +- .../transport/CreateWorkflowTransportAction.java | 4 ++-- .../flowframework/transport/WorkflowRequest.java | 13 +++++++++++-- .../ProvisionWorkflowTransportActionTests.java | 4 ++-- .../transport/WorkflowRequestResponseTests.java | 6 +++--- 5 files changed, 19 insertions(+), 10 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java index 69d853423..9e6eb4d01 100644 --- 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 index 7d3719f7f..f6c5fbdf6 100644 --- 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 { + 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 internalListener) { + protected void checkMaxWorkflows(TimeValue requestTimeOut, Integer maxWorkflow, ActionListener 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 index 92e83282c..d049be8f6 100644 --- 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/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java index 9e22bff01..d3f6fb6fd 100644 --- a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java @@ -102,7 +102,7 @@ public void testProvisionWorkflow() { String workflowId = "1"; @SuppressWarnings("unchecked") ActionListener listener = mock(ActionListener.class); - WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null, null, null); + WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null); doAnswer(invocation -> { ActionListener responseListener = invocation.getArgument(1); @@ -124,7 +124,7 @@ public void testProvisionWorkflow() { public void testFailedToRetrieveTemplateFromGlobalContext() { @SuppressWarnings("unchecked") ActionListener listener = mock(ActionListener.class); - WorkflowRequest request = new WorkflowRequest("1", null, null, null); + WorkflowRequest request = new WorkflowRequest("1", null); doAnswer(invocation -> { ActionListener 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 index 5de21d5ab..d64249e27 100644 --- 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());