From a0edf682838f2d14b2cf6ca4a0d0fcda0769ccf2 Mon Sep 17 00:00:00 2001 From: Owais Kazi Date: Wed, 8 Nov 2023 16:57:29 -0800 Subject: [PATCH] Moved UI_METADATA from state to GC index and excluded it from search response Signed-off-by: Owais Kazi --- .../flowframework/model/Template.java | 167 +++++++++++++++++- .../flowframework/model/WorkflowState.java | 50 ++---- .../CreateWorkflowTransportAction.java | 1 + .../flowframework/util/RestHandlerUtils.java | 5 +- .../flowframework/model/TemplateTests.java | 4 + .../rest/RestCreateWorkflowActionTests.java | 1 + .../CreateWorkflowTransportActionTests.java | 2 + ...ProvisionWorkflowTransportActionTests.java | 1 + .../WorkflowRequestResponseTests.java | 1 + 9 files changed, 188 insertions(+), 44 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/model/Template.java b/src/main/java/org/opensearch/flowframework/model/Template.java index fbafb8fd0..58178c2f4 100644 --- a/src/main/java/org/opensearch/flowframework/model/Template.java +++ b/src/main/java/org/opensearch/flowframework/model/Template.java @@ -30,6 +30,7 @@ import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD; import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD; import static org.opensearch.flowframework.common.CommonValue.TEMPLATE_FIELD; +import static org.opensearch.flowframework.common.CommonValue.UI_METADATA_FIELD; import static org.opensearch.flowframework.common.CommonValue.USER_FIELD; import static org.opensearch.flowframework.common.CommonValue.USE_CASE_FIELD; import static org.opensearch.flowframework.common.CommonValue.VERSION_FIELD; @@ -40,13 +41,14 @@ */ public class Template implements ToXContentObject { - private final String name; - private final String description; - private final String useCase; // probably an ENUM actually - private final Version templateVersion; - private final List compatibilityVersion; - private final Map workflows; - private final User user; + private String name; + private String description; + private String useCase; // probably an ENUM actually + private Version templateVersion; + private List compatibilityVersion; + private Map workflows; + private Map uiMetadata; + private User user; /** * Instantiate the object representing a use case template @@ -57,6 +59,7 @@ public class Template implements ToXContentObject { * @param templateVersion The version of this template * @param compatibilityVersion OpenSearch version compatibility of this template * @param workflows Workflow graph definitions corresponding to the defined operations. + * @param uiMetadata The UI metadata related to the given workflow * @param user The user extracted from the thread context from the request */ public Template( @@ -66,6 +69,7 @@ public Template( Version templateVersion, List compatibilityVersion, Map workflows, + Map uiMetadata, User user ) { this.name = name; @@ -74,9 +78,129 @@ public Template( this.templateVersion = templateVersion; this.compatibilityVersion = List.copyOf(compatibilityVersion); this.workflows = Map.copyOf(workflows); + this.uiMetadata = uiMetadata; this.user = user; } + private Template() {} + + /** + * Class for constructing a Builder for Template + */ + public static class Builder { + private String name = null; + private String description = ""; + private String useCase = ""; + private Version templateVersion = null; + private List compatibilityVersion = new ArrayList<>(); + private Map workflows = new HashMap<>(); + private Map uiMetadata = null; + private User user = null; + + /** + * Empty Constructor for the Builder object + */ + public Builder() {} + + /** + * Builder method for adding template name + * @param name template name + * @return the Builder object + */ + public Builder name(String name) { + this.name = name; + return this; + } + + /** + * Builder method for adding template description + * @param description template description + * @return the Builder object + */ + public Builder description(String description) { + this.description = description; + return this; + } + + /** + * Builder method for adding template useCase + * @param useCase template useCase + * @return the Builder object + */ + public Builder useCase(String useCase) { + this.useCase = useCase; + return this; + } + + /** + * Builder method for adding templateVersion + * @param templateVersion templateVersion + * @return the Builder object + */ + public Builder templateVersion(Version templateVersion) { + this.templateVersion = templateVersion; + return this; + } + + /** + * Builder method for adding compatibilityVersion + * @param compatibilityVersion compatibilityVersion + * @return the Builder object + */ + public Builder compatibilityVersion(List compatibilityVersion) { + this.compatibilityVersion = compatibilityVersion; + return this; + } + + /** + * Builder method for adding workflows + * @param workflows workflows + * @return the Builder object + */ + public Builder workflows(Map workflows) { + this.workflows = workflows; + return this; + } + + /** + * Builder method for adding uiMetadata + * @param uiMetadata uiMetadata + * @return the Builder object + */ + public Builder uiMetadata(Map uiMetadata) { + this.uiMetadata = uiMetadata; + return this; + } + + /** + * Builder method for adding user + * @param user user + * @return the Builder object + */ + public Builder user(User user) { + this.user = user; + return this; + } + + /** + * Allows building a template + * @return Template Object containing all needed fields + */ + public Template build() { + Template template = new Template(); + template.name = this.name; + template.description = this.description; + template.useCase = this.useCase; + template.templateVersion = this.templateVersion; + template.compatibilityVersion = this.compatibilityVersion; + template.workflows = this.workflows; + template.uiMetadata = this.uiMetadata; + template.user = this.user; + return template; + } + + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { XContentBuilder xContentBuilder = builder.startObject(); @@ -103,6 +227,11 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws for (Entry e : workflows.entrySet()) { xContentBuilder.field(e.getKey(), e.getValue(), params); } + + if (uiMetadata != null && !uiMetadata.isEmpty()) { + xContentBuilder.field(UI_METADATA_FIELD, uiMetadata); + } + xContentBuilder.endObject(); if (user != null) { xContentBuilder.field(USER_FIELD, user); @@ -125,6 +254,7 @@ public static Template parse(XContentParser parser) throws IOException { Version templateVersion = null; List compatibilityVersion = new ArrayList<>(); Map workflows = new HashMap<>(); + Map uiMetadata = null; User user = null; ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); @@ -169,6 +299,9 @@ public static Template parse(XContentParser parser) throws IOException { workflows.put(workflowFieldName, Workflow.parse(parser)); } break; + case UI_METADATA_FIELD: + uiMetadata = parser.map(); + break; case USER_FIELD: user = User.parse(parser); break; @@ -180,7 +313,15 @@ public static Template parse(XContentParser parser) throws IOException { throw new IOException("A template object requires a name."); } - return new Template(name, description, useCase, templateVersion, compatibilityVersion, workflows, user); + return new Builder().name(name) + .description(description) + .useCase(useCase) + .templateVersion(templateVersion) + .compatibilityVersion(compatibilityVersion) + .workflows(workflows) + .uiMetadata(uiMetadata) + .user(user) + .build(); } /** @@ -276,6 +417,14 @@ public Map workflows() { return workflows; } + /** + * A map corresponding to the UI metadata + * @return the userOutputs + */ + public Map getUiMetadata() { + return uiMetadata; + } + /** * User that created and owns this template * @return the user @@ -298,6 +447,8 @@ public String toString() { + compatibilityVersion + ", workflows=" + workflows + + ", uiMedata=" + + uiMetadata + "]"; } } diff --git a/src/main/java/org/opensearch/flowframework/model/WorkflowState.java b/src/main/java/org/opensearch/flowframework/model/WorkflowState.java index c2b39f0ec..f51520585 100644 --- a/src/main/java/org/opensearch/flowframework/model/WorkflowState.java +++ b/src/main/java/org/opensearch/flowframework/model/WorkflowState.java @@ -26,7 +26,6 @@ import static org.opensearch.flowframework.common.CommonValue.PROVISION_START_TIME_FIELD; import static org.opensearch.flowframework.common.CommonValue.RESOURCES_CREATED_FIELD; import static org.opensearch.flowframework.common.CommonValue.STATE_FIELD; -import static org.opensearch.flowframework.common.CommonValue.UI_METADATA_FIELD; import static org.opensearch.flowframework.common.CommonValue.USER_FIELD; import static org.opensearch.flowframework.common.CommonValue.USER_OUTPUTS_FIELD; import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID_FIELD; @@ -45,7 +44,6 @@ public class WorkflowState implements ToXContentObject { private Instant provisionStartTime; private Instant provisionEndTime; private User user; - private Map uiMetadata; private Map userOutputs; private Map resourcesCreated; @@ -59,7 +57,6 @@ public class WorkflowState implements ToXContentObject { * @param provisionStartTime Indicates the start time of the whole provisioning flow * @param provisionEndTime Indicates the end time of the whole provisioning flow * @param user The user extracted from the thread context from the request - * @param uiMetadata The UI metadata related to the given workflow * @param userOutputs A map of essential API responses for backend to use and lookup. * @param resourcesCreated A map of all the resources created. */ @@ -71,7 +68,6 @@ public WorkflowState( Instant provisionStartTime, Instant provisionEndTime, User user, - Map uiMetadata, Map userOutputs, Map resourcesCreated ) { @@ -82,7 +78,6 @@ public WorkflowState( this.provisionStartTime = provisionStartTime; this.provisionEndTime = provisionEndTime; this.user = user; - this.uiMetadata = uiMetadata; this.userOutputs = Map.copyOf(userOutputs); this.resourcesCreated = Map.copyOf(resourcesCreated); } @@ -108,7 +103,6 @@ public static class Builder { private Instant provisionStartTime = null; private Instant provisionEndTime = null; private User user = null; - private Map uiMetadata = null; private Map userOutputs = null; private Map resourcesCreated = null; @@ -187,16 +181,6 @@ public Builder user(User user) { return this; } - /** - * Builder method for adding uiMetadata - * @param uiMetadata uiMetadata - * @return the Builder object - */ - public Builder uiMetadata(Map uiMetadata) { - this.uiMetadata = uiMetadata; - return this; - } - /** * Builder method for adding userOutputs * @param userOutputs userOutputs @@ -230,7 +214,6 @@ public WorkflowState build() { workflowState.provisionStartTime = this.provisionStartTime; workflowState.provisionEndTime = this.provisionEndTime; workflowState.user = this.user; - workflowState.uiMetadata = this.uiMetadata; workflowState.userOutputs = this.userOutputs; workflowState.resourcesCreated = this.resourcesCreated; return workflowState; @@ -261,9 +244,6 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws if (user != null) { xContentBuilder.field(USER_FIELD, user); } - if (uiMetadata != null && !uiMetadata.isEmpty()) { - xContentBuilder.field(UI_METADATA_FIELD, uiMetadata); - } if (userOutputs != null && !userOutputs.isEmpty()) { xContentBuilder.field(USER_OUTPUTS_FIELD, userOutputs); } @@ -288,7 +268,6 @@ public static WorkflowState parse(XContentParser parser) throws IOException { Instant provisionStartTime = null; Instant provisionEndTime = null; User user = null; - Map uiMetadata = null; Map userOutputs = new HashMap<>(); Map resourcesCreated = new HashMap<>(); @@ -318,9 +297,6 @@ public static WorkflowState parse(XContentParser parser) throws IOException { case USER_FIELD: user = User.parse(parser); break; - case UI_METADATA_FIELD: - uiMetadata = parser.map(); - break; case USER_OUTPUTS_FIELD: ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); while (parser.nextToken() != XContentParser.Token.END_OBJECT) { @@ -367,7 +343,6 @@ public static WorkflowState parse(XContentParser parser) throws IOException { .provisionStartTime(provisionStartTime) .provisionEndTime(provisionEndTime) .user(user) - .uiMetadata(uiMetadata) .userOutputs(userOutputs) .resourcesCreated(resourcesCreated) .build(); @@ -429,14 +404,6 @@ public User getUser() { return user; } - /** - * A map corresponding to the UI metadata - * @return the userOutputs - */ - public Map getUiMetadata() { - return uiMetadata; - } - /** * A map of essential API responses * @return the userOutputs @@ -452,4 +419,21 @@ public Map userOutputs() { public Map resourcesCreated() { return resourcesCreated; } + + @Override + public String toString() { + return "WorkflowState [workflowId=" + + workflowId + + ", error=" + + error + + ", state=" + + state + + ", provisioningProgress=" + + provisioningProgress + + ", userOutputs=" + + userOutputs + + ", resourcesCreated=" + + resourcesCreated + + "]"; + } } diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java index 5781bb412..a782740c4 100644 --- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java @@ -91,6 +91,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener nodes = List.of(nodeA, nodeB); List edges = List.of(edgeAB); Workflow workflow = new Workflow(Map.of("key", "value"), nodes, edges); + Map uiMetadata = null; Template template = new Template( "test", @@ -43,6 +44,7 @@ public void testTemplate() throws IOException { templateVersion, compatibilityVersion, Map.of("workflow", workflow), + uiMetadata, null ); @@ -51,6 +53,7 @@ public void testTemplate() throws IOException { assertEquals("test use case", template.useCase()); assertEquals(templateVersion, template.templateVersion()); assertEquals(compatibilityVersion, template.compatibilityVersion()); + assertEquals(uiMetadata, template.getUiMetadata()); Workflow wf = template.workflows().get("workflow"); assertNotNull(wf); assertEquals("Workflow [userParams={key=value}, nodes=[A, B], edges=[A->B]]", wf.toString()); @@ -63,6 +66,7 @@ public void testTemplate() throws IOException { assertEquals("test use case", templateX.useCase()); assertEquals(templateVersion, templateX.templateVersion()); assertEquals(compatibilityVersion, templateX.compatibilityVersion()); + assertEquals(uiMetadata, templateX.getUiMetadata()); Workflow wfX = templateX.workflows().get("workflow"); assertNotNull(wfX); assertEquals("Workflow [userParams={key=value}, nodes=[A, B], edges=[A->B]]", wfX.toString()); diff --git a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java index c81498cea..1fbd78641 100644 --- a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java +++ b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java @@ -81,6 +81,7 @@ public void setUp() throws Exception { templateVersion, compatibilityVersions, Map.of("workflow", workflow), + Map.of(), TestHelpers.randomUser() ); diff --git a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java index 18f330c98..5492ad822 100644 --- a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java @@ -108,6 +108,7 @@ public void setUp() throws Exception { templateVersion, compatibilityVersions, Map.of("workflow", workflow), + Map.of(), TestHelpers.randomUser() ); } @@ -160,6 +161,7 @@ public void testFailedDryRunValidation() { Version.fromString("1.0.0"), List.of(Version.fromString("2.0.0"), Version.fromString("3.0.0")), Map.of("workflow", workflow), + Map.of(), TestHelpers.randomUser() ); diff --git a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java index d3f6fb6fd..7d061e572 100644 --- a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java @@ -87,6 +87,7 @@ public void setUp() throws Exception { templateVersion, compatibilityVersions, Map.of("provision", workflow), + Map.of(), TestHelpers.randomUser() ); diff --git a/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java b/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java index d64249e27..e79ed38cc 100644 --- a/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java @@ -50,6 +50,7 @@ public void setUp() throws Exception { templateVersion, compatibilityVersions, Map.of("workflow", workflow), + Map.of(), TestHelpers.randomUser() ); }