diff --git a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java index 1355b2cff..14d94fd8c 100644 --- a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java +++ b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java @@ -81,7 +81,7 @@ public void putTemplateToGlobalContext(Template template, ActionListener context.restore())); } catch (Exception e) { diff --git a/src/main/java/org/opensearch/flowframework/model/Template.java b/src/main/java/org/opensearch/flowframework/model/Template.java index b3f4478b9..bcfbad22d 100644 --- a/src/main/java/org/opensearch/flowframework/model/Template.java +++ b/src/main/java/org/opensearch/flowframework/model/Template.java @@ -159,6 +159,228 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws return xContentBuilder.endObject(); } + /** + * Converts a template object into a Global Context document + * @param builder the XContentBuilder + * @param params the params + * @return the XContentBuilder + * @throws IOException + */ + public XContentBuilder toDocumentSource(XContentBuilder builder, Params params) throws IOException { + XContentBuilder xContentBuilder = builder.startObject(); + xContentBuilder.field(NAME_FIELD, this.name); + xContentBuilder.field(DESCRIPTION_FIELD, this.description); + xContentBuilder.field(USE_CASE_FIELD, this.useCase); + xContentBuilder.startArray(OPERATIONS_FIELD); + for (String op : this.operations) { + xContentBuilder.value(op); + } + xContentBuilder.endArray(); + + if (this.templateVersion != null || !this.compatibilityVersion.isEmpty()) { + xContentBuilder.startObject(VERSION_FIELD); + if (this.templateVersion != null) { + xContentBuilder.field(TEMPLATE_FIELD, this.templateVersion); + } + if (!this.compatibilityVersion.isEmpty()) { + xContentBuilder.startArray(COMPATIBILITY_FIELD); + for (Version v : this.compatibilityVersion) { + xContentBuilder.value(v); + } + xContentBuilder.endArray(); + } + xContentBuilder.endObject(); + } + + if (!this.userInputs.isEmpty()) { + xContentBuilder.startObject(USER_INPUTS_FIELD); + for (Entry e : userInputs.entrySet()) { + xContentBuilder.field(e.getKey(), e.getValue()); + } + xContentBuilder.endObject(); + } + + try (XContentBuilder workflowBuilder = JsonXContent.contentBuilder()) { + workflowBuilder.startObject(); + for (Entry e : workflows.entrySet()) { + workflowBuilder.field(e.getKey(), e.getValue()); + } + workflowBuilder.endObject(); + xContentBuilder.field(WORKFLOWS_FIELD, workflowBuilder.toString()); + } + + try (XContentBuilder userOutputsBuilder = JsonXContent.contentBuilder()) { + userOutputsBuilder.startObject(); + for (Entry e : userOutputs.entrySet()) { + userOutputsBuilder.field(e.getKey(), e.getValue()); + } + userOutputsBuilder.endObject(); + xContentBuilder.field(USER_OUTPUTS_FIELD, userOutputsBuilder.toString()); + } + + try (XContentBuilder resourcesCreatedBuilder = JsonXContent.contentBuilder()) { + resourcesCreatedBuilder.startObject(); + for (Entry e : resourcesCreated.entrySet()) { + resourcesCreatedBuilder.field(e.getKey(), e.getValue()); + } + resourcesCreatedBuilder.endObject(); + xContentBuilder.field(RESOURCES_CREATED_FIELD, resourcesCreatedBuilder.toString()); + } + + xContentBuilder.endObject(); + + return xContentBuilder; + + } + + /** + * Parse global context document source into a Template instance + * + * @param documentSource the document source string + * @return an instance of the template + * @throws IOException if content can't be parsed correctly + */ + public static Template parseFromDocumentSource(String documentSource) throws IOException { + XContentParser parser = jsonToParser(documentSource); + + String name = null; + String description = ""; + String useCase = ""; + List operations = new ArrayList<>(); + Version templateVersion = null; + List compatibilityVersion = new ArrayList<>(); + Map userInputs = new HashMap<>(); + Map workflows = new HashMap<>(); + Map userOutputs = new HashMap<>(); + Map resourcesCreated = new HashMap<>(); + + ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); + while (parser.nextToken() != XContentParser.Token.END_OBJECT) { + String fieldName = parser.currentName(); + parser.nextToken(); + switch (fieldName) { + case NAME_FIELD: + name = parser.text(); + break; + case DESCRIPTION_FIELD: + description = parser.text(); + break; + case USE_CASE_FIELD: + useCase = parser.text(); + break; + case OPERATIONS_FIELD: + ensureExpectedToken(XContentParser.Token.START_ARRAY, parser.currentToken(), parser); + while (parser.nextToken() != XContentParser.Token.END_ARRAY) { + operations.add(parser.text()); + } + break; + case VERSION_FIELD: + ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); + while (parser.nextToken() != XContentParser.Token.END_OBJECT) { + String versionFieldName = parser.currentName(); + parser.nextToken(); + switch (versionFieldName) { + case TEMPLATE_FIELD: + templateVersion = Version.fromString(parser.text()); + break; + case COMPATIBILITY_FIELD: + ensureExpectedToken(XContentParser.Token.START_ARRAY, parser.currentToken(), parser); + while (parser.nextToken() != XContentParser.Token.END_ARRAY) { + compatibilityVersion.add(Version.fromString(parser.text())); + } + break; + default: + throw new IOException("Unable to parse field [" + fieldName + "] in a version object."); + } + } + break; + case USER_INPUTS_FIELD: + ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); + while (parser.nextToken() != XContentParser.Token.END_OBJECT) { + String inputFieldName = parser.currentName(); + switch (parser.nextToken()) { + case VALUE_STRING: + userInputs.put(inputFieldName, parser.text()); + break; + case START_OBJECT: + userInputs.put(inputFieldName, parseStringToStringMap(parser)); + break; + default: + throw new IOException("Unable to parse field [" + inputFieldName + "] in a user inputs object."); + } + } + break; + case WORKFLOWS_FIELD: + String workflowsJson = parser.text(); + XContentParser workflowsParser = jsonToParser(workflowsJson); + while (workflowsParser.nextToken() != XContentParser.Token.END_OBJECT) { + String workflowFieldName = workflowsParser.currentName(); + workflowsParser.nextToken(); + workflows.put(workflowFieldName, Workflow.parse(workflowsParser)); + } + break; + case USER_OUTPUTS_FIELD: + + String userOutputsJson = parser.text(); + XContentParser userOuputsParser = jsonToParser(userOutputsJson); + while (userOuputsParser.nextToken() != XContentParser.Token.END_OBJECT) { + String userOutputsFieldName = userOuputsParser.currentName(); + switch (userOuputsParser.nextToken()) { + case VALUE_STRING: + userOutputs.put(userOutputsFieldName, userOuputsParser.text()); + break; + case START_OBJECT: + userOutputs.put(userOutputsFieldName, parseStringToStringMap(userOuputsParser)); + break; + default: + throw new IOException("Unable to parse field [" + userOutputsFieldName + "] in a user_outputs object."); + } + } + break; + + case RESOURCES_CREATED_FIELD: + + String resourcesCreatedJson = parser.text(); + XContentParser resourcesCreatedParser = jsonToParser(resourcesCreatedJson); + while (resourcesCreatedParser.nextToken() != XContentParser.Token.END_OBJECT) { + String resourcesCreatedField = resourcesCreatedParser.currentName(); + switch (resourcesCreatedParser.nextToken()) { + case VALUE_STRING: + resourcesCreated.put(resourcesCreatedField, resourcesCreatedParser.text()); + break; + case START_OBJECT: + resourcesCreated.put(resourcesCreatedField, parseStringToStringMap(resourcesCreatedParser)); + break; + default: + throw new IOException( + "Unable to parse field [" + resourcesCreatedField + "] in a resources_created object." + ); + } + } + break; + + default: + throw new IOException("Unable to parse field [" + fieldName + "] in a template object."); + } + } + if (name == null) { + throw new IOException("An template object requires a name."); + } + + return new Template( + name, + description, + useCase, + operations, + templateVersion, + compatibilityVersion, + userInputs, + workflows, + userOutputs, + resourcesCreated + ); + } + /** * Parse raw json content into a Template instance. * @@ -300,6 +522,23 @@ public static Template parse(XContentParser parser) throws IOException { ); } + /** + * Converts a JSON string into an XContentParser + * + * @param json the json string + * @return The XContent parser for the json string + * @throws IOException on failure to create the parser + */ + public static XContentParser jsonToParser(String json) throws IOException { + XContentParser parser = JsonXContent.jsonXContent.createParser( + NamedXContentRegistry.EMPTY, + LoggingDeprecationHandler.INSTANCE, + json + ); + ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser); + return parser; + } + /** * Parse a JSON use case template * diff --git a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java index f81a623b7..be4f2ed2d 100644 --- a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java @@ -83,7 +83,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener { + when(template.toDocumentSource(any(XContentBuilder.class), eq(ToXContent.EMPTY_PARAMS))).thenAnswer(invocation -> { XContentBuilder builder = invocation.getArgument(0); return builder; }); diff --git a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java index d190d3dc0..c58e810b9 100644 --- a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java @@ -106,7 +106,7 @@ public void testProvisionWorkflow() { ActionListener responseListener = invocation.getArgument(1); XContentBuilder builder = XContentFactory.jsonBuilder(); - this.template.toXContent(builder, ToXContent.EMPTY_PARAMS); + this.template.toDocumentSource(builder, ToXContent.EMPTY_PARAMS); BytesReference templateBytesRef = BytesReference.bytes(builder); GetResult getResult = new GetResult(GLOBAL_CONTEXT_INDEX, workflowId, 1, 1, 1, true, templateBytesRef, null, null); responseListener.onResponse(new GetResponse(getResult));