adding workflow state and create connector resources created

Signed-off-by: Amit Galitzky <[email protected]>
opensearch-project · Nov 6, 2023 · 8b09716 · 8b09716
1 parent ac76a44
commit 8b09716
Show file tree

Hide file tree

Showing 22 changed files with 745 additions and 55 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java b/src/main/java/org/opensearch/flowframework/FlowFrameworkPlugin.java
@@ -77,10 +77,9 @@ public Collection<Object> createComponents(
         Supplier<RepositoriesService> repositoriesServiceSupplier
     ) {
         MachineLearningNodeClient mlClient = new MachineLearningNodeClient(client);
-        WorkflowStepFactory workflowStepFactory = new WorkflowStepFactory(clusterService, client, mlClient);
-        WorkflowProcessSorter workflowProcessSorter = new WorkflowProcessSorter(workflowStepFactory, threadPool);
-
         FlowFrameworkIndicesHandler flowFrameworkIndicesHandler = new FlowFrameworkIndicesHandler(client, clusterService);
+        WorkflowStepFactory workflowStepFactory = new WorkflowStepFactory(clusterService, client, mlClient, flowFrameworkIndicesHandler);
+        WorkflowProcessSorter workflowProcessSorter = new WorkflowProcessSorter(workflowStepFactory, threadPool);
 
         return ImmutableList.of(workflowStepFactory, workflowProcessSorter, flowFrameworkIndicesHandler);
     }

diff --git a/src/main/java/org/opensearch/flowframework/common/CommonValue.java b/src/main/java/org/opensearch/flowframework/common/CommonValue.java
@@ -46,7 +46,7 @@ private CommonValue() {}
     public static final String USER_FIELD = "user";
 
     /** The transport action name prefix */
-    public static final String TRANSPORT_ACION_NAME_PREFIX = "cluster:admin/opensearch/flow_framework/";
+    public static final String TRANSPORT_ACTION_NAME_PREFIX = "cluster:admin/opensearch/flow_framework/";
     /** The base URI for this plugin's rest actions */
     public static final String FLOW_FRAMEWORK_BASE_URI = "/_plugins/_flow_framework";
     /** The URI for this plugin's workflow rest actions */
@@ -130,4 +130,9 @@ private CommonValue() {}
     public static final String USER_OUTPUTS_FIELD = "user_outputs";
     /** The template field name for template resources created */
     public static final String RESOURCES_CREATED_FIELD = "resources_created";
+    /** The field name for the ResourcesCreated's resource ID */
+    public static final String RESOURCE_ID_FIELD = "resource_id";
+    /** The field name for the ResourcesCreated's resource name */
+    public static final String RESOURCE_NAME_FIELD = "resource_type";
+
 }
diff --git a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java
@@ -36,6 +36,7 @@
 import org.opensearch.flowframework.model.State;
 import org.opensearch.flowframework.model.Template;
 import org.opensearch.flowframework.model.WorkflowState;
+import org.opensearch.script.Script;
 
 import java.io.IOException;
 import java.net.URL;
@@ -292,7 +293,7 @@ public void putInitialStateToWorkflowState(String workflowId, User user, ActionL
             .state(State.NOT_STARTED.name())
             .provisioningProgress(ProvisioningProgress.NOT_STARTED.name())
             .user(user)
-            .resourcesCreated(Collections.emptyMap())
+            .resourcesCreated(Collections.emptyList())
             .userOutputs(Collections.emptyMap())
             .build();
         initWorkflowStateIndexIfAbsent(ActionListener.wrap(indexCreated -> {
@@ -381,4 +382,37 @@ public void updateFlowFrameworkSystemIndexDoc(
             }
         }
     }
+
+    /**
+     * Updates a document in the workflow state index
+     * @param indexName the index that we will be updating a document of.
+     * @param documentId the document ID
+     * @param script the given script to update doc
+     * @param listener action listener
+     */
+    public void updateFlowFrameworkSystemIndexDocWithScript(
+        String indexName,
+        String documentId,
+        Script script,
+        ActionListener<UpdateResponse> listener
+    ) {
+        if (!doesIndexExist(indexName)) {
+            String exceptionMessage = "Failed to update document for given workflow due to missing " + indexName + " index";
+            logger.error(exceptionMessage);
+            listener.onFailure(new Exception(exceptionMessage));
+        } else {
+            try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
+                UpdateRequest updateRequest = new UpdateRequest(indexName, documentId);
+                // TODO: Also add ability to change other fields at the same time when adding detailed provision progress
+                updateRequest.script(script);
+                updateRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
+                // TODO: decide what condition can be considered as an update conflict and add retry strategy
+                client.update(updateRequest, ActionListener.runBefore(listener, () -> context.restore()));
+            } catch (Exception e) {
+                logger.error("Failed to update {} entry : {}. {}", indexName, documentId, e.getMessage());
+                listener.onFailure(e);
+            }
+        }
+    }
+
 }
