Integrate multitenancy into workflow provisioning (#1012)

* Adding tenantId in ProcessNode to be surfaced to the client Signed-off-by: Siddhartha Bingi <sidbingi@amazon.com> * Passing tenantId in Transport Actions Signed-off-by: Siddhartha Bingi <sidbingi@amazon.com> * Updating tests to have tenantId as null for consistency Signed-off-by: Siddhartha Bingi <sidbingi@amazon.com> * Minor code refactoring Signed-off-by: Siddhartha Bingi <sidbingi@amazon.com> --------- Signed-off-by: Siddhartha Bingi <sidbingi@amazon.com> Co-authored-by: Siddhartha Bingi <sidbingi@amazon.com> (cherry picked from commit 055f1ea) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Jan 23, 2025 · dd7dace · dd7dace
1 parent 123ec5d
commit dd7dace
Show file tree

Hide file tree

Showing 57 changed files with 312 additions and 156 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
@@ -517,7 +517,7 @@ void checkMaxWorkflows(TimeValue requestTimeOut, Integer maxWorkflow, ActionList
 
     private void validateWorkflows(Template template) throws Exception {
         for (Workflow workflow : template.workflows().values()) {
-            List<ProcessNode> sortedNodes = workflowProcessSorter.sortProcessNodes(workflow, null, Collections.emptyMap());
+            List<ProcessNode> sortedNodes = workflowProcessSorter.sortProcessNodes(workflow, null, Collections.emptyMap(), "fakeTenantId");
             workflowProcessSorter.validate(sortedNodes, pluginsService);
         }
     }

diff --git a/src/main/java/org/opensearch/flowframework/transport/DeprovisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/DeprovisionWorkflowTransportAction.java
@@ -211,7 +211,8 @@ private void executeDeprovisionSequence(
                     Collections.emptyList(),
                     this.threadPool,
                     DEPROVISION_WORKFLOW_THREAD_POOL,
-                    flowFrameworkSettings.getRequestTimeout()
+                    flowFrameworkSettings.getRequestTimeout(),
+                    "fakeTenantId"
                 )
             );
         }
@@ -274,7 +275,8 @@ private void executeDeprovisionSequence(
                         pn.predecessors(),
                         this.threadPool,
                         DEPROVISION_WORKFLOW_THREAD_POOL,
-                        pn.nodeTimeout()
+                        pn.nodeTimeout(),
+                        "fakeTenantId"
                     );
                 }).collect(Collectors.toList());
                 // Pause briefly before next loop

