From f40890ceb80031ef94ad13d31a078e0d1e436376 Mon Sep 17 00:00:00 2001 From: Daniel Widdis Date: Thu, 18 Jul 2024 11:40:34 -0700 Subject: [PATCH] Improve Template and WorkflowState builders (#778) * Simplify Template builder instantiation Signed-off-by: Daniel Widdis * Add ability to more easily update WorkflowState Signed-off-by: Daniel Widdis * Changelog, initialize template parser Signed-off-by: Daniel Widdis --------- Signed-off-by: Daniel Widdis --- CHANGELOG.md | 1 + .../indices/FlowFrameworkIndicesHandler.java | 3 +- .../flowframework/model/Template.java | 24 +++++- .../flowframework/model/WorkflowState.java | 67 ++++++++++++++++- .../CreateWorkflowTransportAction.java | 3 +- .../transport/GetWorkflowStateResponse.java | 3 +- .../ProvisionWorkflowTransportAction.java | 2 +- .../flowframework/util/EncryptorUtils.java | 6 +- .../flowframework/model/TemplateTests.java | 24 +++++- .../model/WorkflowStateTests.java | 73 +++++++++++++++++++ .../rest/FlowFrameworkRestApiIT.java | 4 +- .../CreateWorkflowTransportActionTests.java | 7 +- 12 files changed, 197 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 63453dc37..dbc57cbe9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,3 +28,4 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) ### Documentation ### Maintenance ### Refactoring +- Improve Template and WorkflowState builders ([#778](https://github.com/opensearch-project/flow-framework/pull/778)) diff --git a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java index 8fcfd1207..18f0a9780 100644 --- a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java +++ b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java @@ -357,7 +357,8 @@ public void initializeConfigIndex(ActionListener listener) { * @param listener action listener */ public void putInitialStateToWorkflowState(String workflowId, User user, ActionListener listener) { - WorkflowState state = new WorkflowState.Builder().workflowId(workflowId) + WorkflowState state = WorkflowState.builder() + .workflowId(workflowId) .state(State.NOT_STARTED.name()) .provisioningProgress(ProvisioningProgress.NOT_STARTED.name()) .user(user) diff --git a/src/main/java/org/opensearch/flowframework/model/Template.java b/src/main/java/org/opensearch/flowframework/model/Template.java index a99b87c4b..f4d2ae600 100644 --- a/src/main/java/org/opensearch/flowframework/model/Template.java +++ b/src/main/java/org/opensearch/flowframework/model/Template.java @@ -20,7 +20,6 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.flowframework.exception.FlowFrameworkException; -import org.opensearch.flowframework.model.Template.Builder; import org.opensearch.flowframework.util.ParseUtils; import java.io.IOException; @@ -137,13 +136,13 @@ public static class Builder { /** * Empty Constructor for the Builder object */ - public Builder() {} + private Builder() {} /** * Construct a Builder from an existing template * @param t The existing template to copy */ - public Builder(Template t) { + private Builder(Template t) { this.name = t.name(); this.description = t.description(); this.useCase = t.useCase(); @@ -294,6 +293,23 @@ public Template build() { } } + /** + * Instantiate a new Template builder + * @return a new builder instance + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Instantiate a new Template builder initialized from an existing template + * @param t The existing template to use as the source + * @return a new builder instance initialized from the existing template + */ + public static Builder builder(Template t) { + return new Builder(t); + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { XContentBuilder xContentBuilder = builder.startObject(); @@ -352,7 +368,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws * @return the updated template. */ public static Template updateExistingTemplate(Template existingTemplate, Template templateWithNewFields) { - Builder builder = new Template.Builder(existingTemplate).lastUpdatedTime(Instant.now()); + Builder builder = Template.builder(existingTemplate).lastUpdatedTime(Instant.now()); if (templateWithNewFields.name() != null) { builder.name(templateWithNewFields.name()); } diff --git a/src/main/java/org/opensearch/flowframework/model/WorkflowState.java b/src/main/java/org/opensearch/flowframework/model/WorkflowState.java index 6b1593b6e..6a4b81a55 100644 --- a/src/main/java/org/opensearch/flowframework/model/WorkflowState.java +++ b/src/main/java/org/opensearch/flowframework/model/WorkflowState.java @@ -126,6 +126,15 @@ public static Builder builder() { return new Builder(); } + /** + * Constructs a builder object for workflowState from an existing state + * @param existingState a WorkflowState object to initialize the builder with + * @return Builder Object initialized with existing state + */ + public static Builder builder(WorkflowState existingState) { + return new Builder(existingState); + } + /** * Class for constructing a Builder for WorkflowState */ @@ -143,7 +152,23 @@ public static class Builder { /** * Empty Constructor for the Builder object */ - public Builder() {} + private Builder() {} + + /** + * Builder from existing state + * @param existingState a WorkflowState object to initialize the builder with + */ + private Builder(WorkflowState existingState) { + this.workflowId = existingState.getWorkflowId(); + this.error = existingState.getError(); + this.state = existingState.getState(); + this.provisioningProgress = existingState.getProvisioningProgress(); + this.provisionStartTime = existingState.getProvisionStartTime(); + this.provisionEndTime = existingState.getProvisionEndTime(); + this.user = existingState.getUser(); + this.userOutputs = existingState.userOutputs(); + this.resourcesCreated = existingState.resourcesCreated(); + } /** * Builder method for adding workflowID @@ -254,6 +279,44 @@ public WorkflowState build() { } } + /** + * Merges two workflow states by updating the fields from an existing state with the (non-null) fields of another one. + * @param existingState An existing Workflow state. + * @param stateWithNewFields A workflow state containing only fields to update. + * @return the updated workflow state. + */ + public static WorkflowState updateExistingWorkflowState(WorkflowState existingState, WorkflowState stateWithNewFields) { + Builder builder = WorkflowState.builder(existingState); + if (stateWithNewFields.getWorkflowId() != null) { + builder.workflowId(stateWithNewFields.getWorkflowId()); + } + if (stateWithNewFields.getError() != null) { + builder.error(stateWithNewFields.getError()); + } + if (stateWithNewFields.getState() != null) { + builder.state(stateWithNewFields.getState()); + } + if (stateWithNewFields.getProvisioningProgress() != null) { + builder.provisioningProgress(stateWithNewFields.getProvisioningProgress()); + } + if (stateWithNewFields.getProvisionStartTime() != null) { + builder.provisionStartTime(stateWithNewFields.getProvisionStartTime()); + } + if (stateWithNewFields.getProvisionEndTime() != null) { + builder.provisionEndTime(stateWithNewFields.getProvisionEndTime()); + } + if (stateWithNewFields.getUser() != null) { + builder.user(stateWithNewFields.getUser()); + } + if (stateWithNewFields.userOutputs() != null) { + builder.userOutputs(stateWithNewFields.userOutputs()); + } + if (stateWithNewFields.resourcesCreated() != null) { + builder.resourcesCreated(stateWithNewFields.resourcesCreated()); + } + return builder.build(); + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { XContentBuilder xContentBuilder = builder.startObject(); @@ -492,7 +555,7 @@ public Map userOutputs() { } /** - * A map of all the resources created + * A list of all the resources created * @return the resources created */ public List resourcesCreated() { diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java index 0732fc106..930feef90 100644 --- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java @@ -243,7 +243,8 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener createWorkflowValidation(client(), cyclicalTemplate)); diff --git a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java index 776ce1460..265d1d52d 100644 --- a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java @@ -396,7 +396,7 @@ public void testUpdateWorkflow() { ActionListener getListener = invocation.getArgument(1); GetResponse getResponse = mock(GetResponse.class); when(getResponse.isExists()).thenReturn(true); - when(getResponse.getSourceAsString()).thenReturn(new Template.Builder().name("test").build().toJson()); + when(getResponse.getSourceAsString()).thenReturn(Template.builder().name("test").build().toJson()); getListener.onResponse(getResponse); return null; }).when(client).get(any(GetRequest.class), any()); @@ -425,7 +425,7 @@ public void testUpdateWorkflowWithField() { ActionListener listener = mock(ActionListener.class); WorkflowRequest updateWorkflow = new WorkflowRequest( "1", - new Template.Builder().name("new name").description("test").useCase(null).uiMetadata(Map.of("foo", "bar")).build(), + Template.builder().name("new name").description("test").useCase(null).uiMetadata(Map.of("foo", "bar")).build(), Map.of(UPDATE_WORKFLOW_FIELDS, "true") ); @@ -463,7 +463,8 @@ public void testUpdateWorkflowWithField() { updateWorkflow = new WorkflowRequest( "1", - new Template.Builder().useCase("foo") + Template.builder() + .useCase("foo") .templateVersion(Version.CURRENT) .compatibilityVersion(List.of(Version.V_2_0_0, Version.CURRENT)) .build(),