From d3be2e4ad485382790ae6b7b927352cf5a4645eb Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Fri, 29 Sep 2023 22:41:44 +0000 Subject: [PATCH 01/23] Inital implementation, set up rest/transport actions, registration, pending global context index, state index implementations Signed-off-by: Joshua Palis --- .../flowframework/FlowFrameworkPlugin.java | 50 +++++++++- .../rest/RestCreateWorkflowAction.java | 69 +++++++++++++ .../rest/RestProvisionWorkflowAction.java | 83 ++++++++++++++++ .../transport/CreateWorkflowAction.java | 27 ++++++ .../CreateWorkflowTransportAction.java | 65 +++++++++++++ .../transport/ProvisionWorkflowAction.java | 27 ++++++ .../ProvisionWorkflowTransportAction.java | 69 +++++++++++++ .../transport/WorkflowRequest.java | 96 +++++++++++++++++++ .../transport/WorkflowResponse.java | 67 +++++++++++++ 9 files changed, 552 insertions(+), 1 deletion(-) create mode 100644 src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java create mode 100644 src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java create mode 100644 src/main/java/org/opensearch/flowframework/transport/CreateWorkflowAction.java create mode 100644 src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java create mode 100644 src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowAction.java create mode 100644 src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java create mode 100644 src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java create mode 100644 src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java diff --git a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java index d701c832e..e930912ce 100644 --- a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java +++ b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java @@ -9,28 +9,54 @@ package org.opensearch.flowframework; import com.google.common.collect.ImmutableList; +import org.opensearch.action.ActionRequest; import org.opensearch.client.Client; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.IndexScopedSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.settings.SettingsFilter; +import org.opensearch.core.action.ActionResponse; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.env.Environment; import org.opensearch.env.NodeEnvironment; +import org.opensearch.flowframework.rest.RestCreateWorkflowAction; +import org.opensearch.flowframework.rest.RestProvisionWorkflowAction; +import org.opensearch.flowframework.transport.CreateWorkflowAction; +import org.opensearch.flowframework.transport.CreateWorkflowTransportAction; +import org.opensearch.flowframework.transport.ProvisionWorkflowAction; +import org.opensearch.flowframework.transport.ProvisionWorkflowTransportAction; import org.opensearch.flowframework.workflow.WorkflowProcessSorter; import org.opensearch.flowframework.workflow.WorkflowStepFactory; +import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.Plugin; import org.opensearch.repositories.RepositoriesService; +import org.opensearch.rest.RestController; +import org.opensearch.rest.RestHandler; import org.opensearch.script.ScriptService; import org.opensearch.threadpool.ThreadPool; import org.opensearch.watcher.ResourceWatcherService; import java.util.Collection; +import java.util.List; import java.util.function.Supplier; /** * An OpenSearch plugin that enables builders to innovate AI apps on OpenSearch. */ -public class FlowFrameworkPlugin extends Plugin { +public class FlowFrameworkPlugin extends Plugin implements ActionPlugin { + + /** + * The base URI for this plugin's rest actions + */ + public static final String AI_FLOW_FRAMEWORK_BASE_URI = "/_plugins/_ai_flow"; + /** + * The URI for this plugin's workflow rest actions + */ + public static final String WORKFLOWS_URI = AI_FLOW_FRAMEWORK_BASE_URI + "/workflows"; @Override public Collection createComponents( @@ -51,4 +77,26 @@ public Collection createComponents( return ImmutableList.of(workflowStepFactory, workflowProcessSorter); } + + @Override + public List getRestHandlers( + Settings settings, + RestController restController, + ClusterSettings clusterSettings, + IndexScopedSettings indexScopedSettings, + SettingsFilter settingsFilter, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier nodesInCluster + ) { + return ImmutableList.of(new RestCreateWorkflowAction(), new RestProvisionWorkflowAction()); + } + + @Override + public List> getActions() { + return ImmutableList.of( + new ActionHandler<>(CreateWorkflowAction.INSTANCE, CreateWorkflowTransportAction.class), + new ActionHandler<>(ProvisionWorkflowAction.INSTANCE, ProvisionWorkflowTransportAction.class) + ); + } + } diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java new file mode 100644 index 000000000..5a2b352fe --- /dev/null +++ b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java @@ -0,0 +1,69 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.rest; + +import com.google.common.collect.ImmutableList; +import org.opensearch.client.node.NodeClient; +import org.opensearch.flowframework.FlowFrameworkPlugin; +import org.opensearch.flowframework.model.Template; +import org.opensearch.flowframework.transport.CreateWorkflowAction; +import org.opensearch.flowframework.transport.WorkflowRequest; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestToXContentListener; + +import java.io.IOException; +import java.util.List; +import java.util.Locale; + +/** + * Rest Action to facilitate requests to create and update a use case template + */ +public class RestCreateWorkflowAction extends BaseRestHandler { + + private static final String CREATE_WORKFLOW_ACTION = "create_workflow_action"; + + // TODO : move to common values class, pending implementation + /** + * Field name for workflow Id, the document Id of the indexed use case template + */ + public static final String WORKFLOW_ID = "workflowId"; + + /** + * Intantiates a new RestCreateWorkflowAction + */ + public RestCreateWorkflowAction() { + // TODO : Pass settings and cluster service to constructor and add settings update consumer for request timeout value + } + + @Override + public String getName() { + return CREATE_WORKFLOW_ACTION; + } + + @Override + public List routes() { + return ImmutableList.of( + // Create new workflow + new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s", FlowFrameworkPlugin.WORKFLOWS_URI)), + // Update workflow + new Route(RestRequest.Method.PUT, String.format(Locale.ROOT, "%s/{%s}", FlowFrameworkPlugin.WORKFLOWS_URI, WORKFLOW_ID)) + ); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + + String workflowId = request.param(WORKFLOW_ID, null); + Template template = Template.parse(request.content().utf8ToString()); + WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template); + return channel -> client.execute(CreateWorkflowAction.INSTANCE, workflowRequest, new RestToXContentListener<>(channel)); + } + +} diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java new file mode 100644 index 000000000..b1d8fb671 --- /dev/null +++ b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java @@ -0,0 +1,83 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.rest; + +import com.google.common.collect.ImmutableList; +import org.opensearch.client.node.NodeClient; +import org.opensearch.flowframework.FlowFrameworkPlugin; +import org.opensearch.flowframework.model.Template; +import org.opensearch.flowframework.transport.ProvisionWorkflowAction; +import org.opensearch.flowframework.transport.WorkflowRequest; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestToXContentListener; + +import java.io.IOException; +import java.util.List; +import java.util.Locale; + +/** + * Rest action to facilitate requests to provision a workflow from an inline defined or stored use case template + */ +public class RestProvisionWorkflowAction extends BaseRestHandler { + + private static final String PROVISION_WORKFLOW_ACTION = "provision_workflow_action"; + + // TODO : move to common values class, pending implementation + /** + * Field name for workflow Id, the document Id of the indexed use case template + */ + public static final String WORKFLOW_ID = "workflowId"; + + /** + * Instantiates a new RestProvisionWorkflowAction + */ + public RestProvisionWorkflowAction() { + // TODO : Pass settings and cluster service to constructor and add settings update consumer for request timeout value + } + + @Override + public String getName() { + return PROVISION_WORKFLOW_ACTION; + } + + @Override + public List routes() { + return ImmutableList.of( + // Provision workflow from inline use case template + new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s/%s", FlowFrameworkPlugin.WORKFLOWS_URI, "_provision")), + // Provision workflow from indexed use case template + new Route( + RestRequest.Method.POST, + String.format(Locale.ROOT, "%s/{%s}/%s", FlowFrameworkPlugin.WORKFLOWS_URI, WORKFLOW_ID, "_provision") + ) + ); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + + String workflowId = request.param(WORKFLOW_ID, null); + Template template = null; + + if (request.hasContent()) { + template = Template.parse(request.content().utf8ToString()); + } + + // Validate workflow request inputs + if (workflowId == null && template == null) { + throw new IOException("WorkflowId and template cannot be both null"); + } + + // Create request and provision + WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template); + return channel -> client.execute(ProvisionWorkflowAction.INSTANCE, workflowRequest, new RestToXContentListener<>(channel)); + } + +} diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowAction.java new file mode 100644 index 000000000..8e38c4189 --- /dev/null +++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowAction.java @@ -0,0 +1,27 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.transport; + +import org.opensearch.action.ActionType; + +/** + * External Action for public facing RestCreateWorkflowActiom + */ +public class CreateWorkflowAction extends ActionType { + + // TODO : Determine external action prefix for plugin + /** The name of this action */ + public static final String NAME = "workflows/create"; + /** An instance of this action */ + public static final CreateWorkflowAction INSTANCE = new CreateWorkflowAction(); + + private CreateWorkflowAction() { + super(NAME, WorkflowResponse::new); + } +} diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java new file mode 100644 index 000000000..555b8f767 --- /dev/null +++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java @@ -0,0 +1,65 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.transport; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.client.Client; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.action.ActionListener; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +/** + * Transport Action to index or update a use case template within the Global Context + */ +public class CreateWorkflowTransportAction extends HandledTransportAction { + + private final Logger logger = LogManager.getLogger(CreateWorkflowTransportAction.class); + + private final Client client; + + /** + * Intantiates a new CreateWorkflowTransportAction + * @param transportService the TransportService + * @param actionFilters action filters + * @param client the node client to interact with an index + */ + @Inject + public CreateWorkflowTransportAction(TransportService transportService, ActionFilters actionFilters, Client client) { + super(CreateWorkflowAction.NAME, transportService, actionFilters, WorkflowRequest::new); + this.client = client; + } + + @Override + protected void doExecute(Task task, WorkflowRequest request, ActionListener listener) { + + String workflowId; + // TODO : Check if global context index exists, and if it does not then create + + if (request.getWorkflowId() == null) { + // TODO : Create new entry + // TODO : Insert doc + + // TODO : get document ID + workflowId = ""; + // TODO : check if state index exists, and if it does not, then create + // TODO : insert state index doc, mapped with documentId, defaulted to NOT_STARTED + } else { + // TODO : Update existing entry + workflowId = request.getWorkflowId(); + // TODO : Update state index entry, default back to NOT_STARTED + } + + listener.onResponse(new WorkflowResponse(workflowId)); + } + +} diff --git a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowAction.java new file mode 100644 index 000000000..7167c4357 --- /dev/null +++ b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowAction.java @@ -0,0 +1,27 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.transport; + +import org.opensearch.action.ActionType; + +/** + * External Action for public facing RestProvisionWorkflowAction + */ +public class ProvisionWorkflowAction extends ActionType { + + // TODO : Determine external action prefix for plugin + /** The name of this action */ + public static final String NAME = "workflows/provision"; + /** An instance of this action */ + public static final ProvisionWorkflowAction INSTANCE = new ProvisionWorkflowAction(); + + private ProvisionWorkflowAction() { + super(NAME, WorkflowResponse::new); + } +} diff --git a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java new file mode 100644 index 000000000..855905af2 --- /dev/null +++ b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java @@ -0,0 +1,69 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.transport; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.client.Client; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.action.ActionListener; +import org.opensearch.flowframework.model.Template; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +/** + * Transport Action to provision a workflow from a stored use case template + */ +public class ProvisionWorkflowTransportAction extends HandledTransportAction { + private final Logger logger = LogManager.getLogger(ProvisionWorkflowTransportAction.class); + private final Client client; + + /** + * Instantiates a new ProvisionWorkflowTransportAction + * @param transportService The TransportService + * @param actionFilters action filters + * @param client The node client to retrieve a stored use case template + */ + @Inject + public ProvisionWorkflowTransportAction(TransportService transportService, ActionFilters actionFilters, Client client) { + super(ProvisionWorkflowAction.NAME, transportService, actionFilters, WorkflowRequest::new); + this.client = client; + } + + @Override + protected void doExecute(Task task, WorkflowRequest request, ActionListener listener) { + + if (request.getWorkflowId() == null) { + // Workflow provisioning from inline template, first parse and then index the given use case template + client.execute(CreateWorkflowAction.INSTANCE, request, ActionListener.wrap(workflowResponse -> { + String workflowId = workflowResponse.getWorkflowId(); + Template template = request.getTemplate(); + + // TODO : Use node client to update state index, given workflowId + // TODO : Pass workflowId and template to execution + + listener.onResponse(new WorkflowResponse(workflowId)); + + }, exception -> { listener.onFailure(exception); })); + + } else { + + // Use case template has been previously saved, retrieve entry and execute + String workflowId = request.getWorkflowId(); + + // TODO : use node client to update state index, given workflowId + // TODO : Retrieve template from global context index using node client execute + + listener.onResponse(new WorkflowResponse(workflowId)); + } + } + +} diff --git a/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java b/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java new file mode 100644 index 000000000..5e17248f1 --- /dev/null +++ b/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java @@ -0,0 +1,96 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.transport; + +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.common.Nullable; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.flowframework.model.Template; + +import java.io.IOException; + +/** + * Transport Request to create and provision a workflow + */ +public class WorkflowRequest extends ActionRequest { + + /** + * The documentId of the workflow entry within the Global Context index + */ + @Nullable + private String workflowId; + /** + * The use case template to index + */ + @Nullable + private Template template; + + /** + * Instantiates a new WorkflowRequest + * @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 = workflowId; + this.template = template; + } + + /** + * Instantiates a new Workflow request + * @param in The input stream to read from + * @throws IOException If the stream cannot be read properly + */ + public WorkflowRequest(StreamInput in) throws IOException { + super(in); + this.workflowId = in.readOptionalString(); + if (in.readBoolean()) { + this.template = Template.parse(in.readString()); + } else { + this.template = null; + } + } + + /** + * Gets the workflow Id of the request + * @return the workflow Id + */ + @Nullable + public String getWorkflowId() { + return this.workflowId; + } + + /** + * Gets the use case template of the request + * @return the use case template + */ + @Nullable + public Template getTemplate() { + return this.template; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeOptionalString(workflowId); + if (template != null) { + out.writeBoolean(true); + out.writeString(template.toJson()); + } else { + out.writeBoolean(false); + } + } + + @Override + public ActionRequestValidationException validate() { + return null; + } + +} diff --git a/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java b/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java new file mode 100644 index 000000000..eb4bfe089 --- /dev/null +++ b/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java @@ -0,0 +1,67 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.transport; + +import org.opensearch.core.action.ActionResponse; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.flowframework.rest.RestCreateWorkflowAction; + +import java.io.IOException; + +/** + * Transport Response from creating or provisioning a workflow + */ +public class WorkflowResponse extends ActionResponse implements ToXContentObject { + + /** + * The documentId of the workflow entry within the Global Context index + */ + private String workflowId; + + /** + * Instantiates a new WorkflowResponse from params + * @param workflowId the documentId of the indexed use case template + */ + public WorkflowResponse(String workflowId) { + this.workflowId = workflowId; + } + + /** + * Instatiates a new WorkflowResponse from an input stream + * @param in the input stream to read from + * @throws IOException if the workflowId cannot be read from the input stream + */ + public WorkflowResponse(StreamInput in) throws IOException { + super(in); + this.workflowId = in.readString(); + } + + /** + * Gets the workflowId of this repsonse + * @return the workflowId + */ + public String getWorkflowId() { + return this.workflowId; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeString(workflowId); + } + + // TODO : Replace WORKFLOW_ID after string is moved to common values class + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + return builder.startObject().field(RestCreateWorkflowAction.WORKFLOW_ID, this.workflowId).endObject(); + } + +} From c32bf3101c9b97e02d9788fe8935f8d376c5a0c5 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Mon, 2 Oct 2023 23:06:23 +0000 Subject: [PATCH 02/23] Addressing PR comments, seting params to snake case, removing redundant param default value, setting workflow request to read/write optional string Signed-off-by: Joshua Palis --- .../rest/RestCreateWorkflowAction.java | 4 ++-- .../rest/RestProvisionWorkflowAction.java | 6 +++--- .../flowframework/transport/WorkflowRequest.java | 14 +++----------- 3 files changed, 8 insertions(+), 16 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java index 5a2b352fe..926ebd6be 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java @@ -33,7 +33,7 @@ public class RestCreateWorkflowAction extends BaseRestHandler { /** * Field name for workflow Id, the document Id of the indexed use case template */ - public static final String WORKFLOW_ID = "workflowId"; + public static final String WORKFLOW_ID = "workflow_id"; /** * Intantiates a new RestCreateWorkflowAction @@ -60,7 +60,7 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - String workflowId = request.param(WORKFLOW_ID, null); + String workflowId = request.param(WORKFLOW_ID); Template template = Template.parse(request.content().utf8ToString()); WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template); return channel -> client.execute(CreateWorkflowAction.INSTANCE, workflowRequest, new RestToXContentListener<>(channel)); diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java index b1d8fb671..f4cb55e51 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java @@ -33,7 +33,7 @@ public class RestProvisionWorkflowAction extends BaseRestHandler { /** * Field name for workflow Id, the document Id of the indexed use case template */ - public static final String WORKFLOW_ID = "workflowId"; + public static final String WORKFLOW_ID = "workflow_id"; /** * Instantiates a new RestProvisionWorkflowAction @@ -63,7 +63,7 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - String workflowId = request.param(WORKFLOW_ID, null); + String workflowId = request.param(WORKFLOW_ID); Template template = null; if (request.hasContent()) { @@ -72,7 +72,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli // Validate workflow request inputs if (workflowId == null && template == null) { - throw new IOException("WorkflowId and template cannot be both null"); + throw new IOException("workflow_id and template cannot be both null"); } // Create request and provision diff --git a/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java b/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java index 5e17248f1..0b105552f 100644 --- a/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java +++ b/src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java @@ -51,11 +51,8 @@ public WorkflowRequest(@Nullable String workflowId, @Nullable Template template) public WorkflowRequest(StreamInput in) throws IOException { super(in); this.workflowId = in.readOptionalString(); - if (in.readBoolean()) { - this.template = Template.parse(in.readString()); - } else { - this.template = null; - } + String templateJson = in.readOptionalString(); + this.template = templateJson == null ? null : Template.parse(templateJson); } /** @@ -80,12 +77,7 @@ public Template getTemplate() { public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeOptionalString(workflowId); - if (template != null) { - out.writeBoolean(true); - out.writeString(template.toJson()); - } else { - out.writeBoolean(false); - } + out.writeOptionalString(template == null ? null : template.toJson()); } @Override From cc32079134f718db2cd66d863be187dc6f9ab7af Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Wed, 4 Oct 2023 22:37:13 +0000 Subject: [PATCH 03/23] Introducing getExecutorBuilders extension point to FlowFramworkPlugin, added FixedExecutorBuilder thread pool for provisioning tasks, set up async workflow execution, added TODOs for state/GC index handling Signed-off-by: Joshua Palis --- .../flowframework/FlowFrameworkPlugin.java | 27 ++++- .../ProvisionWorkflowTransportAction.java | 112 ++++++++++++++++-- 2 files changed, 131 insertions(+), 8 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java index 3c8b15f49..a5b0cea13 100644 --- a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java +++ b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java @@ -37,6 +37,8 @@ import org.opensearch.rest.RestController; import org.opensearch.rest.RestHandler; import org.opensearch.script.ScriptService; +import org.opensearch.threadpool.ExecutorBuilder; +import org.opensearch.threadpool.FixedExecutorBuilder; import org.opensearch.threadpool.ThreadPool; import org.opensearch.watcher.ResourceWatcherService; @@ -49,14 +51,23 @@ */ public class FlowFrameworkPlugin extends Plugin implements ActionPlugin { + // TODO : Move names to common values class /** * The base URI for this plugin's rest actions */ - public static final String AI_FLOW_FRAMEWORK_BASE_URI = "/_plugins/_ai_flow"; + public static final String AI_FLOW_FRAMEWORK_BASE_URI = "/_plugins/_flow_framework"; /** * The URI for this plugin's workflow rest actions */ public static final String WORKFLOWS_URI = AI_FLOW_FRAMEWORK_BASE_URI + "/workflows"; + /** + * Flow Framework plugin thread pool name prefix + */ + public static final String FLOW_FRAMEWORK_THREAD_POOL_PREFIX = "thread_pool.flow_framework."; + /** + * The provision workflow thread pool name + */ + public static final String PROVISION_THREAD_POOL = "opensearch_workflow_provision"; /** * Instantiate this plugin. @@ -104,4 +115,18 @@ public List getRestHandlers( ); } + @Override + public List> getExecutorBuilders(Settings settings) { + // TODO : Determine final size/queueSize values for the provision thread pool + FixedExecutorBuilder provisionThreadPool = new FixedExecutorBuilder( + settings, + PROVISION_THREAD_POOL, + 1, + 10, + FLOW_FRAMEWORK_THREAD_POOL_PREFIX + PROVISION_THREAD_POOL, + false + ); + return ImmutableList.of(provisionThreadPool); + } + } diff --git a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java index 855905af2..bc30cae3a 100644 --- a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java @@ -16,26 +16,60 @@ import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; import org.opensearch.flowframework.model.Template; +import org.opensearch.flowframework.model.Workflow; +import org.opensearch.flowframework.workflow.ProcessNode; +import org.opensearch.flowframework.workflow.WorkflowProcessSorter; import org.opensearch.tasks.Task; +import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.stream.Collectors; + +import static org.opensearch.flowframework.FlowFrameworkPlugin.PROVISION_THREAD_POOL; + /** * Transport Action to provision a workflow from a stored use case template */ public class ProvisionWorkflowTransportAction extends HandledTransportAction { + private final Logger logger = LogManager.getLogger(ProvisionWorkflowTransportAction.class); + + // TODO : Move to common values class, pending implementation + /** + * The name of the provision workflow within the use case template + */ + private static final String PROVISION_WORKFLOW = "provision"; + + private final ThreadPool threadPool; private final Client client; + private final WorkflowProcessSorter workflowProcessSorter; /** * Instantiates a new ProvisionWorkflowTransportAction * @param transportService The TransportService * @param actionFilters action filters + * @param threadPool The OpenSearch thread pool * @param client The node client to retrieve a stored use case template + * @param workflowProcessSorter Utility class to generate a togologically sorted list of Process nodes */ @Inject - public ProvisionWorkflowTransportAction(TransportService transportService, ActionFilters actionFilters, Client client) { + public ProvisionWorkflowTransportAction( + TransportService transportService, + ActionFilters actionFilters, + ThreadPool threadPool, + Client client, + WorkflowProcessSorter workflowProcessSorter + ) { super(ProvisionWorkflowAction.NAME, transportService, actionFilters, WorkflowRequest::new); + this.threadPool = threadPool; this.client = client; + this.workflowProcessSorter = workflowProcessSorter; } @Override @@ -47,22 +81,86 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener { listener.onFailure(exception); })); + // Asychronously begin provision workflow excecution + executeWorkflowAsync(workflowId, template.workflows().get(PROVISION_WORKFLOW)); + }, exception -> { listener.onFailure(exception); })); } else { - // Use case template has been previously saved, retrieve entry and execute String workflowId = request.getWorkflowId(); - // TODO : use node client to update state index, given workflowId - // TODO : Retrieve template from global context index using node client execute + // TODO : Retrieve template from global context index using node client + Template template = null; // temporary, remove later + + // TODO : use node client to update state index entry to PROVISIONING, given workflowId listener.onResponse(new WorkflowResponse(workflowId)); + executeWorkflowAsync(workflowId, template.workflows().get(PROVISION_WORKFLOW)); + } + } + + /** + * Retrieves a thread from the provision thread pool to execute a workflow + * @param workflowId The id of the workflow + * @param workflow The workflow to execute + */ + private void executeWorkflowAsync(String workflowId, Workflow workflow) { + // TODO : Update Action listener type to State index Request + ActionListener provisionWorkflowListener = ActionListener.wrap(response -> { + logger.info("Provisioning completed successuflly for workflow {}", workflowId); + + // TODO : Create State index request to update STATE entry status to READY + }, exception -> { + logger.error("Provisioning failed for workflow {} : {}", workflowId, exception); + + // TODO : Create State index request to update STATE entry status to FAILED + }); + try { + threadPool.executor(PROVISION_THREAD_POOL).execute(() -> { executeWorkflow(workflow, provisionWorkflowListener); }); + } catch (Exception exception) { + provisionWorkflowListener.onFailure(exception); + } + } + + /** + * Topologically sorts a given workflow into a sequence of ProcessNodes and executes the workflow + * @param workflow The workflow to execute + * @param workflowListener The listener that updates the status of a workflow execution + */ + private void executeWorkflow(Workflow workflow, ActionListener workflowListener) { + + List processSequence = workflowProcessSorter.sortProcessNodes(workflow); + List> workflowFutureList = new ArrayList<>(); + + for (ProcessNode processNode : processSequence) { + List predecessors = processNode.predecessors(); + + logger.info( + "Queueing process [{}].{}", + processNode.id(), + predecessors.isEmpty() + ? " Can start immediately!" + : String.format( + Locale.getDefault(), + " Must wait for [%s] to complete first.", + predecessors.stream().map(p -> p.id()).collect(Collectors.joining(", ")) + ) + ); + + workflowFutureList.add(processNode.execute()); + } + try { + // Attempt to join each workflow step future, may throw a CompletionException if any step completes exceptionally + workflowFutureList.forEach(CompletableFuture::join); + + // TODO : Create State Index request with provisioning state, start time, end time, etc, pending implementation. String for now + workflowListener.onResponse("READY"); + } catch (CancellationException | CompletionException ex) { + workflowListener.onFailure(ex); } } From a51c68150c9d8026e3a93e3648eedc2d8601e5d6 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 5 Oct 2023 20:58:09 +0000 Subject: [PATCH 04/23] updating unit tests for FlowFrameworkPluginTests, adding WorkflowRequestResponse unit tests Signed-off-by: Joshua Palis --- .../FlowFrameworkPluginTests.java | 3 + .../WorkflowRequestResponseTests.java | 113 ++++++++++++++++++ 2 files changed, 116 insertions(+) create mode 100644 src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java index d211e3928..0d09fe239 100644 --- a/src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java +++ b/src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java @@ -42,6 +42,9 @@ public void tearDown() throws Exception { public void testPlugin() throws IOException { try (FlowFrameworkPlugin ffp = new FlowFrameworkPlugin()) { assertEquals(2, ffp.createComponents(client, null, threadPool, null, null, null, null, null, null, null, null).size()); + assertEquals(2, ffp.getRestHandlers(null, null, null, null, null, null, null).size()); + assertEquals(2, ffp.getActions().size()); + assertEquals(1, ffp.getExecutorBuilders(null).size()); } } } diff --git a/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java b/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java new file mode 100644 index 000000000..85541e0c8 --- /dev/null +++ b/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java @@ -0,0 +1,113 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.transport; + +import org.opensearch.Version; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.flowframework.model.Template; +import org.opensearch.flowframework.model.Workflow; +import org.opensearch.flowframework.model.WorkflowEdge; +import org.opensearch.flowframework.model.WorkflowNode; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +public class WorkflowRequestResponseTests extends OpenSearchTestCase { + + private Template template; + + @Override + public void setUp() throws Exception { + super.setUp(); + List operations = List.of("operation"); + Version templateVersion = Version.fromString("1.0.0"); + List compatibilityVersions = List.of(Version.fromString("2.0.0"), Version.fromString("3.0.0")); + WorkflowNode nodeA = new WorkflowNode("A", "a-type", Map.of("foo", "bar")); + WorkflowNode nodeB = new WorkflowNode("B", "b-type", Map.of("baz", "qux")); + WorkflowEdge edgeAB = new WorkflowEdge("A", "B"); + List nodes = List.of(nodeA, nodeB); + List edges = List.of(edgeAB); + Workflow workflow = new Workflow(Map.of("key", "value"), nodes, edges); + + this.template = new Template( + "test", + "description", + "use case", + operations, + templateVersion, + compatibilityVersions, + Map.ofEntries(Map.entry("userKey", "userValue"), Map.entry("userMapKey", Map.of("nestedKey", "nestedValue"))), + Map.of("workflow", workflow) + ); + } + + public void testNullIdWorkflowRequest() throws IOException { + WorkflowRequest nullIdRequest = new WorkflowRequest(null, template); + assertNull(nullIdRequest.getWorkflowId()); + assertEquals(template, nullIdRequest.getTemplate()); + + BytesStreamOutput out = new BytesStreamOutput(); + nullIdRequest.writeTo(out); + BytesStreamInput in = new BytesStreamInput(BytesReference.toBytes(out.bytes())); + + WorkflowRequest streamInputRequest = new WorkflowRequest(in); + + assertEquals(nullIdRequest.getWorkflowId(), streamInputRequest.getWorkflowId()); + assertEquals(nullIdRequest.getTemplate().toJson(), streamInputRequest.getTemplate().toJson()); + } + + public void testNullTemplateWorkflowRequest() throws IOException { + WorkflowRequest nullTemplateRequest = new WorkflowRequest("123", null); + assertNotNull(nullTemplateRequest.getWorkflowId()); + assertNull(nullTemplateRequest.getTemplate()); + + BytesStreamOutput out = new BytesStreamOutput(); + nullTemplateRequest.writeTo(out); + BytesStreamInput in = new BytesStreamInput(BytesReference.toBytes(out.bytes())); + + WorkflowRequest streamInputRequest = new WorkflowRequest(in); + + assertEquals(nullTemplateRequest.getWorkflowId(), streamInputRequest.getWorkflowId()); + assertEquals(nullTemplateRequest.getTemplate(), streamInputRequest.getTemplate()); + } + + public void testWorkflowRequest() throws IOException { + WorkflowRequest workflowRequest = new WorkflowRequest("123", template); + assertNotNull(workflowRequest.getWorkflowId()); + assertEquals(template, workflowRequest.getTemplate()); + + BytesStreamOutput out = new BytesStreamOutput(); + workflowRequest.writeTo(out); + BytesStreamInput in = new BytesStreamInput(BytesReference.toBytes(out.bytes())); + + WorkflowRequest streamInputRequest = new WorkflowRequest(in); + + assertEquals(workflowRequest.getWorkflowId(), streamInputRequest.getWorkflowId()); + assertEquals(workflowRequest.getTemplate().toJson(), streamInputRequest.getTemplate().toJson()); + + } + + public void testWorkflowResponse() throws IOException { + WorkflowResponse response = new WorkflowResponse("123"); + assertEquals("123", response.getWorkflowId()); + + BytesStreamOutput out = new BytesStreamOutput(); + response.writeTo(out); + BytesStreamInput in = new BytesStreamInput(BytesReference.toBytes(out.bytes())); + + WorkflowResponse streamInputResponse = new WorkflowResponse(in); + + assertEquals(response.getWorkflowId(), streamInputResponse.getWorkflowId()); + } + +} From 524407e260fd7abb1de08018bcd0c36c67f87ee3 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 5 Oct 2023 22:04:24 +0000 Subject: [PATCH 05/23] Adding validate/toXContent tests for workflow request/responses Signed-off-by: Joshua Palis --- .../transport/WorkflowRequestResponseTests.java | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java b/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java index 85541e0c8..ab4d4a068 100644 --- a/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java @@ -10,8 +10,12 @@ import org.opensearch.Version; import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.xcontent.XContentType; import org.opensearch.core.common.bytes.BytesReference; import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.flowframework.model.Template; import org.opensearch.flowframework.model.Workflow; import org.opensearch.flowframework.model.WorkflowEdge; @@ -55,6 +59,7 @@ public void testNullIdWorkflowRequest() throws IOException { WorkflowRequest nullIdRequest = new WorkflowRequest(null, template); assertNull(nullIdRequest.getWorkflowId()); assertEquals(template, nullIdRequest.getTemplate()); + assertNull(nullIdRequest.validate()); BytesStreamOutput out = new BytesStreamOutput(); nullIdRequest.writeTo(out); @@ -70,6 +75,7 @@ public void testNullTemplateWorkflowRequest() throws IOException { WorkflowRequest nullTemplateRequest = new WorkflowRequest("123", null); assertNotNull(nullTemplateRequest.getWorkflowId()); assertNull(nullTemplateRequest.getTemplate()); + assertNull(nullTemplateRequest.validate()); BytesStreamOutput out = new BytesStreamOutput(); nullTemplateRequest.writeTo(out); @@ -85,6 +91,7 @@ public void testWorkflowRequest() throws IOException { WorkflowRequest workflowRequest = new WorkflowRequest("123", template); assertNotNull(workflowRequest.getWorkflowId()); assertEquals(template, workflowRequest.getTemplate()); + assertNull(workflowRequest.validate()); BytesStreamOutput out = new BytesStreamOutput(); workflowRequest.writeTo(out); @@ -106,8 +113,12 @@ public void testWorkflowResponse() throws IOException { BytesStreamInput in = new BytesStreamInput(BytesReference.toBytes(out.bytes())); WorkflowResponse streamInputResponse = new WorkflowResponse(in); - assertEquals(response.getWorkflowId(), streamInputResponse.getWorkflowId()); + + XContentBuilder builder = MediaTypeRegistry.contentBuilder(XContentType.JSON); + response.toXContent(builder, ToXContent.EMPTY_PARAMS); + assertNotNull(builder); + assertEquals("{\"workflow_id\":\"123\"}", builder.toString()); } } From 6aaa8c79b26ed9e1acfddaaee68b3f9ceaf12625 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 5 Oct 2023 23:34:03 +0000 Subject: [PATCH 06/23] Adding unit tests for create and provision rest actions Signed-off-by: Joshua Palis --- .../rest/RestCreateWorkflowActionTests.java | 74 +++++++++++++++ .../RestProvisionWorkflowActionTests.java | 91 +++++++++++++++++++ 2 files changed, 165 insertions(+) create mode 100644 src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java create mode 100644 src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java diff --git a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java new file mode 100644 index 000000000..07ba1dc25 --- /dev/null +++ b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java @@ -0,0 +1,74 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.rest; + +import org.opensearch.client.node.NodeClient; +import org.opensearch.core.common.bytes.BytesArray; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.flowframework.FlowFrameworkPlugin; +import org.opensearch.rest.RestHandler.Route; +import org.opensearch.rest.RestRequest; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.rest.FakeRestRequest; + +import java.io.IOException; +import java.util.List; +import java.util.Locale; + +import static org.mockito.Mockito.mock; + +public class RestCreateWorkflowActionTests extends OpenSearchTestCase { + + private String invalidTemplate; + private RestCreateWorkflowAction createWorkflowRestAction; + private String createWorkflowPath; + private String updateWorkflowPath; + private NodeClient nodeClient; + + @Override + public void setUp() throws Exception { + super.setUp(); + // Invalid template configuration, missing name field + this.invalidTemplate = "{\"description\":\"description\"," + + "\"use_case\":\"use case\"," + + "\"operations\":[\"operation\"]," + + "\"version\":{\"template\":\"1.0.0\",\"compatibility\":[\"2.0.0\",\"3.0.0\"]}," + + "\"user_inputs\":{\"userKey\":\"userValue\",\"userMapKey\":{\"nestedKey\":\"nestedValue\"}}," + + "\"workflows\":{\"workflow\":{\"user_params\":{\"key\":\"value\"},\"nodes\":[{\"id\":\"A\",\"type\":\"a-type\",\"inputs\":{\"foo\":\"bar\"}},{\"id\":\"B\",\"type\":\"b-type\",\"inputs\":{\"baz\":\"qux\"}}],\"edges\":[{\"source\":\"A\",\"dest\":\"B\"}]}}}"; + this.createWorkflowRestAction = new RestCreateWorkflowAction(); + this.createWorkflowPath = String.format(Locale.ROOT, "%s", FlowFrameworkPlugin.WORKFLOWS_URI); + this.updateWorkflowPath = String.format(Locale.ROOT, "%s/{%s}", FlowFrameworkPlugin.WORKFLOWS_URI, "workflow_id"); + this.nodeClient = mock(NodeClient.class); + } + + public void testRestCreateWorkflowActionName() { + String name = createWorkflowRestAction.getName(); + assertEquals("create_workflow_action", name); + } + + public void testRestCreateWorkflowActionRoutes() { + List routes = createWorkflowRestAction.routes(); + assertEquals(2, routes.size()); + assertEquals(RestRequest.Method.POST, routes.get(0).getMethod()); + assertEquals(RestRequest.Method.PUT, routes.get(1).getMethod()); + assertEquals(this.createWorkflowPath, routes.get(0).getPath()); + assertEquals(this.updateWorkflowPath, routes.get(1).getPath()); + + } + + public void testInvalidCreateWorkflowRequest() throws IOException { + RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.POST) + .withPath(this.createWorkflowPath) + .withContent(new BytesArray(invalidTemplate), MediaTypeRegistry.JSON) + .build(); + + IOException ex = expectThrows(IOException.class, () -> { createWorkflowRestAction.prepareRequest(request, nodeClient); }); + assertEquals("An template object requires a name.", ex.getMessage()); + } +} diff --git a/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java new file mode 100644 index 000000000..75fec5bae --- /dev/null +++ b/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java @@ -0,0 +1,91 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.rest; + +import org.opensearch.client.node.NodeClient; +import org.opensearch.core.common.bytes.BytesArray; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.flowframework.FlowFrameworkPlugin; +import org.opensearch.rest.RestHandler.Route; +import org.opensearch.rest.RestRequest; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.rest.FakeRestRequest; + +import java.io.IOException; +import java.util.List; +import java.util.Locale; + +import static org.mockito.Mockito.mock; + +public class RestProvisionWorkflowActionTests extends OpenSearchTestCase { + + private String invalidTemplate; + private RestProvisionWorkflowAction provisionWorkflowRestAction; + private String provisionInlineWorkflowPath; + private String provisionSavedWorkflowPath; + private NodeClient nodeClient; + + @Override + public void setUp() throws Exception { + super.setUp(); + // Invalid template configuration, missing name field + this.invalidTemplate = "{\"description\":\"description\"," + + "\"use_case\":\"use case\"," + + "\"operations\":[\"operation\"]," + + "\"version\":{\"template\":\"1.0.0\",\"compatibility\":[\"2.0.0\",\"3.0.0\"]}," + + "\"user_inputs\":{\"userKey\":\"userValue\",\"userMapKey\":{\"nestedKey\":\"nestedValue\"}}," + + "\"workflows\":{\"workflow\":{\"user_params\":{\"key\":\"value\"},\"nodes\":[{\"id\":\"A\",\"type\":\"a-type\",\"inputs\":{\"foo\":\"bar\"}},{\"id\":\"B\",\"type\":\"b-type\",\"inputs\":{\"baz\":\"qux\"}}],\"edges\":[{\"source\":\"A\",\"dest\":\"B\"}]}}}"; + this.provisionWorkflowRestAction = new RestProvisionWorkflowAction(); + this.provisionInlineWorkflowPath = String.format(Locale.ROOT, "%s/%s", FlowFrameworkPlugin.WORKFLOWS_URI, "_provision"); + this.provisionSavedWorkflowPath = String.format( + Locale.ROOT, + "%s/{%s}/%s", + FlowFrameworkPlugin.WORKFLOWS_URI, + "workflow_id", + "_provision" + ); + this.nodeClient = mock(NodeClient.class); + } + + public void testRestProvisionWorkflowActionName() { + String name = provisionWorkflowRestAction.getName(); + assertEquals("provision_workflow_action", name); + } + + public void testRestProvisiionWorkflowActionRoutes() { + List routes = provisionWorkflowRestAction.routes(); + assertEquals(2, routes.size()); + assertEquals(RestRequest.Method.POST, routes.get(0).getMethod()); + assertEquals(RestRequest.Method.POST, routes.get(1).getMethod()); + assertEquals(this.provisionInlineWorkflowPath, routes.get(0).getPath()); + assertEquals(this.provisionSavedWorkflowPath, routes.get(1).getPath()); + } + + public void testNullWorkflowIdAndTemplate() throws IOException { + + // Request with no content or params + RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.POST) + .withPath(this.provisionInlineWorkflowPath) + .build(); + + IOException ex = expectThrows(IOException.class, () -> { provisionWorkflowRestAction.prepareRequest(request, nodeClient); }); + assertEquals("workflow_id and template cannot be both null", ex.getMessage()); + } + + public void testInvalidProvisionInlineWorkflowRequest() throws IOException { + RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.POST) + .withPath(this.provisionInlineWorkflowPath) + .withContent(new BytesArray(invalidTemplate), MediaTypeRegistry.JSON) + .build(); + + IOException ex = expectThrows(IOException.class, () -> { provisionWorkflowRestAction.prepareRequest(request, nodeClient); }); + assertEquals("An template object requires a name.", ex.getMessage()); + } + +} From 791f943094c713d70865c24c97520881a69532c4 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Mon, 9 Oct 2023 18:39:14 +0000 Subject: [PATCH 07/23] Addressing PR comments (Part 1). Moving common vlaues to CommonValue class Signed-off-by: Joshua Palis --- .../flowframework/FlowFrameworkPlugin.java | 21 +++---------------- .../flowframework/common/CommonValue.java | 11 ++++++++++ .../rest/RestCreateWorkflowAction.java | 16 ++++++-------- .../rest/RestProvisionWorkflowAction.java | 11 +++++----- .../ProvisionWorkflowTransportAction.java | 2 +- .../transport/WorkflowResponse.java | 5 +++-- .../rest/RestCreateWorkflowActionTests.java | 6 +++--- .../RestProvisionWorkflowActionTests.java | 12 +++-------- .../WorkflowRequestResponseTests.java | 4 +++- 9 files changed, 38 insertions(+), 50 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java index cb4d3f068..6f0ca9000 100644 --- a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java +++ b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java @@ -46,29 +46,14 @@ import java.util.List; import java.util.function.Supplier; +import static org.opensearch.flowframework.common.CommonValue.FLOW_FRAMEWORK_THREAD_POOL_PREFIX; +import static org.opensearch.flowframework.common.CommonValue.PROVISION_THREAD_POOL; + /** * An OpenSearch plugin that enables builders to innovate AI apps on OpenSearch. */ public class FlowFrameworkPlugin extends Plugin implements ActionPlugin { - // TODO : Move names to common values class - /** - * The base URI for this plugin's rest actions - */ - public static final String AI_FLOW_FRAMEWORK_BASE_URI = "/_plugins/_flow_framework"; - /** - * The URI for this plugin's workflow rest actions - */ - public static final String WORKFLOWS_URI = AI_FLOW_FRAMEWORK_BASE_URI + "/workflows"; - /** - * Flow Framework plugin thread pool name prefix - */ - public static final String FLOW_FRAMEWORK_THREAD_POOL_PREFIX = "thread_pool.flow_framework."; - /** - * The provision workflow thread pool name - */ - public static final String PROVISION_THREAD_POOL = "opensearch_workflow_provision"; - /** * Instantiate this plugin. */ diff --git a/src/main/java/org/opensearch/flowframework/common/CommonValue.java b/src/main/java/org/opensearch/flowframework/common/CommonValue.java index a8fdf2929..0aa231ca2 100644 --- a/src/main/java/org/opensearch/flowframework/common/CommonValue.java +++ b/src/main/java/org/opensearch/flowframework/common/CommonValue.java @@ -19,4 +19,15 @@ public class CommonValue { public static final String GLOBAL_CONTEXT_INDEX = ".plugins-ai-global-context"; public static final String GLOBAL_CONTEXT_INDEX_MAPPING = "mappings/global-context.json"; public static final Integer GLOBAL_CONTEXT_INDEX_VERSION = 1; + + /** The base URI for this plugin's rest actions */ + public static final String AI_FLOW_FRAMEWORK_BASE_URI = "/_plugins/_flow_framework"; + /** The URI for this plugin's workflow rest actions */ + public static final String WORKFLOWS_URI = AI_FLOW_FRAMEWORK_BASE_URI + "/workflows"; + /** Flow Framework plugin thread pool name prefix */ + public static final String FLOW_FRAMEWORK_THREAD_POOL_PREFIX = "thread_pool.flow_framework."; + /** The provision workflow thread pool name */ + public static final String PROVISION_THREAD_POOL = "opensearch_workflow_provision"; + /** Field name for workflow Id, the document Id of the indexed use case template */ + public static final String WORKFLOW_ID = "workflow_id"; } diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java index 926ebd6be..29038405e 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java @@ -10,7 +10,6 @@ import com.google.common.collect.ImmutableList; import org.opensearch.client.node.NodeClient; -import org.opensearch.flowframework.FlowFrameworkPlugin; import org.opensearch.flowframework.model.Template; import org.opensearch.flowframework.transport.CreateWorkflowAction; import org.opensearch.flowframework.transport.WorkflowRequest; @@ -22,6 +21,9 @@ import java.util.List; import java.util.Locale; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOWS_URI; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID; + /** * Rest Action to facilitate requests to create and update a use case template */ @@ -29,12 +31,6 @@ public class RestCreateWorkflowAction extends BaseRestHandler { private static final String CREATE_WORKFLOW_ACTION = "create_workflow_action"; - // TODO : move to common values class, pending implementation - /** - * Field name for workflow Id, the document Id of the indexed use case template - */ - public static final String WORKFLOW_ID = "workflow_id"; - /** * Intantiates a new RestCreateWorkflowAction */ @@ -51,9 +47,9 @@ public String getName() { public List routes() { return ImmutableList.of( // Create new workflow - new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s", FlowFrameworkPlugin.WORKFLOWS_URI)), - // Update workflow - new Route(RestRequest.Method.PUT, String.format(Locale.ROOT, "%s/{%s}", FlowFrameworkPlugin.WORKFLOWS_URI, WORKFLOW_ID)) + new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s", WORKFLOWS_URI)), + // Update use case template + new Route(RestRequest.Method.PUT, String.format(Locale.ROOT, "%s/{%s}", WORKFLOWS_URI, WORKFLOW_ID)) ); } diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java index f4cb55e51..2c2b532b4 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java @@ -10,7 +10,6 @@ import com.google.common.collect.ImmutableList; import org.opensearch.client.node.NodeClient; -import org.opensearch.flowframework.FlowFrameworkPlugin; import org.opensearch.flowframework.model.Template; import org.opensearch.flowframework.transport.ProvisionWorkflowAction; import org.opensearch.flowframework.transport.WorkflowRequest; @@ -22,6 +21,9 @@ import java.util.List; import java.util.Locale; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOWS_URI; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID; + /** * Rest action to facilitate requests to provision a workflow from an inline defined or stored use case template */ @@ -51,12 +53,9 @@ public String getName() { public List routes() { return ImmutableList.of( // Provision workflow from inline use case template - new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s/%s", FlowFrameworkPlugin.WORKFLOWS_URI, "_provision")), + new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s/%s", WORKFLOWS_URI, "_provision")), // Provision workflow from indexed use case template - new Route( - RestRequest.Method.POST, - String.format(Locale.ROOT, "%s/{%s}/%s", FlowFrameworkPlugin.WORKFLOWS_URI, WORKFLOW_ID, "_provision") - ) + new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s/{%s}/%s", WORKFLOWS_URI, WORKFLOW_ID, "_provision")) ); } diff --git a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java index bc30cae3a..bf7a725a7 100644 --- a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java @@ -31,7 +31,7 @@ import java.util.concurrent.CompletionException; import java.util.stream.Collectors; -import static org.opensearch.flowframework.FlowFrameworkPlugin.PROVISION_THREAD_POOL; +import static org.opensearch.flowframework.common.CommonValue.PROVISION_THREAD_POOL; /** * Transport Action to provision a workflow from a stored use case template diff --git a/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java b/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java index eb4bfe089..05c26d66f 100644 --- a/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java +++ b/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java @@ -13,10 +13,11 @@ import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.flowframework.rest.RestCreateWorkflowAction; import java.io.IOException; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID; + /** * Transport Response from creating or provisioning a workflow */ @@ -61,7 +62,7 @@ public void writeTo(StreamOutput out) throws IOException { // TODO : Replace WORKFLOW_ID after string is moved to common values class @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { - return builder.startObject().field(RestCreateWorkflowAction.WORKFLOW_ID, this.workflowId).endObject(); + return builder.startObject().field(WORKFLOW_ID, this.workflowId).endObject(); } } diff --git a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java index 07ba1dc25..057c99b35 100644 --- a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java +++ b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java @@ -11,7 +11,6 @@ import org.opensearch.client.node.NodeClient; import org.opensearch.core.common.bytes.BytesArray; import org.opensearch.core.xcontent.MediaTypeRegistry; -import org.opensearch.flowframework.FlowFrameworkPlugin; import org.opensearch.rest.RestHandler.Route; import org.opensearch.rest.RestRequest; import org.opensearch.test.OpenSearchTestCase; @@ -21,6 +20,7 @@ import java.util.List; import java.util.Locale; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOWS_URI; import static org.mockito.Mockito.mock; public class RestCreateWorkflowActionTests extends OpenSearchTestCase { @@ -42,8 +42,8 @@ public void setUp() throws Exception { + "\"user_inputs\":{\"userKey\":\"userValue\",\"userMapKey\":{\"nestedKey\":\"nestedValue\"}}," + "\"workflows\":{\"workflow\":{\"user_params\":{\"key\":\"value\"},\"nodes\":[{\"id\":\"A\",\"type\":\"a-type\",\"inputs\":{\"foo\":\"bar\"}},{\"id\":\"B\",\"type\":\"b-type\",\"inputs\":{\"baz\":\"qux\"}}],\"edges\":[{\"source\":\"A\",\"dest\":\"B\"}]}}}"; this.createWorkflowRestAction = new RestCreateWorkflowAction(); - this.createWorkflowPath = String.format(Locale.ROOT, "%s", FlowFrameworkPlugin.WORKFLOWS_URI); - this.updateWorkflowPath = String.format(Locale.ROOT, "%s/{%s}", FlowFrameworkPlugin.WORKFLOWS_URI, "workflow_id"); + this.createWorkflowPath = String.format(Locale.ROOT, "%s", WORKFLOWS_URI); + this.updateWorkflowPath = String.format(Locale.ROOT, "%s/{%s}", WORKFLOWS_URI, "workflow_id"); this.nodeClient = mock(NodeClient.class); } diff --git a/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java index 75fec5bae..f69dfe401 100644 --- a/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java +++ b/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java @@ -11,7 +11,6 @@ import org.opensearch.client.node.NodeClient; import org.opensearch.core.common.bytes.BytesArray; import org.opensearch.core.xcontent.MediaTypeRegistry; -import org.opensearch.flowframework.FlowFrameworkPlugin; import org.opensearch.rest.RestHandler.Route; import org.opensearch.rest.RestRequest; import org.opensearch.test.OpenSearchTestCase; @@ -21,6 +20,7 @@ import java.util.List; import java.util.Locale; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOWS_URI; import static org.mockito.Mockito.mock; public class RestProvisionWorkflowActionTests extends OpenSearchTestCase { @@ -42,14 +42,8 @@ public void setUp() throws Exception { + "\"user_inputs\":{\"userKey\":\"userValue\",\"userMapKey\":{\"nestedKey\":\"nestedValue\"}}," + "\"workflows\":{\"workflow\":{\"user_params\":{\"key\":\"value\"},\"nodes\":[{\"id\":\"A\",\"type\":\"a-type\",\"inputs\":{\"foo\":\"bar\"}},{\"id\":\"B\",\"type\":\"b-type\",\"inputs\":{\"baz\":\"qux\"}}],\"edges\":[{\"source\":\"A\",\"dest\":\"B\"}]}}}"; this.provisionWorkflowRestAction = new RestProvisionWorkflowAction(); - this.provisionInlineWorkflowPath = String.format(Locale.ROOT, "%s/%s", FlowFrameworkPlugin.WORKFLOWS_URI, "_provision"); - this.provisionSavedWorkflowPath = String.format( - Locale.ROOT, - "%s/{%s}/%s", - FlowFrameworkPlugin.WORKFLOWS_URI, - "workflow_id", - "_provision" - ); + this.provisionInlineWorkflowPath = String.format(Locale.ROOT, "%s/%s", WORKFLOWS_URI, "_provision"); + this.provisionSavedWorkflowPath = String.format(Locale.ROOT, "%s/{%s}/%s", WORKFLOWS_URI, "workflow_id", "_provision"); this.nodeClient = mock(NodeClient.class); } diff --git a/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java b/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java index ab4d4a068..cc5c19a09 100644 --- a/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/WorkflowRequestResponseTests.java @@ -51,7 +51,9 @@ public void setUp() throws Exception { templateVersion, compatibilityVersions, Map.ofEntries(Map.entry("userKey", "userValue"), Map.entry("userMapKey", Map.of("nestedKey", "nestedValue"))), - Map.of("workflow", workflow) + Map.of("workflow", workflow), + Map.of("outputKey", "outputValue"), + Map.of("resourceKey", "resourceValue") ); } From c7c819b6f193a90cd46e9b585e56b50b56717c1d Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Mon, 9 Oct 2023 21:10:32 +0000 Subject: [PATCH 08/23] Addressing PR comments (Part 2), adding globalcontexthandler to create components, added updateTemplate(), indexExists() methods to handler and createIndex step respecitvely. Implemented CreateWorkflow/ProvisionWorkflow transport actions Signed-off-by: Joshua Palis --- .../flowframework/FlowFrameworkPlugin.java | 13 +++-- .../flowframework/common/CommonValue.java | 8 ++- .../indices/GlobalContextHandler.java | 32 ++++++++++++ .../rest/RestProvisionWorkflowAction.java | 24 +++------ .../CreateWorkflowTransportAction.java | 47 +++++++++-------- .../ProvisionWorkflowTransportAction.java | 51 +++++++++---------- .../workflow/CreateIndexStep.java | 10 ++++ .../FlowFrameworkPluginTests.java | 7 ++- .../RestProvisionWorkflowActionTests.java | 32 ++++-------- 9 files changed, 129 insertions(+), 95 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java index 6f0ca9000..810674c47 100644 --- a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java +++ b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java @@ -18,17 +18,20 @@ import org.opensearch.common.settings.IndexScopedSettings; import org.opensearch.common.settings.Settings; import org.opensearch.common.settings.SettingsFilter; +import org.opensearch.common.util.concurrent.OpenSearchExecutors; import org.opensearch.core.action.ActionResponse; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.env.Environment; import org.opensearch.env.NodeEnvironment; +import org.opensearch.flowframework.indices.GlobalContextHandler; import org.opensearch.flowframework.rest.RestCreateWorkflowAction; import org.opensearch.flowframework.rest.RestProvisionWorkflowAction; import org.opensearch.flowframework.transport.CreateWorkflowAction; import org.opensearch.flowframework.transport.CreateWorkflowTransportAction; import org.opensearch.flowframework.transport.ProvisionWorkflowAction; import org.opensearch.flowframework.transport.ProvisionWorkflowTransportAction; +import org.opensearch.flowframework.workflow.CreateIndexStep; import org.opensearch.flowframework.workflow.WorkflowProcessSorter; import org.opensearch.flowframework.workflow.WorkflowStepFactory; import org.opensearch.plugins.ActionPlugin; @@ -76,7 +79,10 @@ public Collection createComponents( WorkflowStepFactory workflowStepFactory = new WorkflowStepFactory(clusterService, client); WorkflowProcessSorter workflowProcessSorter = new WorkflowProcessSorter(workflowStepFactory, threadPool); - return ImmutableList.of(workflowStepFactory, workflowProcessSorter); + // TODO : Refactor, move system index creation/associated methods outside of the CreateIndexStep + GlobalContextHandler globalContextHandler = new GlobalContextHandler(client, new CreateIndexStep(clusterService, client)); + + return ImmutableList.of(workflowStepFactory, workflowProcessSorter, globalContextHandler); } @Override @@ -106,10 +112,9 @@ public List> getExecutorBuilders(Settings settings) { FixedExecutorBuilder provisionThreadPool = new FixedExecutorBuilder( settings, PROVISION_THREAD_POOL, - 1, + OpenSearchExecutors.allocatedProcessors(settings), 10, - FLOW_FRAMEWORK_THREAD_POOL_PREFIX + PROVISION_THREAD_POOL, - false + FLOW_FRAMEWORK_THREAD_POOL_PREFIX + PROVISION_THREAD_POOL ); return ImmutableList.of(provisionThreadPool); } diff --git a/src/main/java/org/opensearch/flowframework/common/CommonValue.java b/src/main/java/org/opensearch/flowframework/common/CommonValue.java index 0aa231ca2..da9aa3d2f 100644 --- a/src/main/java/org/opensearch/flowframework/common/CommonValue.java +++ b/src/main/java/org/opensearch/flowframework/common/CommonValue.java @@ -24,10 +24,14 @@ public class CommonValue { public static final String AI_FLOW_FRAMEWORK_BASE_URI = "/_plugins/_flow_framework"; /** The URI for this plugin's workflow rest actions */ public static final String WORKFLOWS_URI = AI_FLOW_FRAMEWORK_BASE_URI + "/workflows"; + /** Field name for workflow Id, the document Id of the indexed use case template */ + public static final String WORKFLOW_ID = "workflow_id"; + /** The field name for provision workflow within a use case template*/ + public static final String PROVISION_WORKFLOW = "provision"; + /** Flow Framework plugin thread pool name prefix */ public static final String FLOW_FRAMEWORK_THREAD_POOL_PREFIX = "thread_pool.flow_framework."; /** The provision workflow thread pool name */ public static final String PROVISION_THREAD_POOL = "opensearch_workflow_provision"; - /** Field name for workflow Id, the document Id of the indexed use case template */ - public static final String WORKFLOW_ID = "workflow_id"; + } diff --git a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java index 994cdaeda..6d6565eb5 100644 --- a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java +++ b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java @@ -27,6 +27,7 @@ import java.io.IOException; import java.util.HashMap; +import java.util.Locale; import java.util.Map; import static org.opensearch.core.rest.RestStatus.INTERNAL_SERVER_ERROR; @@ -94,6 +95,37 @@ public void putTemplateToGlobalContext(Template template, ActionListener listener) { + if (!createIndexStep.doesIndexExist(GLOBAL_CONTEXT_INDEX)) { + String exceptionMessage = String.format( + Locale.ROOT, + "Failed to update template {}, global_context index does not exist.", + documentId + ); + logger.error(exceptionMessage); + listener.onFailure(new Exception(exceptionMessage)); + } else { + IndexRequest request = new IndexRequest(GLOBAL_CONTEXT_INDEX).id(documentId); + try ( + XContentBuilder builder = XContentFactory.jsonBuilder(); + ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext() + ) { + request.source(template.toXContent(builder, ToXContent.EMPTY_PARAMS)) + .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE); + client.index(request, ActionListener.runBefore(listener, () -> context.restore())); + } catch (Exception e) { + logger.error("Failed to update global_context entry : {}. {}", documentId, e.getMessage()); + listener.onFailure(e); + } + } + } + /** * Update global context index for specific fields * @param documentId global context index document id diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java index 2c2b532b4..0c77d654f 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java @@ -10,7 +10,6 @@ import com.google.common.collect.ImmutableList; import org.opensearch.client.node.NodeClient; -import org.opensearch.flowframework.model.Template; import org.opensearch.flowframework.transport.ProvisionWorkflowAction; import org.opensearch.flowframework.transport.WorkflowRequest; import org.opensearch.rest.BaseRestHandler; @@ -31,12 +30,6 @@ public class RestProvisionWorkflowAction extends BaseRestHandler { private static final String PROVISION_WORKFLOW_ACTION = "provision_workflow_action"; - // TODO : move to common values class, pending implementation - /** - * Field name for workflow Id, the document Id of the indexed use case template - */ - public static final String WORKFLOW_ID = "workflow_id"; - /** * Instantiates a new RestProvisionWorkflowAction */ @@ -52,8 +45,6 @@ public String getName() { @Override public List routes() { return ImmutableList.of( - // Provision workflow from inline use case template - new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s/%s", WORKFLOWS_URI, "_provision")), // Provision workflow from indexed use case template new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s/{%s}/%s", WORKFLOWS_URI, WORKFLOW_ID, "_provision")) ); @@ -62,20 +53,19 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - String workflowId = request.param(WORKFLOW_ID); - Template template = null; - + // Validate content if (request.hasContent()) { - template = Template.parse(request.content().utf8ToString()); + throw new IOException("Invalid request format"); } - // Validate workflow request inputs - if (workflowId == null && template == null) { - throw new IOException("workflow_id and template cannot be both null"); + // Validate params + String workflowId = request.param(WORKFLOW_ID); + if (workflowId == null) { + throw new IOException("workflow_id cannot be null"); } // Create request and provision - WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template); + WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null); return channel -> client.execute(ProvisionWorkflowAction.INSTANCE, workflowRequest, new RestToXContentListener<>(channel)); } diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java index 555b8f767..51a668344 100644 --- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java @@ -12,9 +12,9 @@ import org.apache.logging.log4j.Logger; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.HandledTransportAction; -import org.opensearch.client.Client; import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; +import org.opensearch.flowframework.indices.GlobalContextHandler; import org.opensearch.tasks.Task; import org.opensearch.transport.TransportService; @@ -25,41 +25,46 @@ public class CreateWorkflowTransportAction extends HandledTransportAction listener) { - - String workflowId; - // TODO : Check if global context index exists, and if it does not then create - if (request.getWorkflowId() == null) { - // TODO : Create new entry - // TODO : Insert doc - - // TODO : get document ID - workflowId = ""; - // TODO : check if state index exists, and if it does not, then create - // TODO : insert state index doc, mapped with documentId, defaulted to NOT_STARTED + // Create new global context and state index entries + globalContextHandler.putTemplateToGlobalContext(request.getTemplate(), ActionListener.wrap(response -> { + // TODO : Check if state index exists, create if not + // TODO : Create StateIndexRequest for workflowId, default to NOT_STARTED + listener.onResponse(new WorkflowResponse(response.getId())); + }, exception -> { + logger.error("Failed to save use case template : {}", exception.getMessage()); + listener.onFailure(exception); + })); } else { - // TODO : Update existing entry - workflowId = request.getWorkflowId(); - // TODO : Update state index entry, default back to NOT_STARTED + // Update existing entry, full document replacement + globalContextHandler.updateTemplate(request.getWorkflowId(), request.getTemplate(), ActionListener.wrap(response -> { + // TODO : Create StateIndexRequest for workflowId to reset entry to NOT_STARTED + listener.onResponse(new WorkflowResponse(request.getWorkflowId())); + }, exception -> { + logger.error("Failed to updated use case template {} : {}", request.getWorkflowId(), exception.getMessage()); + listener.onFailure(exception); + })); } - - listener.onResponse(new WorkflowResponse(workflowId)); } } diff --git a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java index bf7a725a7..f81a623b7 100644 --- a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java @@ -10,10 +10,12 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.action.get.GetRequest; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.HandledTransportAction; import org.opensearch.client.Client; import org.opensearch.common.inject.Inject; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; import org.opensearch.flowframework.model.Template; import org.opensearch.flowframework.model.Workflow; @@ -31,7 +33,9 @@ import java.util.concurrent.CompletionException; import java.util.stream.Collectors; +import static org.opensearch.flowframework.common.CommonValue.GLOBAL_CONTEXT_INDEX; import static org.opensearch.flowframework.common.CommonValue.PROVISION_THREAD_POOL; +import static org.opensearch.flowframework.common.CommonValue.PROVISION_WORKFLOW; /** * Transport Action to provision a workflow from a stored use case template @@ -40,12 +44,6 @@ public class ProvisionWorkflowTransportAction extends HandledTransportAction listener) { - if (request.getWorkflowId() == null) { - // Workflow provisioning from inline template, first parse and then index the given use case template - client.execute(CreateWorkflowAction.INSTANCE, request, ActionListener.wrap(workflowResponse -> { - String workflowId = workflowResponse.getWorkflowId(); - Template template = request.getTemplate(); - - // TODO : Use node client to update state index to PROVISIONING, given workflowId - - listener.onResponse(new WorkflowResponse(workflowId)); - - // Asychronously begin provision workflow excecution - executeWorkflowAsync(workflowId, template.workflows().get(PROVISION_WORKFLOW)); + // Retrieve use case template from global context + String workflowId = request.getWorkflowId(); + GetRequest getRequest = new GetRequest(GLOBAL_CONTEXT_INDEX, workflowId); - }, exception -> { listener.onFailure(exception); })); - } else { - // Use case template has been previously saved, retrieve entry and execute - String workflowId = request.getWorkflowId(); + // Stash thread context to interact with system index + try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { + client.get(getRequest, ActionListener.wrap(response -> { + context.restore(); - // TODO : Retrieve template from global context index using node client - Template template = null; // temporary, remove later + // Parse template from document source + Template template = Template.parse(response.getSourceAsString()); - // TODO : use node client to update state index entry to PROVISIONING, given workflowId + // TODO : Update state index entry to PROVISIONING, given workflowId - listener.onResponse(new WorkflowResponse(workflowId)); - executeWorkflowAsync(workflowId, template.workflows().get(PROVISION_WORKFLOW)); + // Respond to rest action then execute provisioning workflow async + listener.onResponse(new WorkflowResponse(workflowId)); + executeWorkflowAsync(workflowId, template.workflows().get(PROVISION_WORKFLOW)); + }, exception -> { + logger.error("Failed to retrieve template from global context.", exception); + listener.onFailure(exception); + })); + } catch (Exception e) { + logger.error("Failed to retrieve template from global context.", e); + listener.onFailure(e); } } diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateIndexStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateIndexStep.java index 848f621a2..2b2f7338d 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/CreateIndexStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/CreateIndexStep.java @@ -117,6 +117,16 @@ public String getName() { return NAME; } + // TODO : Move to index management class, pending implementation + /** + * Checks if the given index exists + * @param indexName the name of the index + * @return boolean indicating the existence of an index + */ + public boolean doesIndexExist(String indexName) { + return clusterService.state().metadata().hasIndex(indexName); + } + /** * Create Index if it's absent * @param index The index that needs to be created diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java index 0d09fe239..76b2bfa6f 100644 --- a/src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java +++ b/src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java @@ -10,6 +10,7 @@ import org.opensearch.client.AdminClient; import org.opensearch.client.Client; +import org.opensearch.common.settings.Settings; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; @@ -24,6 +25,7 @@ public class FlowFrameworkPluginTests extends OpenSearchTestCase { private Client client; private ThreadPool threadPool; + private Settings settings; @Override public void setUp() throws Exception { @@ -31,6 +33,7 @@ public void setUp() throws Exception { client = mock(Client.class); when(client.admin()).thenReturn(mock(AdminClient.class)); threadPool = new TestThreadPool(FlowFrameworkPluginTests.class.getName()); + settings = Settings.EMPTY; } @Override @@ -41,10 +44,10 @@ public void tearDown() throws Exception { public void testPlugin() throws IOException { try (FlowFrameworkPlugin ffp = new FlowFrameworkPlugin()) { - assertEquals(2, ffp.createComponents(client, null, threadPool, null, null, null, null, null, null, null, null).size()); + assertEquals(3, ffp.createComponents(client, null, threadPool, null, null, null, null, null, null, null, null).size()); assertEquals(2, ffp.getRestHandlers(null, null, null, null, null, null, null).size()); assertEquals(2, ffp.getActions().size()); - assertEquals(1, ffp.getExecutorBuilders(null).size()); + assertEquals(1, ffp.getExecutorBuilders(settings).size()); } } } diff --git a/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java index f69dfe401..6e868f00d 100644 --- a/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java +++ b/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java @@ -25,25 +25,15 @@ public class RestProvisionWorkflowActionTests extends OpenSearchTestCase { - private String invalidTemplate; private RestProvisionWorkflowAction provisionWorkflowRestAction; - private String provisionInlineWorkflowPath; - private String provisionSavedWorkflowPath; + private String provisionWorkflowPath; private NodeClient nodeClient; @Override public void setUp() throws Exception { super.setUp(); - // Invalid template configuration, missing name field - this.invalidTemplate = "{\"description\":\"description\"," - + "\"use_case\":\"use case\"," - + "\"operations\":[\"operation\"]," - + "\"version\":{\"template\":\"1.0.0\",\"compatibility\":[\"2.0.0\",\"3.0.0\"]}," - + "\"user_inputs\":{\"userKey\":\"userValue\",\"userMapKey\":{\"nestedKey\":\"nestedValue\"}}," - + "\"workflows\":{\"workflow\":{\"user_params\":{\"key\":\"value\"},\"nodes\":[{\"id\":\"A\",\"type\":\"a-type\",\"inputs\":{\"foo\":\"bar\"}},{\"id\":\"B\",\"type\":\"b-type\",\"inputs\":{\"baz\":\"qux\"}}],\"edges\":[{\"source\":\"A\",\"dest\":\"B\"}]}}}"; this.provisionWorkflowRestAction = new RestProvisionWorkflowAction(); - this.provisionInlineWorkflowPath = String.format(Locale.ROOT, "%s/%s", WORKFLOWS_URI, "_provision"); - this.provisionSavedWorkflowPath = String.format(Locale.ROOT, "%s/{%s}/%s", WORKFLOWS_URI, "workflow_id", "_provision"); + this.provisionWorkflowPath = String.format(Locale.ROOT, "%s/{%s}/%s", WORKFLOWS_URI, "workflow_id", "_provision"); this.nodeClient = mock(NodeClient.class); } @@ -54,32 +44,30 @@ public void testRestProvisionWorkflowActionName() { public void testRestProvisiionWorkflowActionRoutes() { List routes = provisionWorkflowRestAction.routes(); - assertEquals(2, routes.size()); + assertEquals(1, routes.size()); assertEquals(RestRequest.Method.POST, routes.get(0).getMethod()); - assertEquals(RestRequest.Method.POST, routes.get(1).getMethod()); - assertEquals(this.provisionInlineWorkflowPath, routes.get(0).getPath()); - assertEquals(this.provisionSavedWorkflowPath, routes.get(1).getPath()); + assertEquals(this.provisionWorkflowPath, routes.get(0).getPath()); } public void testNullWorkflowIdAndTemplate() throws IOException { // Request with no content or params RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.POST) - .withPath(this.provisionInlineWorkflowPath) + .withPath(this.provisionWorkflowPath) .build(); IOException ex = expectThrows(IOException.class, () -> { provisionWorkflowRestAction.prepareRequest(request, nodeClient); }); - assertEquals("workflow_id and template cannot be both null", ex.getMessage()); + assertEquals("workflow_id cannot be null", ex.getMessage()); } - public void testInvalidProvisionInlineWorkflowRequest() throws IOException { + public void testInvalidRequestWithContent() throws IOException { RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.POST) - .withPath(this.provisionInlineWorkflowPath) - .withContent(new BytesArray(invalidTemplate), MediaTypeRegistry.JSON) + .withPath(this.provisionWorkflowPath) + .withContent(new BytesArray("request body"), MediaTypeRegistry.JSON) .build(); IOException ex = expectThrows(IOException.class, () -> { provisionWorkflowRestAction.prepareRequest(request, nodeClient); }); - assertEquals("An template object requires a name.", ex.getMessage()); + assertEquals("Invalid request format", ex.getMessage()); } } From 6c3d3db81094de2e27be4288c32ba99f24356356 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Mon, 9 Oct 2023 21:18:49 +0000 Subject: [PATCH 09/23] Addressing PR comments (Part 3) Signed-off-by: Joshua Palis --- .../flowframework/FlowFrameworkPlugin.java | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java index 810674c47..0bac15c61 100644 --- a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java +++ b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java @@ -109,14 +109,15 @@ public List getRestHandlers( @Override public List> getExecutorBuilders(Settings settings) { // TODO : Determine final size/queueSize values for the provision thread pool - FixedExecutorBuilder provisionThreadPool = new FixedExecutorBuilder( - settings, - PROVISION_THREAD_POOL, - OpenSearchExecutors.allocatedProcessors(settings), - 10, - FLOW_FRAMEWORK_THREAD_POOL_PREFIX + PROVISION_THREAD_POOL + return ImmutableList.of( + new FixedExecutorBuilder( + settings, + PROVISION_THREAD_POOL, + OpenSearchExecutors.allocatedProcessors(settings), + 10, + FLOW_FRAMEWORK_THREAD_POOL_PREFIX + PROVISION_THREAD_POOL + ) ); - return ImmutableList.of(provisionThreadPool); } } From 288a8ae69c02d61e1b6430effeb3fde3fb6aa408 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Mon, 9 Oct 2023 22:46:47 +0000 Subject: [PATCH 10/23] Removing TODOs for RestAction constructors, adding basic unit tests for added methods in CreateIndexStep, GlobalContextHandler Signed-off-by: Joshua Palis --- .../indices/GlobalContextHandler.java | 9 ++--- .../rest/RestCreateWorkflowAction.java | 4 +-- .../rest/RestProvisionWorkflowAction.java | 4 +-- .../indices/GlobalContextHandlerTests.java | 34 +++++++++++++++++++ .../workflow/CreateIndexStepTests.java | 14 ++++++++ 5 files changed, 53 insertions(+), 12 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java index 6d6565eb5..1355b2cff 100644 --- a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java +++ b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java @@ -27,7 +27,6 @@ import java.io.IOException; import java.util.HashMap; -import java.util.Locale; import java.util.Map; import static org.opensearch.core.rest.RestStatus.INTERNAL_SERVER_ERROR; @@ -103,11 +102,9 @@ public void putTemplateToGlobalContext(Template template, ActionListener listener) { if (!createIndexStep.doesIndexExist(GLOBAL_CONTEXT_INDEX)) { - String exceptionMessage = String.format( - Locale.ROOT, - "Failed to update template {}, global_context index does not exist.", - documentId - ); + String exceptionMessage = "Failed to update template for workflow_id : " + + documentId + + ", global_context index does not exist."; logger.error(exceptionMessage); listener.onFailure(new Exception(exceptionMessage)); } else { diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java index 29038405e..091ee3766 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java @@ -34,9 +34,7 @@ public class RestCreateWorkflowAction extends BaseRestHandler { /** * Intantiates a new RestCreateWorkflowAction */ - public RestCreateWorkflowAction() { - // TODO : Pass settings and cluster service to constructor and add settings update consumer for request timeout value - } + public RestCreateWorkflowAction() {} @Override public String getName() { diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java index 0c77d654f..725c94e19 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java @@ -33,9 +33,7 @@ public class RestProvisionWorkflowAction extends BaseRestHandler { /** * Instantiates a new RestProvisionWorkflowAction */ - public RestProvisionWorkflowAction() { - // TODO : Pass settings and cluster service to constructor and add settings update consumer for request timeout value - } + public RestProvisionWorkflowAction() {} @Override public String getName() { diff --git a/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java b/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java index 0380e4808..5c06c1274 100644 --- a/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java +++ b/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java @@ -109,4 +109,38 @@ public void testStoreResponseToGlobalContext() { assertEquals(GLOBAL_CONTEXT_INDEX, requestCaptor.getValue().index()); assertEquals(documentId, requestCaptor.getValue().id()); } + + public void testUpdateTemplate() throws IOException { + Template template = mock(Template.class); + ActionListener listener = mock(ActionListener.class); + when(template.toXContent(any(XContentBuilder.class), eq(ToXContent.EMPTY_PARAMS))).thenAnswer(invocation -> { + XContentBuilder builder = invocation.getArgument(0); + return builder; + }); + when(createIndexStep.doesIndexExist(any())).thenReturn(true); + + globalContextHandler.updateTemplate("1", template, null); + + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(IndexRequest.class); + verify(client, times(1)).index(requestCaptor.capture(), any()); + + assertEquals("1", requestCaptor.getValue().id()); + } + + public void testFailedUpdateTemplate() throws IOException { + Template template = mock(Template.class); + ActionListener listener = mock(ActionListener.class); + when(createIndexStep.doesIndexExist(any())).thenReturn(false); + + globalContextHandler.updateTemplate("1", template, listener); + ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); + + verify(listener, times(1)).onFailure(exceptionCaptor.capture()); + + assertEquals( + "Failed to update template for workflow_id : 1, global_context index does not exist.", + exceptionCaptor.getValue().getMessage() + ); + + } } diff --git a/src/test/java/org/opensearch/flowframework/workflow/CreateIndexStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/CreateIndexStepTests.java index 72371095c..171ab8ae8 100644 --- a/src/test/java/org/opensearch/flowframework/workflow/CreateIndexStepTests.java +++ b/src/test/java/org/opensearch/flowframework/workflow/CreateIndexStepTests.java @@ -187,4 +187,18 @@ public void testInitIndexIfAbsent_IndexExist_returnFalse() { createIndexStep.initIndexIfAbsent(index, listener); assertTrue(indexMappingUpdated.get(index.getIndexName()).get()); } + + public void testDoesIndexExist() { + ClusterState mockClusterState = mock(ClusterState.class); + Metadata mockMetaData = mock(Metadata.class); + when(clusterService.state()).thenReturn(mockClusterState); + when(mockClusterState.metadata()).thenReturn(mockMetaData); + + createIndexStep.doesIndexExist(GLOBAL_CONTEXT_INDEX); + + ArgumentCaptor indexExistsCaptor = ArgumentCaptor.forClass(String.class); + verify(mockMetaData, times(1)).hasIndex(indexExistsCaptor.capture()); + + assertEquals(GLOBAL_CONTEXT_INDEX, indexExistsCaptor.getValue()); + } } From e200d9ac048c7d6f1bc0de8e921eb632ecc690b5 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Mon, 9 Oct 2023 23:38:29 +0000 Subject: [PATCH 11/23] Adding CreateWorkflowTransportAction unit tests Signed-off-by: Joshua Palis --- .../CreateWorkflowTransportAction.java | 2 +- .../CreateWorkflowTransportActionTests.java | 145 ++++++++++++++++++ 2 files changed, 146 insertions(+), 1 deletion(-) create mode 100644 src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java index 51a668344..eace47037 100644 --- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java @@ -59,7 +59,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener { // TODO : Create StateIndexRequest for workflowId to reset entry to NOT_STARTED - listener.onResponse(new WorkflowResponse(request.getWorkflowId())); + listener.onResponse(new WorkflowResponse(response.getId())); }, exception -> { logger.error("Failed to updated use case template {} : {}", request.getWorkflowId(), exception.getMessage()); listener.onFailure(exception); diff --git a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java new file mode 100644 index 000000000..1df6ef23e --- /dev/null +++ b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java @@ -0,0 +1,145 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.transport; + +import org.opensearch.Version; +import org.opensearch.action.index.IndexResponse; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.flowframework.indices.GlobalContextHandler; +import org.opensearch.flowframework.model.Template; +import org.opensearch.flowframework.model.Workflow; +import org.opensearch.flowframework.model.WorkflowEdge; +import org.opensearch.flowframework.model.WorkflowNode; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.transport.TransportService; + +import java.util.List; +import java.util.Map; + +import org.mockito.ArgumentCaptor; + +import static org.opensearch.flowframework.common.CommonValue.GLOBAL_CONTEXT_INDEX; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +public class CreateWorkflowTransportActionTests extends OpenSearchTestCase { + + private CreateWorkflowTransportAction createWorkflowTransportAction; + private GlobalContextHandler globalContextHandler; + private Template template; + + @Override + public void setUp() throws Exception { + super.setUp(); + this.globalContextHandler = mock(GlobalContextHandler.class); + this.createWorkflowTransportAction = new CreateWorkflowTransportAction( + mock(TransportService.class), + mock(ActionFilters.class), + globalContextHandler + ); + + List operations = List.of("operation"); + Version templateVersion = Version.fromString("1.0.0"); + List compatibilityVersions = List.of(Version.fromString("2.0.0"), Version.fromString("3.0.0")); + WorkflowNode nodeA = new WorkflowNode("A", "a-type", Map.of("foo", "bar")); + WorkflowNode nodeB = new WorkflowNode("B", "b-type", Map.of("baz", "qux")); + WorkflowEdge edgeAB = new WorkflowEdge("A", "B"); + List nodes = List.of(nodeA, nodeB); + List edges = List.of(edgeAB); + Workflow workflow = new Workflow(Map.of("key", "value"), nodes, edges); + + this.template = new Template( + "test", + "description", + "use case", + operations, + templateVersion, + compatibilityVersions, + Map.ofEntries(Map.entry("userKey", "userValue"), Map.entry("userMapKey", Map.of("nestedKey", "nestedValue"))), + Map.of("workflow", workflow), + Map.of("outputKey", "outputValue"), + Map.of("resourceKey", "resourceValue") + ); + } + + public void testCreateNewWorkflow() { + + ActionListener listener = mock(ActionListener.class); + WorkflowRequest createNewWorkflow = new WorkflowRequest(null, template); + + doAnswer(invocation -> { + ActionListener responseListener = invocation.getArgument(1); + responseListener.onResponse(new IndexResponse(new ShardId(GLOBAL_CONTEXT_INDEX, "", 1), "1", 1L, 1L, 1L, true)); + return null; + }).when(globalContextHandler).putTemplateToGlobalContext(any(), any()); + + createWorkflowTransportAction.doExecute(mock(Task.class), createNewWorkflow, listener); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(WorkflowResponse.class); + verify(listener, times(1)).onResponse(responseCaptor.capture()); + + assertEquals("1", responseCaptor.getValue().getWorkflowId()); + + } + + public void testFailedToCreateNewWorkflow() { + ActionListener listener = mock(ActionListener.class); + WorkflowRequest createNewWorkflow = new WorkflowRequest(null, template); + + doAnswer(invocation -> { + ActionListener responseListener = invocation.getArgument(1); + responseListener.onFailure(new Exception("Failed to create global_context index")); + return null; + }).when(globalContextHandler).putTemplateToGlobalContext(any(), any()); + + createWorkflowTransportAction.doExecute(mock(Task.class), createNewWorkflow, listener); + ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); + verify(listener, times(1)).onFailure(exceptionCaptor.capture()); + assertEquals("Failed to create global_context index", exceptionCaptor.getValue().getMessage()); + } + + public void testUpdateWorkflow() { + + ActionListener listener = mock(ActionListener.class); + WorkflowRequest updateWorkflow = new WorkflowRequest("1", template); + + doAnswer(invocation -> { + ActionListener responseListener = invocation.getArgument(2); + responseListener.onResponse(new IndexResponse(new ShardId(GLOBAL_CONTEXT_INDEX, "", 1), "1", 1L, 1L, 1L, true)); + return null; + }).when(globalContextHandler).updateTemplate(any(), any(), any()); + + createWorkflowTransportAction.doExecute(mock(Task.class), updateWorkflow, listener); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(WorkflowResponse.class); + verify(listener, times(1)).onResponse(responseCaptor.capture()); + + assertEquals("1", responseCaptor.getValue().getWorkflowId()); + } + + public void testFailedToUpdateWorkflow() { + ActionListener listener = mock(ActionListener.class); + WorkflowRequest updateWorkflow = new WorkflowRequest("1", template); + + doAnswer(invocation -> { + ActionListener responseListener = invocation.getArgument(2); + responseListener.onFailure(new Exception("Failed to update use case template")); + return null; + }).when(globalContextHandler).updateTemplate(any(), any(), any()); + + createWorkflowTransportAction.doExecute(mock(Task.class), updateWorkflow, listener); + ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); + verify(listener, times(1)).onFailure(exceptionCaptor.capture()); + assertEquals("Failed to update use case template", exceptionCaptor.getValue().getMessage()); + } +} From 956e8233d571dc2313c26594b1bb34cb3a70e694 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Tue, 10 Oct 2023 00:13:47 +0000 Subject: [PATCH 12/23] Adding intial failure test case for the ProvisionWorkflowTransportAction. Still need to add tests for success case Signed-off-by: Joshua Palis --- ...ProvisionWorkflowTransportActionTests.java | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java diff --git a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java new file mode 100644 index 000000000..e2ae8208e --- /dev/null +++ b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java @@ -0,0 +1,81 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.transport; + +import org.opensearch.action.get.GetResponse; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.client.Client; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.action.ActionListener; +import org.opensearch.flowframework.workflow.WorkflowProcessSorter; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import org.mockito.ArgumentCaptor; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class ProvisionWorkflowTransportActionTests extends OpenSearchTestCase { + + private ThreadPool threadPool; + private Client client; + private WorkflowProcessSorter workflowProcessSorter; + private ProvisionWorkflowTransportAction provisionWorkflowTransportAction; + + @Override + public void setUp() throws Exception { + super.setUp(); + this.threadPool = mock(ThreadPool.class); + this.client = mock(Client.class); + this.workflowProcessSorter = mock(WorkflowProcessSorter.class); + + this.provisionWorkflowTransportAction = new ProvisionWorkflowTransportAction( + mock(TransportService.class), + mock(ActionFilters.class), + threadPool, + client, + workflowProcessSorter + ); + + ThreadPool clientThreadPool = mock(ThreadPool.class); + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + + when(client.threadPool()).thenReturn(clientThreadPool); + when(clientThreadPool.getThreadContext()).thenReturn(threadContext); + } + + public void testProvisionWorkflow() { + // TODO : Add tests for success case + } + + public void testFailedToRetrieveTemplateFromGlobalContext() { + ActionListener listener = mock(ActionListener.class); + WorkflowRequest request = new WorkflowRequest("1", null); + doAnswer(invocation -> { + ActionListener responseListener = invocation.getArgument(1); + responseListener.onFailure(new Exception("Failed to retrieve template from global context.")); + return null; + }).when(client).get(any(), any()); + + provisionWorkflowTransportAction.doExecute(mock(Task.class), request, listener); + ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); + + verify(listener, times(1)).onFailure(exceptionCaptor.capture()); + assertEquals("Failed to retrieve template from global context.", exceptionCaptor.getValue().getMessage()); + } + +} From 92369660e606efe4e6bedee05183df84e96caf13 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Tue, 10 Oct 2023 17:18:03 +0000 Subject: [PATCH 13/23] Updating base URI namespace to workflow instead of workflows Signed-off-by: Joshua Palis --- .../org/opensearch/flowframework/common/CommonValue.java | 2 +- .../flowframework/rest/RestCreateWorkflowAction.java | 6 +++--- .../flowframework/rest/RestProvisionWorkflowAction.java | 4 ++-- .../flowframework/rest/RestCreateWorkflowActionTests.java | 6 +++--- .../rest/RestProvisionWorkflowActionTests.java | 4 ++-- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/common/CommonValue.java b/src/main/java/org/opensearch/flowframework/common/CommonValue.java index da9aa3d2f..e478d1542 100644 --- a/src/main/java/org/opensearch/flowframework/common/CommonValue.java +++ b/src/main/java/org/opensearch/flowframework/common/CommonValue.java @@ -23,7 +23,7 @@ public class CommonValue { /** The base URI for this plugin's rest actions */ public static final String AI_FLOW_FRAMEWORK_BASE_URI = "/_plugins/_flow_framework"; /** The URI for this plugin's workflow rest actions */ - public static final String WORKFLOWS_URI = AI_FLOW_FRAMEWORK_BASE_URI + "/workflows"; + public static final String WORKFLOW_URI = AI_FLOW_FRAMEWORK_BASE_URI + "/workflow"; /** Field name for workflow Id, the document Id of the indexed use case template */ public static final String WORKFLOW_ID = "workflow_id"; /** The field name for provision workflow within a use case template*/ diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java index 091ee3766..ace440f75 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java @@ -21,8 +21,8 @@ import java.util.List; import java.util.Locale; -import static org.opensearch.flowframework.common.CommonValue.WORKFLOWS_URI; import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI; /** * Rest Action to facilitate requests to create and update a use case template @@ -45,9 +45,9 @@ public String getName() { public List routes() { return ImmutableList.of( // Create new workflow - new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s", WORKFLOWS_URI)), + new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s", WORKFLOW_URI)), // Update use case template - new Route(RestRequest.Method.PUT, String.format(Locale.ROOT, "%s/{%s}", WORKFLOWS_URI, WORKFLOW_ID)) + new Route(RestRequest.Method.PUT, String.format(Locale.ROOT, "%s/{%s}", WORKFLOW_URI, WORKFLOW_ID)) ); } diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java index 725c94e19..180248df8 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java @@ -20,8 +20,8 @@ import java.util.List; import java.util.Locale; -import static org.opensearch.flowframework.common.CommonValue.WORKFLOWS_URI; import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI; /** * Rest action to facilitate requests to provision a workflow from an inline defined or stored use case template @@ -44,7 +44,7 @@ public String getName() { public List routes() { return ImmutableList.of( // Provision workflow from indexed use case template - new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s/{%s}/%s", WORKFLOWS_URI, WORKFLOW_ID, "_provision")) + new Route(RestRequest.Method.POST, String.format(Locale.ROOT, "%s/{%s}/%s", WORKFLOW_URI, WORKFLOW_ID, "_provision")) ); } diff --git a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java index 057c99b35..ddab7f177 100644 --- a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java +++ b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java @@ -20,7 +20,7 @@ import java.util.List; import java.util.Locale; -import static org.opensearch.flowframework.common.CommonValue.WORKFLOWS_URI; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI; import static org.mockito.Mockito.mock; public class RestCreateWorkflowActionTests extends OpenSearchTestCase { @@ -42,8 +42,8 @@ public void setUp() throws Exception { + "\"user_inputs\":{\"userKey\":\"userValue\",\"userMapKey\":{\"nestedKey\":\"nestedValue\"}}," + "\"workflows\":{\"workflow\":{\"user_params\":{\"key\":\"value\"},\"nodes\":[{\"id\":\"A\",\"type\":\"a-type\",\"inputs\":{\"foo\":\"bar\"}},{\"id\":\"B\",\"type\":\"b-type\",\"inputs\":{\"baz\":\"qux\"}}],\"edges\":[{\"source\":\"A\",\"dest\":\"B\"}]}}}"; this.createWorkflowRestAction = new RestCreateWorkflowAction(); - this.createWorkflowPath = String.format(Locale.ROOT, "%s", WORKFLOWS_URI); - this.updateWorkflowPath = String.format(Locale.ROOT, "%s/{%s}", WORKFLOWS_URI, "workflow_id"); + this.createWorkflowPath = String.format(Locale.ROOT, "%s", WORKFLOW_URI); + this.updateWorkflowPath = String.format(Locale.ROOT, "%s/{%s}", WORKFLOW_URI, "workflow_id"); this.nodeClient = mock(NodeClient.class); } diff --git a/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java index 6e868f00d..267b4f6ce 100644 --- a/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java +++ b/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java @@ -20,7 +20,7 @@ import java.util.List; import java.util.Locale; -import static org.opensearch.flowframework.common.CommonValue.WORKFLOWS_URI; +import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI; import static org.mockito.Mockito.mock; public class RestProvisionWorkflowActionTests extends OpenSearchTestCase { @@ -33,7 +33,7 @@ public class RestProvisionWorkflowActionTests extends OpenSearchTestCase { public void setUp() throws Exception { super.setUp(); this.provisionWorkflowRestAction = new RestProvisionWorkflowAction(); - this.provisionWorkflowPath = String.format(Locale.ROOT, "%s/{%s}/%s", WORKFLOWS_URI, "workflow_id", "_provision"); + this.provisionWorkflowPath = String.format(Locale.ROOT, "%s/{%s}/%s", WORKFLOW_URI, "workflow_id", "_provision"); this.nodeClient = mock(NodeClient.class); } From 0ac787359969ea29897f57f4693620a7dc3b1063 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Tue, 10 Oct 2023 17:28:09 +0000 Subject: [PATCH 14/23] Addressing PR comment, updating invalid template config test, removing field via string replacement Signed-off-by: Joshua Palis --- .../rest/RestCreateWorkflowActionTests.java | 41 +++++++++++++++---- 1 file changed, 33 insertions(+), 8 deletions(-) diff --git a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java index ddab7f177..fd07a91eb 100644 --- a/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java +++ b/src/test/java/org/opensearch/flowframework/rest/RestCreateWorkflowActionTests.java @@ -8,9 +8,14 @@ */ package org.opensearch.flowframework.rest; +import org.opensearch.Version; import org.opensearch.client.node.NodeClient; import org.opensearch.core.common.bytes.BytesArray; import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.flowframework.model.Template; +import org.opensearch.flowframework.model.Workflow; +import org.opensearch.flowframework.model.WorkflowEdge; +import org.opensearch.flowframework.model.WorkflowNode; import org.opensearch.rest.RestHandler.Route; import org.opensearch.rest.RestRequest; import org.opensearch.test.OpenSearchTestCase; @@ -19,6 +24,7 @@ import java.io.IOException; import java.util.List; import java.util.Locale; +import java.util.Map; import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI; import static org.mockito.Mockito.mock; @@ -34,13 +40,32 @@ public class RestCreateWorkflowActionTests extends OpenSearchTestCase { @Override public void setUp() throws Exception { super.setUp(); - // Invalid template configuration, missing name field - this.invalidTemplate = "{\"description\":\"description\"," - + "\"use_case\":\"use case\"," - + "\"operations\":[\"operation\"]," - + "\"version\":{\"template\":\"1.0.0\",\"compatibility\":[\"2.0.0\",\"3.0.0\"]}," - + "\"user_inputs\":{\"userKey\":\"userValue\",\"userMapKey\":{\"nestedKey\":\"nestedValue\"}}," - + "\"workflows\":{\"workflow\":{\"user_params\":{\"key\":\"value\"},\"nodes\":[{\"id\":\"A\",\"type\":\"a-type\",\"inputs\":{\"foo\":\"bar\"}},{\"id\":\"B\",\"type\":\"b-type\",\"inputs\":{\"baz\":\"qux\"}}],\"edges\":[{\"source\":\"A\",\"dest\":\"B\"}]}}}"; + + List operations = List.of("operation"); + Version templateVersion = Version.fromString("1.0.0"); + List compatibilityVersions = List.of(Version.fromString("2.0.0"), Version.fromString("3.0.0")); + WorkflowNode nodeA = new WorkflowNode("A", "a-type", Map.of("foo", "bar")); + WorkflowNode nodeB = new WorkflowNode("B", "b-type", Map.of("baz", "qux")); + WorkflowEdge edgeAB = new WorkflowEdge("A", "B"); + List nodes = List.of(nodeA, nodeB); + List edges = List.of(edgeAB); + Workflow workflow = new Workflow(Map.of("key", "value"), nodes, edges); + + Template template = new Template( + "test", + "description", + "use case", + operations, + templateVersion, + compatibilityVersions, + Map.ofEntries(Map.entry("userKey", "userValue"), Map.entry("userMapKey", Map.of("nestedKey", "nestedValue"))), + Map.of("workflow", workflow), + Map.of("outputKey", "outputValue"), + Map.of("resourceKey", "resourceValue") + ); + + // Invalid template configuration, wrong field name + this.invalidTemplate = template.toJson().replace("use_case", "invalid"); this.createWorkflowRestAction = new RestCreateWorkflowAction(); this.createWorkflowPath = String.format(Locale.ROOT, "%s", WORKFLOW_URI); this.updateWorkflowPath = String.format(Locale.ROOT, "%s/{%s}", WORKFLOW_URI, "workflow_id"); @@ -69,6 +94,6 @@ public void testInvalidCreateWorkflowRequest() throws IOException { .build(); IOException ex = expectThrows(IOException.class, () -> { createWorkflowRestAction.prepareRequest(request, nodeClient); }); - assertEquals("An template object requires a name.", ex.getMessage()); + assertEquals("Unable to parse field [invalid] in a template object.", ex.getMessage()); } } From b945bdd1a744d9b4fdd176f61f49fcd86c67d8dd Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Tue, 10 Oct 2023 18:34:34 +0000 Subject: [PATCH 15/23] Add success test case for ProvisionWorkflowTransportAction Signed-off-by: Joshua Palis --- ...ProvisionWorkflowTransportActionTests.java | 59 ++++++++++++++++++- 1 file changed, 58 insertions(+), 1 deletion(-) diff --git a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java index e2ae8208e..d190d3dc0 100644 --- a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java @@ -8,20 +8,34 @@ */ package org.opensearch.flowframework.transport; +import org.opensearch.Version; import org.opensearch.action.get.GetResponse; import org.opensearch.action.support.ActionFilters; import org.opensearch.client.Client; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.flowframework.model.Template; +import org.opensearch.flowframework.model.Workflow; +import org.opensearch.flowframework.model.WorkflowEdge; +import org.opensearch.flowframework.model.WorkflowNode; import org.opensearch.flowframework.workflow.WorkflowProcessSorter; +import org.opensearch.index.get.GetResult; import org.opensearch.tasks.Task; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; +import java.util.List; +import java.util.Map; + import org.mockito.ArgumentCaptor; +import static org.opensearch.flowframework.common.CommonValue.GLOBAL_CONTEXT_INDEX; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; @@ -35,6 +49,7 @@ public class ProvisionWorkflowTransportActionTests extends OpenSearchTestCase { private Client client; private WorkflowProcessSorter workflowProcessSorter; private ProvisionWorkflowTransportAction provisionWorkflowTransportAction; + private Template template; @Override public void setUp() throws Exception { @@ -51,6 +66,29 @@ public void setUp() throws Exception { workflowProcessSorter ); + List operations = List.of("operation"); + Version templateVersion = Version.fromString("1.0.0"); + List compatibilityVersions = List.of(Version.fromString("2.0.0"), Version.fromString("3.0.0")); + WorkflowNode nodeA = new WorkflowNode("A", "a-type", Map.of("foo", "bar")); + WorkflowNode nodeB = new WorkflowNode("B", "b-type", Map.of("baz", "qux")); + WorkflowEdge edgeAB = new WorkflowEdge("A", "B"); + List nodes = List.of(nodeA, nodeB); + List edges = List.of(edgeAB); + Workflow workflow = new Workflow(Map.of("key", "value"), nodes, edges); + + this.template = new Template( + "test", + "description", + "use case", + operations, + templateVersion, + compatibilityVersions, + Map.ofEntries(Map.entry("userKey", "userValue"), Map.entry("userMapKey", Map.of("nestedKey", "nestedValue"))), + Map.of("provision", workflow), + Map.of("outputKey", "outputValue"), + Map.of("resourceKey", "resourceValue") + ); + ThreadPool clientThreadPool = mock(ThreadPool.class); ThreadContext threadContext = new ThreadContext(Settings.EMPTY); @@ -59,7 +97,26 @@ public void setUp() throws Exception { } public void testProvisionWorkflow() { - // TODO : Add tests for success case + + String workflowId = "1"; + ActionListener listener = mock(ActionListener.class); + WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null); + + doAnswer(invocation -> { + ActionListener responseListener = invocation.getArgument(1); + + XContentBuilder builder = XContentFactory.jsonBuilder(); + this.template.toXContent(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)); + return null; + }).when(client).get(any(), any()); + + provisionWorkflowTransportAction.doExecute(mock(Task.class), workflowRequest, listener); + ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(WorkflowResponse.class); + verify(listener, times(1)).onResponse(responseCaptor.capture()); + assertEquals(workflowId, responseCaptor.getValue().getWorkflowId()); } public void testFailedToRetrieveTemplateFromGlobalContext() { From 1dac1eeb2a28a8479dd9f3de14b662167aa27fda Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Wed, 11 Oct 2023 19:44:25 +0000 Subject: [PATCH 16/23] Updating global context index mapping for template version and compatibility version fields from int to text Signed-off-by: Joshua Palis --- src/main/resources/mappings/global-context.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/main/resources/mappings/global-context.json b/src/main/resources/mappings/global-context.json index 86e952942..cd3ce4e6b 100644 --- a/src/main/resources/mappings/global-context.json +++ b/src/main/resources/mappings/global-context.json @@ -29,10 +29,10 @@ "type": "nested", "properties": { "template": { - "type": "integer" + "type": "text" }, "compatibility": { - "type": "integer" + "type": "text" } } }, From afeb2b68aa6d0531ff685078352aed1279b9a980 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 12 Oct 2023 00:52:07 +0000 Subject: [PATCH 17/23] Fixing bugs, changed GC index mapping so that template/compatibility versions are of type text, added GC template document readers/writers, modified tests. Still need to add test cases for the new readers/writers Signed-off-by: Joshua Palis --- .../indices/GlobalContextHandler.java | 2 +- .../flowframework/model/Template.java | 239 ++++++++++++++++++ .../ProvisionWorkflowTransportAction.java | 2 +- .../indices/GlobalContextHandlerTests.java | 2 +- ...ProvisionWorkflowTransportActionTests.java | 2 +- 5 files changed, 243 insertions(+), 4 deletions(-) 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)); From bbe8effd7ab472c32445331c3b0a7a8a5f6303b1 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 12 Oct 2023 00:55:31 +0000 Subject: [PATCH 18/23] Updating GlobalContextHandler.updateTemplate() to use toDocumentSource instead of toXContent() Signed-off-by: Joshua Palis --- .../opensearch/flowframework/indices/GlobalContextHandler.java | 2 +- .../flowframework/indices/GlobalContextHandlerTests.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java index 14d94fd8c..49f0888e6 100644 --- a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java +++ b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java @@ -113,7 +113,7 @@ public void updateTemplate(String documentId, Template template, ActionListener< XContentBuilder builder = XContentFactory.jsonBuilder(); ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext() ) { - request.source(template.toXContent(builder, ToXContent.EMPTY_PARAMS)) + request.source(template.toDocumentSource(builder, ToXContent.EMPTY_PARAMS)) .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE); client.index(request, ActionListener.runBefore(listener, () -> context.restore())); } catch (Exception e) { diff --git a/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java b/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java index 9d6a1544e..a09f27b99 100644 --- a/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java +++ b/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java @@ -113,7 +113,7 @@ public void testStoreResponseToGlobalContext() { public void testUpdateTemplate() throws IOException { Template template = mock(Template.class); ActionListener listener = mock(ActionListener.class); - when(template.toXContent(any(XContentBuilder.class), eq(ToXContent.EMPTY_PARAMS))).thenAnswer(invocation -> { + when(template.toDocumentSource(any(XContentBuilder.class), eq(ToXContent.EMPTY_PARAMS))).thenAnswer(invocation -> { XContentBuilder builder = invocation.getArgument(0); return builder; }); From b910ebc85bb251d37c961afbaa21fcdebfdaab42 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 12 Oct 2023 16:26:55 +0000 Subject: [PATCH 19/23] Replacing exceptions with FlowFrameworException Signed-off-by: Joshua Palis --- .../rest/RestProvisionWorkflowAction.java | 6 ++++-- .../transport/CreateWorkflowTransportAction.java | 6 ++++-- .../transport/ProvisionWorkflowTransportAction.java | 10 ++++++---- .../rest/RestProvisionWorkflowActionTests.java | 12 ++++++++++-- 4 files changed, 24 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 180248df8..89471ee00 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java @@ -10,6 +10,8 @@ import com.google.common.collect.ImmutableList; import org.opensearch.client.node.NodeClient; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.flowframework.transport.ProvisionWorkflowAction; import org.opensearch.flowframework.transport.WorkflowRequest; import org.opensearch.rest.BaseRestHandler; @@ -53,13 +55,13 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli // Validate content if (request.hasContent()) { - throw new IOException("Invalid request format"); + throw new FlowFrameworkException("Invalid request format", RestStatus.BAD_REQUEST); } // Validate params String workflowId = request.param(WORKFLOW_ID); if (workflowId == null) { - throw new IOException("workflow_id cannot be null"); + throw new FlowFrameworkException("workflow_id cannot be null", RestStatus.BAD_REQUEST); } // Create request and provision diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java index eace47037..abc704064 100644 --- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java @@ -14,6 +14,8 @@ import org.opensearch.action.support.HandledTransportAction; import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.flowframework.indices.GlobalContextHandler; import org.opensearch.tasks.Task; import org.opensearch.transport.TransportService; @@ -53,7 +55,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener { logger.error("Failed to save use case template : {}", exception.getMessage()); - listener.onFailure(exception); + listener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.INTERNAL_SERVER_ERROR)); })); } else { // Update existing entry, full document replacement @@ -62,7 +64,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener { logger.error("Failed to updated use case template {} : {}", request.getWorkflowId(), exception.getMessage()); - listener.onFailure(exception); + listener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.INTERNAL_SERVER_ERROR)); })); } } diff --git a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java index be4f2ed2d..0dbec5bf2 100644 --- a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java @@ -17,6 +17,8 @@ import org.opensearch.common.inject.Inject; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.flowframework.model.Template; import org.opensearch.flowframework.model.Workflow; import org.opensearch.flowframework.workflow.ProcessNode; @@ -92,11 +94,11 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener { logger.error("Failed to retrieve template from global context.", exception); - listener.onFailure(exception); + listener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.INTERNAL_SERVER_ERROR)); })); } catch (Exception e) { logger.error("Failed to retrieve template from global context.", e); - listener.onFailure(e); + listener.onFailure(new FlowFrameworkException(e.getMessage(), RestStatus.INTERNAL_SERVER_ERROR)); } } @@ -119,7 +121,7 @@ private void executeWorkflowAsync(String workflowId, Workflow workflow) { try { threadPool.executor(PROVISION_THREAD_POOL).execute(() -> { executeWorkflow(workflow, provisionWorkflowListener); }); } catch (Exception exception) { - provisionWorkflowListener.onFailure(exception); + provisionWorkflowListener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.INTERNAL_SERVER_ERROR)); } } @@ -157,7 +159,7 @@ private void executeWorkflow(Workflow workflow, ActionListener workflowL // TODO : Create State Index request with provisioning state, start time, end time, etc, pending implementation. String for now workflowListener.onResponse("READY"); } catch (CancellationException | CompletionException ex) { - workflowListener.onFailure(ex); + workflowListener.onFailure(new FlowFrameworkException(ex.getMessage(), RestStatus.INTERNAL_SERVER_ERROR)); } } diff --git a/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java b/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java index 267b4f6ce..a44817cec 100644 --- a/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java +++ b/src/test/java/org/opensearch/flowframework/rest/RestProvisionWorkflowActionTests.java @@ -10,7 +10,9 @@ import org.opensearch.client.node.NodeClient; import org.opensearch.core.common.bytes.BytesArray; +import org.opensearch.core.rest.RestStatus; import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.rest.RestHandler.Route; import org.opensearch.rest.RestRequest; import org.opensearch.test.OpenSearchTestCase; @@ -56,8 +58,11 @@ public void testNullWorkflowIdAndTemplate() throws IOException { .withPath(this.provisionWorkflowPath) .build(); - IOException ex = expectThrows(IOException.class, () -> { provisionWorkflowRestAction.prepareRequest(request, nodeClient); }); + FlowFrameworkException ex = expectThrows(FlowFrameworkException.class, () -> { + provisionWorkflowRestAction.prepareRequest(request, nodeClient); + }); assertEquals("workflow_id cannot be null", ex.getMessage()); + assertEquals(RestStatus.BAD_REQUEST, ex.getRestStatus()); } public void testInvalidRequestWithContent() throws IOException { @@ -66,8 +71,11 @@ public void testInvalidRequestWithContent() throws IOException { .withContent(new BytesArray("request body"), MediaTypeRegistry.JSON) .build(); - IOException ex = expectThrows(IOException.class, () -> { provisionWorkflowRestAction.prepareRequest(request, nodeClient); }); + FlowFrameworkException ex = expectThrows(FlowFrameworkException.class, () -> { + provisionWorkflowRestAction.prepareRequest(request, nodeClient); + }); assertEquals("Invalid request format", ex.getMessage()); + assertEquals(RestStatus.BAD_REQUEST, ex.getRestStatus()); } } From decb31f620e3e7fcdc3725b5ea0accd2dc26b7cf Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 12 Oct 2023 16:36:11 +0000 Subject: [PATCH 20/23] Resolving javadoc warnings Signed-off-by: Joshua Palis --- .../flowframework/common/CommonValue.java | 15 +++++++++++++++ .../exception/FlowFrameworkException.java | 1 + .../flowframework/indices/FlowFrameworkIndex.java | 15 +++++++++++++++ .../opensearch/flowframework/model/Template.java | 2 +- 4 files changed, 32 insertions(+), 1 deletion(-) diff --git a/src/main/java/org/opensearch/flowframework/common/CommonValue.java b/src/main/java/org/opensearch/flowframework/common/CommonValue.java index 83f7546b9..f6d515f43 100644 --- a/src/main/java/org/opensearch/flowframework/common/CommonValue.java +++ b/src/main/java/org/opensearch/flowframework/common/CommonValue.java @@ -13,11 +13,17 @@ */ public class CommonValue { + /** Default value for no schema version */ public static Integer NO_SCHEMA_VERSION = 0; + /** Index mapping meta field name*/ public static final String META = "_meta"; + /** Schema Version field name */ public static final String SCHEMA_VERSION_FIELD = "schema_version"; + /** Global Context Index Name */ public static final String GLOBAL_CONTEXT_INDEX = ".plugins-ai-global-context"; + /** Global Context index mapping file path */ public static final String GLOBAL_CONTEXT_INDEX_MAPPING = "mappings/global-context.json"; + /** Global Context index mapping version */ public static final Integer GLOBAL_CONTEXT_INDEX_VERSION = 1; /** The base URI for this plugin's rest actions */ @@ -34,13 +40,22 @@ public class CommonValue { /** The provision workflow thread pool name */ public static final String PROVISION_THREAD_POOL = "opensearch_workflow_provision"; + /** Model Id field */ public static final String MODEL_ID = "model_id"; + /** Function Name field */ public static final String FUNCTION_NAME = "function_name"; + /** Model Name field */ public static final String MODEL_NAME = "name"; + /** Model Version field */ public static final String MODEL_VERSION = "model_version"; + /** Model Group Id field */ public static final String MODEL_GROUP_ID = "model_group_id"; + /** Description field */ public static final String DESCRIPTION = "description"; + /** Connector Id field */ public static final String CONNECTOR_ID = "connector_id"; + /** Model format field */ public static final String MODEL_FORMAT = "model_format"; + /** Model config field */ public static final String MODEL_CONFIG = "model_config"; } diff --git a/src/main/java/org/opensearch/flowframework/exception/FlowFrameworkException.java b/src/main/java/org/opensearch/flowframework/exception/FlowFrameworkException.java index 6508fb9f7..f3cb55950 100644 --- a/src/main/java/org/opensearch/flowframework/exception/FlowFrameworkException.java +++ b/src/main/java/org/opensearch/flowframework/exception/FlowFrameworkException.java @@ -17,6 +17,7 @@ public class FlowFrameworkException extends RuntimeException { private static final long serialVersionUID = 1L; + /** The rest status code of this exception */ private final RestStatus restStatus; /** diff --git a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndex.java b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndex.java index 30261ae0e..d0ef3503c 100644 --- a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndex.java +++ b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndex.java @@ -19,6 +19,9 @@ * An enumeration of Flow Framework indices */ public enum FlowFrameworkIndex { + /** + * Global Context Index + */ GLOBAL_CONTEXT( GLOBAL_CONTEXT_INDEX, ThrowingSupplierWrapper.throwingSupplierWrapper(GlobalContextHandler::getGlobalContextMappings), @@ -35,14 +38,26 @@ public enum FlowFrameworkIndex { this.version = version; } + /** + * Retrieves the index name + * @return the index name + */ public String getIndexName() { return indexName; } + /** + * Retrieves the index mapping + * @return the index mapping + */ public String getMapping() { return mapping; } + /** + * Retrieves the index version + * @return the index version + */ public Integer getVersion() { return version; } diff --git a/src/main/java/org/opensearch/flowframework/model/Template.java b/src/main/java/org/opensearch/flowframework/model/Template.java index bcfbad22d..ece92535d 100644 --- a/src/main/java/org/opensearch/flowframework/model/Template.java +++ b/src/main/java/org/opensearch/flowframework/model/Template.java @@ -164,7 +164,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws * @param builder the XContentBuilder * @param params the params * @return the XContentBuilder - * @throws IOException + * @throws IOException if the document source fails to be generated */ public XContentBuilder toDocumentSource(XContentBuilder builder, Params params) throws IOException { XContentBuilder xContentBuilder = builder.startObject(); From 09dd4711611684f777c1ccb1d08e2693ab2aa29d Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 12 Oct 2023 18:22:21 +0000 Subject: [PATCH 21/23] Cleaning up TODOs Signed-off-by: Joshua Palis --- .../org/opensearch/flowframework/transport/WorkflowResponse.java | 1 - 1 file changed, 1 deletion(-) diff --git a/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java b/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java index 05c26d66f..20a7700a3 100644 --- a/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java +++ b/src/main/java/org/opensearch/flowframework/transport/WorkflowResponse.java @@ -59,7 +59,6 @@ public void writeTo(StreamOutput out) throws IOException { out.writeString(workflowId); } - // TODO : Replace WORKFLOW_ID after string is moved to common values class @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { return builder.startObject().field(WORKFLOW_ID, this.workflowId).endObject(); From 4d96e5014078b0313a569e24e274f62f21125aec Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 12 Oct 2023 22:14:28 +0000 Subject: [PATCH 22/23] Addressing PR comments Signed-off-by: Joshua Palis --- .../flowframework/common/CommonValue.java | 6 ++++-- .../indices/GlobalContextHandler.java | 2 +- .../transport/CreateWorkflowAction.java | 5 +++-- .../CreateWorkflowTransportAction.java | 18 +++++++++++------- .../transport/ProvisionWorkflowAction.java | 6 +++--- .../indices/GlobalContextHandlerTests.java | 10 +++++----- .../CreateWorkflowTransportActionTests.java | 8 ++++---- .../ProvisionWorkflowTransportActionTests.java | 5 +++-- .../workflow/CreateIndexStepTests.java | 2 +- 9 files changed, 35 insertions(+), 27 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/common/CommonValue.java b/src/main/java/org/opensearch/flowframework/common/CommonValue.java index f6d515f43..eb67f2dc4 100644 --- a/src/main/java/org/opensearch/flowframework/common/CommonValue.java +++ b/src/main/java/org/opensearch/flowframework/common/CommonValue.java @@ -26,10 +26,12 @@ public class CommonValue { /** Global Context index mapping version */ public static final Integer GLOBAL_CONTEXT_INDEX_VERSION = 1; + /** The transport action name prefix */ + public static final String TRANSPORT_ACION_NAME_PREFIX = "cluster:admin/opensearch/flow_framework/"; /** The base URI for this plugin's rest actions */ - public static final String AI_FLOW_FRAMEWORK_BASE_URI = "/_plugins/_flow_framework"; + public static final String FLOW_FRAMEWORK_BASE_URI = "/_plugins/_flow_framework"; /** The URI for this plugin's workflow rest actions */ - public static final String WORKFLOW_URI = AI_FLOW_FRAMEWORK_BASE_URI + "/workflow"; + public static final String WORKFLOW_URI = FLOW_FRAMEWORK_BASE_URI + "/workflow"; /** Field name for workflow Id, the document Id of the indexed use case template */ public static final String WORKFLOW_ID = "workflow_id"; /** The field name for provision workflow within a use case template*/ diff --git a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java index 49f0888e6..53037d7ce 100644 --- a/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java +++ b/src/main/java/org/opensearch/flowframework/indices/GlobalContextHandler.java @@ -100,7 +100,7 @@ public void putTemplateToGlobalContext(Template template, ActionListener listener) { + public void updateTemplateInGlobalContext(String documentId, Template template, ActionListener listener) { if (!createIndexStep.doesIndexExist(GLOBAL_CONTEXT_INDEX)) { String exceptionMessage = "Failed to update template for workflow_id : " + documentId diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowAction.java index 8e38c4189..0f49c826f 100644 --- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowAction.java @@ -10,14 +10,15 @@ import org.opensearch.action.ActionType; +import static org.opensearch.flowframework.common.CommonValue.TRANSPORT_ACION_NAME_PREFIX; + /** * External Action for public facing RestCreateWorkflowActiom */ public class CreateWorkflowAction extends ActionType { - // TODO : Determine external action prefix for plugin /** The name of this action */ - public static final String NAME = "workflows/create"; + public static final String NAME = TRANSPORT_ACION_NAME_PREFIX + "workflow/create"; /** An instance of this action */ public static final CreateWorkflowAction INSTANCE = new CreateWorkflowAction(); diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java index abc704064..f4147b144 100644 --- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java @@ -59,13 +59,17 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener { - // TODO : Create StateIndexRequest for workflowId to reset entry to NOT_STARTED - listener.onResponse(new WorkflowResponse(response.getId())); - }, exception -> { - logger.error("Failed to updated use case template {} : {}", request.getWorkflowId(), exception.getMessage()); - listener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.INTERNAL_SERVER_ERROR)); - })); + globalContextHandler.updateTemplateInGlobalContext( + request.getWorkflowId(), + request.getTemplate(), + ActionListener.wrap(response -> { + // TODO : Create StateIndexRequest for workflowId to reset entry to NOT_STARTED + listener.onResponse(new WorkflowResponse(response.getId())); + }, exception -> { + logger.error("Failed to updated use case template {} : {}", request.getWorkflowId(), exception.getMessage()); + listener.onFailure(new FlowFrameworkException(exception.getMessage(), RestStatus.INTERNAL_SERVER_ERROR)); + }) + ); } } diff --git a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowAction.java index 7167c4357..022e73488 100644 --- a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowAction.java @@ -10,14 +10,14 @@ import org.opensearch.action.ActionType; +import static org.opensearch.flowframework.common.CommonValue.TRANSPORT_ACION_NAME_PREFIX; + /** * External Action for public facing RestProvisionWorkflowAction */ public class ProvisionWorkflowAction extends ActionType { - - // TODO : Determine external action prefix for plugin /** The name of this action */ - public static final String NAME = "workflows/provision"; + public static final String NAME = TRANSPORT_ACION_NAME_PREFIX + "workflow/provision"; /** An instance of this action */ public static final ProvisionWorkflowAction INSTANCE = new ProvisionWorkflowAction(); diff --git a/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java b/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java index a09f27b99..800f1d49e 100644 --- a/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java +++ b/src/test/java/org/opensearch/flowframework/indices/GlobalContextHandlerTests.java @@ -84,7 +84,7 @@ public void testPutTemplateToGlobalContext() throws IOException { ActionListener callback = invocation.getArgument(1); callback.onResponse(true); return null; - }).when(createIndexStep).initIndexIfAbsent(any(), any()); + }).when(createIndexStep).initIndexIfAbsent(any(FlowFrameworkIndex.class), any()); globalContextHandler.putTemplateToGlobalContext(template, listener); @@ -110,7 +110,7 @@ public void testStoreResponseToGlobalContext() { assertEquals(documentId, requestCaptor.getValue().id()); } - public void testUpdateTemplate() throws IOException { + public void testUpdateTemplateInGlobalContext() throws IOException { Template template = mock(Template.class); ActionListener listener = mock(ActionListener.class); when(template.toDocumentSource(any(XContentBuilder.class), eq(ToXContent.EMPTY_PARAMS))).thenAnswer(invocation -> { @@ -119,7 +119,7 @@ public void testUpdateTemplate() throws IOException { }); when(createIndexStep.doesIndexExist(any())).thenReturn(true); - globalContextHandler.updateTemplate("1", template, null); + globalContextHandler.updateTemplateInGlobalContext("1", template, null); ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(IndexRequest.class); verify(client, times(1)).index(requestCaptor.capture(), any()); @@ -127,12 +127,12 @@ public void testUpdateTemplate() throws IOException { assertEquals("1", requestCaptor.getValue().id()); } - public void testFailedUpdateTemplate() throws IOException { + public void testFailedUpdateTemplateInGlobalContext() throws IOException { Template template = mock(Template.class); ActionListener listener = mock(ActionListener.class); when(createIndexStep.doesIndexExist(any())).thenReturn(false); - globalContextHandler.updateTemplate("1", template, listener); + globalContextHandler.updateTemplateInGlobalContext("1", template, listener); ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); verify(listener, times(1)).onFailure(exceptionCaptor.capture()); diff --git a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java index 1df6ef23e..1b937570b 100644 --- a/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/CreateWorkflowTransportActionTests.java @@ -83,7 +83,7 @@ public void testCreateNewWorkflow() { ActionListener responseListener = invocation.getArgument(1); responseListener.onResponse(new IndexResponse(new ShardId(GLOBAL_CONTEXT_INDEX, "", 1), "1", 1L, 1L, 1L, true)); return null; - }).when(globalContextHandler).putTemplateToGlobalContext(any(), any()); + }).when(globalContextHandler).putTemplateToGlobalContext(any(Template.class), any()); createWorkflowTransportAction.doExecute(mock(Task.class), createNewWorkflow, listener); ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(WorkflowResponse.class); @@ -101,7 +101,7 @@ public void testFailedToCreateNewWorkflow() { ActionListener responseListener = invocation.getArgument(1); responseListener.onFailure(new Exception("Failed to create global_context index")); return null; - }).when(globalContextHandler).putTemplateToGlobalContext(any(), any()); + }).when(globalContextHandler).putTemplateToGlobalContext(any(Template.class), any()); createWorkflowTransportAction.doExecute(mock(Task.class), createNewWorkflow, listener); ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); @@ -118,7 +118,7 @@ public void testUpdateWorkflow() { ActionListener responseListener = invocation.getArgument(2); responseListener.onResponse(new IndexResponse(new ShardId(GLOBAL_CONTEXT_INDEX, "", 1), "1", 1L, 1L, 1L, true)); return null; - }).when(globalContextHandler).updateTemplate(any(), any(), any()); + }).when(globalContextHandler).updateTemplateInGlobalContext(any(), any(Template.class), any()); createWorkflowTransportAction.doExecute(mock(Task.class), updateWorkflow, listener); ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(WorkflowResponse.class); @@ -135,7 +135,7 @@ public void testFailedToUpdateWorkflow() { ActionListener responseListener = invocation.getArgument(2); responseListener.onFailure(new Exception("Failed to update use case template")); return null; - }).when(globalContextHandler).updateTemplate(any(), any(), any()); + }).when(globalContextHandler).updateTemplateInGlobalContext(any(), any(Template.class), any()); createWorkflowTransportAction.doExecute(mock(Task.class), updateWorkflow, listener); ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); diff --git a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java index c58e810b9..7e1e13e03 100644 --- a/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java +++ b/src/test/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportActionTests.java @@ -9,6 +9,7 @@ package org.opensearch.flowframework.transport; import org.opensearch.Version; +import org.opensearch.action.get.GetRequest; import org.opensearch.action.get.GetResponse; import org.opensearch.action.support.ActionFilters; import org.opensearch.client.Client; @@ -111,7 +112,7 @@ public void testProvisionWorkflow() { GetResult getResult = new GetResult(GLOBAL_CONTEXT_INDEX, workflowId, 1, 1, 1, true, templateBytesRef, null, null); responseListener.onResponse(new GetResponse(getResult)); return null; - }).when(client).get(any(), any()); + }).when(client).get(any(GetRequest.class), any()); provisionWorkflowTransportAction.doExecute(mock(Task.class), workflowRequest, listener); ArgumentCaptor responseCaptor = ArgumentCaptor.forClass(WorkflowResponse.class); @@ -126,7 +127,7 @@ public void testFailedToRetrieveTemplateFromGlobalContext() { ActionListener responseListener = invocation.getArgument(1); responseListener.onFailure(new Exception("Failed to retrieve template from global context.")); return null; - }).when(client).get(any(), any()); + }).when(client).get(any(GetRequest.class), any()); provisionWorkflowTransportAction.doExecute(mock(Task.class), request, listener); ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); diff --git a/src/test/java/org/opensearch/flowframework/workflow/CreateIndexStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/CreateIndexStepTests.java index 171ab8ae8..a952f6fe7 100644 --- a/src/test/java/org/opensearch/flowframework/workflow/CreateIndexStepTests.java +++ b/src/test/java/org/opensearch/flowframework/workflow/CreateIndexStepTests.java @@ -128,7 +128,7 @@ public void testInitIndexIfAbsent_IndexNotPresent() { ActionListener listener = mock(ActionListener.class); createIndexStep.initIndexIfAbsent(FlowFrameworkIndex.GLOBAL_CONTEXT, listener); - verify(indicesAdminClient, times(1)).create(any(), any()); + verify(indicesAdminClient, times(1)).create(any(CreateIndexRequest.class), any()); } public void testInitIndexIfAbsent_IndexExist() { From 820044112bb4128421cf326024fe8e5dd0058694 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Thu, 12 Oct 2023 22:28:41 +0000 Subject: [PATCH 23/23] Addressing PR comments, moving some common template parsing methods to a common TemplateUtil class Signed-off-by: Joshua Palis --- .../flowframework/common/TemplateUtil.java | 79 +++++++++++++++++++ .../model/PipelineProcessor.java | 6 +- .../flowframework/model/Template.java | 52 +----------- .../flowframework/model/WorkflowNode.java | 10 ++- 4 files changed, 91 insertions(+), 56 deletions(-) create mode 100644 src/main/java/org/opensearch/flowframework/common/TemplateUtil.java diff --git a/src/main/java/org/opensearch/flowframework/common/TemplateUtil.java b/src/main/java/org/opensearch/flowframework/common/TemplateUtil.java new file mode 100644 index 000000000..80e59ce81 --- /dev/null +++ b/src/main/java/org/opensearch/flowframework/common/TemplateUtil.java @@ -0,0 +1,79 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.common; + +import org.opensearch.common.xcontent.LoggingDeprecationHandler; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.core.xcontent.XContentParser; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.Map.Entry; + +import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken; + +/** + * Utility methods for Template parsing + */ +public class TemplateUtil { + + /** + * 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; + } + + /** + * Builds an XContent object representing a map of String keys to String values. + * + * @param xContentBuilder An XContent builder whose position is at the start of the map object to build + * @param map A map as key-value String pairs. + * @throws IOException on a build failure + */ + public static void buildStringToStringMap(XContentBuilder xContentBuilder, Map map) throws IOException { + xContentBuilder.startObject(); + for (Entry e : map.entrySet()) { + xContentBuilder.field((String) e.getKey(), (String) e.getValue()); + } + xContentBuilder.endObject(); + } + + /** + * Parses an XContent object representing a map of String keys to String values. + * + * @param parser An XContent parser whose position is at the start of the map object to parse + * @return A map as identified by the key-value pairs in the XContent + * @throws IOException on a parse failure + */ + public static Map parseStringToStringMap(XContentParser parser) throws IOException { + ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); + Map map = new HashMap<>(); + while (parser.nextToken() != XContentParser.Token.END_OBJECT) { + String fieldName = parser.currentName(); + parser.nextToken(); + map.put(fieldName, parser.text()); + } + return map; + } + +} diff --git a/src/main/java/org/opensearch/flowframework/model/PipelineProcessor.java b/src/main/java/org/opensearch/flowframework/model/PipelineProcessor.java index 1407036b3..b6da0abe5 100644 --- a/src/main/java/org/opensearch/flowframework/model/PipelineProcessor.java +++ b/src/main/java/org/opensearch/flowframework/model/PipelineProcessor.java @@ -17,6 +17,8 @@ import java.util.Map; import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken; +import static org.opensearch.flowframework.common.TemplateUtil.buildStringToStringMap; +import static org.opensearch.flowframework.common.TemplateUtil.parseStringToStringMap; /** * This represents a processor associated with search and ingest pipelines in the {@link Template}. @@ -46,7 +48,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws XContentBuilder xContentBuilder = builder.startObject(); xContentBuilder.field(TYPE_FIELD, this.type); xContentBuilder.field(PARAMS_FIELD); - Template.buildStringToStringMap(xContentBuilder, this.params); + buildStringToStringMap(xContentBuilder, this.params); return xContentBuilder.endObject(); } @@ -70,7 +72,7 @@ public static PipelineProcessor parse(XContentParser parser) throws IOException type = parser.text(); break; case PARAMS_FIELD: - params = Template.parseStringToStringMap(parser); + params = parseStringToStringMap(parser); break; default: throw new IOException("Unable to parse field [" + fieldName + "] in a pipeline processor object."); diff --git a/src/main/java/org/opensearch/flowframework/model/Template.java b/src/main/java/org/opensearch/flowframework/model/Template.java index ece92535d..7d08ef240 100644 --- a/src/main/java/org/opensearch/flowframework/model/Template.java +++ b/src/main/java/org/opensearch/flowframework/model/Template.java @@ -25,6 +25,8 @@ import java.util.Map.Entry; import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken; +import static org.opensearch.flowframework.common.TemplateUtil.jsonToParser; +import static org.opensearch.flowframework.common.TemplateUtil.parseStringToStringMap; /** * The Template is the central data structure which configures workflows. This object is used to parse JSON communicated via REST API. @@ -522,23 +524,6 @@ 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 * @@ -556,39 +541,6 @@ public static Template parse(String json) throws IOException { return parse(parser); } - /** - * Builds an XContent object representing a map of String keys to String values. - * - * @param xContentBuilder An XContent builder whose position is at the start of the map object to build - * @param map A map as key-value String pairs. - * @throws IOException on a build failure - */ - public static void buildStringToStringMap(XContentBuilder xContentBuilder, Map map) throws IOException { - xContentBuilder.startObject(); - for (Entry e : map.entrySet()) { - xContentBuilder.field((String) e.getKey(), (String) e.getValue()); - } - xContentBuilder.endObject(); - } - - /** - * Parses an XContent object representing a map of String keys to String values. - * - * @param parser An XContent parser whose position is at the start of the map object to parse - * @return A map as identified by the key-value pairs in the XContent - * @throws IOException on a parse failure - */ - public static Map parseStringToStringMap(XContentParser parser) throws IOException { - ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); - Map map = new HashMap<>(); - while (parser.nextToken() != XContentParser.Token.END_OBJECT) { - String fieldName = parser.currentName(); - parser.nextToken(); - map.put(fieldName, parser.text()); - } - return map; - } - /** * Output this object in a compact JSON string. * diff --git a/src/main/java/org/opensearch/flowframework/model/WorkflowNode.java b/src/main/java/org/opensearch/flowframework/model/WorkflowNode.java index 8c4a6ae52..e34c4ddec 100644 --- a/src/main/java/org/opensearch/flowframework/model/WorkflowNode.java +++ b/src/main/java/org/opensearch/flowframework/model/WorkflowNode.java @@ -24,6 +24,8 @@ import java.util.Objects; import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken; +import static org.opensearch.flowframework.common.TemplateUtil.buildStringToStringMap; +import static org.opensearch.flowframework.common.TemplateUtil.parseStringToStringMap; /** * This represents a process node (step) in a workflow graph in the {@link Template}. @@ -75,7 +77,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws if (e.getValue() instanceof String) { xContentBuilder.value(e.getValue()); } else if (e.getValue() instanceof Map) { - Template.buildStringToStringMap(xContentBuilder, (Map) e.getValue()); + buildStringToStringMap(xContentBuilder, (Map) e.getValue()); } else if (e.getValue() instanceof Object[]) { xContentBuilder.startArray(); if (PROCESSORS_FIELD.equals(e.getKey())) { @@ -84,7 +86,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws } } else { for (Map map : (Map[]) e.getValue()) { - Template.buildStringToStringMap(xContentBuilder, map); + buildStringToStringMap(xContentBuilder, map); } } xContentBuilder.endArray(); @@ -127,7 +129,7 @@ public static WorkflowNode parse(XContentParser parser) throws IOException { inputs.put(inputFieldName, parser.text()); break; case START_OBJECT: - inputs.put(inputFieldName, Template.parseStringToStringMap(parser)); + inputs.put(inputFieldName, parseStringToStringMap(parser)); break; case START_ARRAY: if (PROCESSORS_FIELD.equals(inputFieldName)) { @@ -139,7 +141,7 @@ public static WorkflowNode parse(XContentParser parser) throws IOException { } else { List> mapList = new ArrayList<>(); while (parser.nextToken() != XContentParser.Token.END_ARRAY) { - mapList.add(Template.parseStringToStringMap(parser)); + mapList.add(parseStringToStringMap(parser)); } inputs.put(inputFieldName, mapList.toArray(new Map[0])); }