diff --git a/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/ProvisionWorkflowTransportAction.java
@@ -198,7 +198,8 @@ private void executeProvisionRequest(
             List<ProcessNode> provisionProcessSequence = workflowProcessSorter.sortProcessNodes(
                 provisionWorkflow,
                 workflowId,
-                request.getParams()
+                request.getParams(),
+                "fakeTenantId"
             );
             workflowProcessSorter.validate(provisionProcessSequence, pluginsService);
 

diff --git a/src/main/java/org/opensearch/flowframework/transport/ReprovisionWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/ReprovisionWorkflowTransportAction.java
@@ -198,7 +198,8 @@ private void executeReprovisionRequest(
             List<ProcessNode> updatedProcessSequence = workflowProcessSorter.sortProcessNodes(
                 provisionWorkflow,
                 request.getWorkflowId(),
-                Collections.emptyMap() // TODO : Add suport to reprovision substitution templates
+                Collections.emptyMap(), // TODO : Add suport to reprovision substitution templates
+                "fakeTenantId"
             );
 
             try {

diff --git a/src/main/java/org/opensearch/flowframework/workflow/AbstractCreatePipelineStep.java b/src/main/java/org/opensearch/flowframework/workflow/AbstractCreatePipelineStep.java
@@ -63,7 +63,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
 
         PlainActionFuture<WorkflowData> createPipelineFuture = PlainActionFuture.newFuture();

diff --git a/src/main/java/org/opensearch/flowframework/workflow/AbstractRegisterLocalModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/AbstractRegisterLocalModelStep.java
@@ -89,7 +89,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
 
         PlainActionFuture<WorkflowData> registerLocalModelFuture = PlainActionFuture.newFuture();

diff --git a/src/main/java/org/opensearch/flowframework/workflow/AbstractUpdatePipelineStep.java b/src/main/java/org/opensearch/flowframework/workflow/AbstractUpdatePipelineStep.java
@@ -70,7 +70,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> createPipelineFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java
@@ -76,7 +76,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> createConnectorFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateIndexStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateIndexStep.java
@@ -67,7 +67,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> createIndexFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteAgentStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteAgentStep.java
@@ -55,7 +55,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> deleteAgentFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteConnectorStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteConnectorStep.java
@@ -53,7 +53,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> deleteConnectorFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteIndexStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteIndexStep.java
@@ -58,7 +58,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> deleteIndexFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteIngestPipelineStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteIngestPipelineStep.java
@@ -58,7 +58,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> deletePipelineFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteModelStep.java
@@ -55,7 +55,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> deleteModelFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteSearchPipelineStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteSearchPipelineStep.java
@@ -58,7 +58,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> deleteSearchPipelineFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java
@@ -65,7 +65,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
 
         PlainActionFuture<WorkflowData> deployModelFuture = PlainActionFuture.newFuture();

diff --git a/src/main/java/org/opensearch/flowframework/workflow/NoOpStep.java b/src/main/java/org/opensearch/flowframework/workflow/NoOpStep.java
@@ -38,7 +38,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> future = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/ProcessNode.java b/src/main/java/org/opensearch/flowframework/workflow/ProcessNode.java
@@ -36,6 +36,7 @@ public class ProcessNode {
     private final ThreadPool threadPool;
     private final String threadPoolName;
     private final TimeValue nodeTimeout;
+    private final String tenantId;
 
     private final PlainActionFuture<WorkflowData> future = PlainActionFuture.newFuture();
 
@@ -51,6 +52,7 @@ public class ProcessNode {
      * @param threadPool The OpenSearch thread pool
      * @param threadPoolName The thread pool to use
      * @param nodeTimeout The timeout value for executing on this node
+     * @param tenantId The tenantId
      */
     public ProcessNode(
         String id,
@@ -61,7 +63,8 @@ public ProcessNode(
         List<ProcessNode> predecessors,
         ThreadPool threadPool,
         String threadPoolName,
-        TimeValue nodeTimeout
+        TimeValue nodeTimeout,
+        String tenantId
     ) {
         this.id = id;
         this.workflowStep = workflowStep;
@@ -72,6 +75,7 @@ public ProcessNode(
         this.threadPool = threadPool;
         this.threadPoolName = threadPoolName;
         this.nodeTimeout = nodeTimeout;
+        this.tenantId = tenantId;
     }
 
     /**
@@ -144,6 +148,14 @@ public TimeValue nodeTimeout() {
         return nodeTimeout;
     }
 
+    /**
+     * Returns the tenantId value for this node in the workflow.
+     * @return The node's tenantId value
+     */
+    public String tenantId() {
+        return tenantId;
+    }
+
     /**
      * Execute this node in the sequence.
      * Initializes the node's {@link CompletableFuture} and completes it when the process completes.
@@ -172,7 +184,8 @@ public PlainActionFuture<WorkflowData> execute() {
                     this.input,
                     inputMap,
                     this.previousNodeInputs,
-                    this.params
+                    this.params,
+                    this.tenantId
                 );
                 // If completed exceptionally, this is a no-op
                 future.onResponse(stepFuture.actionGet(this.nodeTimeout));

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java
@@ -84,7 +84,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
 
         String workflowId = currentNodeInputs.getWorkflowId();

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java
@@ -70,7 +70,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> registerModelGroupFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterRemoteModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterRemoteModelStep.java
@@ -74,7 +74,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
 
         PlainActionFuture<WorkflowData> registerRemoteModelFuture = PlainActionFuture.newFuture();

diff --git a/src/main/java/org/opensearch/flowframework/workflow/ReindexStep.java b/src/main/java/org/opensearch/flowframework/workflow/ReindexStep.java
@@ -68,7 +68,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
 
         PlainActionFuture<WorkflowData> reIndexFuture = PlainActionFuture.newFuture();

diff --git a/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java b/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java
@@ -63,7 +63,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         try {
             Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(

diff --git a/src/main/java/org/opensearch/flowframework/workflow/UndeployModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/UndeployModelStep.java
@@ -58,7 +58,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> undeployModelFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/UpdateIndexStep.java b/src/main/java/org/opensearch/flowframework/workflow/UpdateIndexStep.java
@@ -64,7 +64,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> updateIndexFuture = PlainActionFuture.newFuture();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowDataStep.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowDataStep.java
@@ -40,7 +40,8 @@ public PlainActionFuture<WorkflowData> execute(
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     ) {
         PlainActionFuture<WorkflowData> workflowDataFuture = PlainActionFuture.newFuture();
         workflowDataFuture.onResponse(

diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowProcessSorter.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowProcessSorter.java
@@ -90,9 +90,10 @@ public WorkflowProcessSorter(
      * @param workflow A workflow with (unsorted) nodes and edges which define predecessors and successors
      * @param workflowId The workflowId associated with the step
      * @param params Parameters passed on the REST path
+     * @param tenantId The tenantId associated with the step
      * @return A list of Process Nodes sorted topologically.  All predecessors of any node will occur prior to it in the list.
      */
-    public List<ProcessNode> sortProcessNodes(Workflow workflow, String workflowId, Map<String, String> params) {
+    public List<ProcessNode> sortProcessNodes(Workflow workflow, String workflowId, Map<String, String> params, String tenantId) {
         if (workflow.nodes().size() > this.maxWorkflowSteps) {
             throw new FlowFrameworkException(
                 "Workflow "
@@ -140,7 +141,8 @@ public List<ProcessNode> sortProcessNodes(Workflow workflow, String workflowId,
                 predecessorNodes,
                 threadPool,
                 PROVISION_WORKFLOW_THREAD_POOL,
-                nodeTimeout
+                nodeTimeout,
+                tenantId
             );
             idToNodeMap.put(processNode.id(), processNode);
             nodes.add(processNode);
@@ -319,7 +321,8 @@ private ProcessNode createNewProcessNode(
             predecessorNodes,
             threadPool,
             PROVISION_WORKFLOW_THREAD_POOL,
-            nodeTimeout
+            nodeTimeout,
+            "fakeTenantId"
         );
     }
 
@@ -350,7 +353,8 @@ private ProcessNode createUpdateProcessNode(
                 predecessorNodes,
                 threadPool,
                 PROVISION_WORKFLOW_THREAD_POOL,
-                nodeTimeout
+                nodeTimeout,
+                "fakeTenantId"
             );
         } else {
             // Case 3 : Cannot update step (not supported)
@@ -392,7 +396,8 @@ private ProcessNode createWorkflowDataStepNode(
                 predecessorNodes,
                 threadPool,
                 PROVISION_WORKFLOW_THREAD_POOL,
-                nodeTimeout
+                nodeTimeout,
+                "fakeTenantId"
             );
         } else {
             return null;

diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStep.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStep.java
@@ -24,14 +24,16 @@ public interface WorkflowStep {
      * @param outputs WorkflowData content of previous steps.
      * @param previousNodeInputs Input params for this node that come from previous steps
      * @param params Params passed on the REST path
+     * @param tenantId The tenantId
      * @return A CompletableFuture of the building block. This block should return immediately, but not be completed until the step executes, containing either the step's output data or {@link WorkflowData#EMPTY} which may be passed to follow-on steps.
      */
     PlainActionFuture<WorkflowData> execute(
         String currentNodeId,
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
         Map<String, String> previousNodeInputs,
-        Map<String, String> params
+        Map<String, String> params,
+        String tenantId
     );
 
     /**