diff --git a/src/main/java/org/opensearch/flowframework/model/ProvisioningProgress.java b/src/main/java/org/opensearch/flowframework/model/ProvisioningProgress.java
@@ -18,5 +18,7 @@ public enum ProvisioningProgress {
     /** In Progress State */
     IN_PROGRESS,
     /** Done State */
-    DONE
+    DONE,
+    /** Failed State */
+    FAILED
 }
diff --git a/src/main/java/org/opensearch/flowframework/model/ResourcesCreated.java b/src/main/java/org/opensearch/flowframework/model/ResourcesCreated.java
@@ -0,0 +1,76 @@
+/*
+ * 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.model;
+
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.xcontent.ToXContentObject;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.core.xcontent.XContentParser;
+
+import java.io.IOException;
+
+import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+import static org.opensearch.flowframework.common.CommonValue.RESOURCE_ID_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.RESOURCE_NAME_FIELD;
+
+public class ResourcesCreated implements ToXContentObject, Writeable {
+
+    public String resourceName;
+    public String resourceId;
+
+    public ResourcesCreated(String resourceName, String resourceId) {
+        this.resourceName = resourceName;
+        this.resourceId = resourceId;
+    }
+
+    public ResourcesCreated(StreamInput input) throws IOException {
+        this.resourceName = input.readString();
+        this.resourceId = input.readString();
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        XContentBuilder xContentBuilder = builder.startObject()
+            .field(RESOURCE_NAME_FIELD, resourceName)
+            .field(RESOURCE_ID_FIELD, resourceId);
+        return xContentBuilder.endObject();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeString(resourceName);
+        out.writeString(resourceId);
+    }
+
+    public static ResourcesCreated parse(XContentParser parser) throws IOException {
+        String resourceName = null;
+        String resourceId = null;
+
+        ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
+        while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
+            String fieldName = parser.currentName();
+            parser.nextToken();
+
+            switch (fieldName) {
+                case RESOURCE_NAME_FIELD:
+                    resourceName = parser.text();
+                    break;
+                case RESOURCE_ID_FIELD:
+                    resourceId = parser.text();
+                    break;
+                default:
+                    throw new IOException("Unable to parse field [" + fieldName + "] in a resources_created object.");
+            }
+        }
+        return new ResourcesCreated(resourceName, resourceId);
+    }
+
+}
diff --git a/src/main/java/org/opensearch/flowframework/model/Template.java b/src/main/java/org/opensearch/flowframework/model/Template.java
@@ -13,6 +13,8 @@
 import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.common.xcontent.yaml.YamlXContent;
 import org.opensearch.commons.authuser.User;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -38,7 +40,7 @@
 /**
  * The Template is the central data structure which configures workflows. This object is used to parse JSON communicated via REST API.
  */
-public class Template implements ToXContentObject {
+public class Template implements ToXContentObject, Writeable {
 
     private final String name;
     private final String description;
@@ -77,6 +79,18 @@ public Template(
         this.user = user;
     }
 
+    // public Template(StreamInput input) throws IOException {
+    // this.name = input.readString();
+    // this.description = input.readOptionalString();
+    // this.useCase = input.readString();
+    // this.templateVersion = input.readVersion();
+    // this.compatibilityVersion = input.readList(Version::new); // Replace with actual method if different
+    // this.workflows = input.readMap(StreamInput::readString, WorkFlow::new); // Replace with the actual function to read WorkFlow objects
+    // if (input.readBoolean()) {
+    // this.user = new User(input); // Replace with the actual constructor or factory method for User
+    // }
+    // }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         XContentBuilder xContentBuilder = builder.startObject();
@@ -111,6 +125,22 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         return xContentBuilder.endObject();
     }
 
+    @Override
+    public void writeTo(StreamOutput output) throws IOException {
+        output.writeString(name);
+        output.writeOptionalString(description);
+        output.writeString(useCase);
+        output.writeVersion(templateVersion);
+        output.writeList((List<? extends Writeable>) compatibilityVersion);
+        output.writeMapWithConsistentOrder(workflows);
+        if (user != null) {
+            output.writeBoolean(true); // user exists
+            user.writeTo(output);
+        } else {
+            output.writeBoolean(false); // user does not exist
+        }
+    }
+
     /**
      * Parse raw json content into a Template instance.
      *