Pass tenant id to transport actions in template

Signed-off-by: Daniel Widdis <[email protected]>

src/main/java/org/opensearch/flowframework/model/Template.java

@@ -38,6 +38,7 @@
 import static org.opensearch.flowframework.common.CommonValue.LAST_PROVISIONED_TIME_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.LAST_UPDATED_TIME_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.TENANT_ID_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.UI_METADATA_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.USER_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.VERSION_FIELD;
@@ -75,6 +76,7 @@ public class Template implements ToXContentObject {
     private final Instant createdTime;
     private final Instant lastUpdatedTime;
     private final Instant lastProvisionedTime;
+    private String tenantId;
 
     /**
      * Instantiate the object representing a use case template
@@ -358,6 +360,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             xContentBuilder.field(LAST_PROVISIONED_TIME_FIELD, lastProvisionedTime.toEpochMilli());
         }
 
+        if (tenantId != null) {
+            xContentBuilder.field(TENANT_ID_FIELD, tenantId);
+        }
+
         return xContentBuilder.endObject();
     }
 
@@ -421,6 +427,7 @@ public static Template parse(XContentParser parser, boolean fieldUpdate) throws
         Instant createdTime = null;
         Instant lastUpdatedTime = null;
         Instant lastProvisionedTime = null;
+        String tenantId = null;
 
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -488,6 +495,9 @@ public static Template parse(XContentParser parser, boolean fieldUpdate) throws
                 case LAST_PROVISIONED_TIME_FIELD:
                     lastProvisionedTime = ParseUtils.parseInstant(parser);
                     break;
+                case TENANT_ID_FIELD:
+                    tenantId = parser.text();
+                    break;
                 default:
                     throw new FlowFrameworkException(
                         "Unable to parse field [" + fieldName + "] in a template object.",
@@ -507,7 +517,7 @@ public static Template parse(XContentParser parser, boolean fieldUpdate) throws
             }
         }
 
-        return new Builder().name(name)
+        Template template = new Builder().name(name)
             .description(description)
             .useCase(useCase)
             .templateVersion(templateVersion)
@@ -519,6 +529,10 @@ public static Template parse(XContentParser parser, boolean fieldUpdate) throws
             .lastUpdatedTime(lastUpdatedTime)
             .lastProvisionedTime(lastProvisionedTime)
             .build();
+        if (tenantId != null) {
+            template.setTenantId(tenantId);
+        }
+        return template;
     }
 
     /**
@@ -541,6 +555,21 @@ public static Template parse(String json) throws IOException {
         }
     }
 
+    /**
+     * Creates an empty template with the given tenant ID
+     *
+     * @param tenantId the tenantID
+     * @return an empty template containing the tenant id if it's not null, null otherwise
+     */
+    public static Template createEmptyTemplateWithTenantId(String tenantId) {
+        if (tenantId == null) {
+            return null;
+        }
+        Template emptyTemplate = builder().name("").build();
+        emptyTemplate.setTenantId(tenantId);
+        return emptyTemplate;
+    }
+
     /**
      * Output this object in a compact JSON string.
      *
@@ -657,6 +686,22 @@ public Instant lastProvisionedTime() {
         return lastProvisionedTime;
     }
 
+    /**
+     * The tenant id
+     * @return the tenant id
+     */
+    public String getTenantId() {
+        return tenantId;
+    }
+
+    /**
+     * Sets the tenant id
+     * @param tenantId the tenant id to set
+     */
+    public void setTenantId(String tenantId) {
+        this.tenantId = tenantId;
+    }
+
     @Override
     public String toString() {
         return "Template [name="

src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java

@@ -111,6 +111,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
                 new BytesRestResponse(ffe.getRestStatus(), ffe.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS))
             );
         }
+        String tenantId = RestActionUtils.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request);
         if (!provision && !params.isEmpty()) {
             FlowFrameworkException ffe = new FlowFrameworkException(
                 "Only the parameters " + request.consumedParams() + " are permitted unless the provision parameter is set to true.",
@@ -146,7 +147,6 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             );
             return processError(ffe, params, request);
         }
-        String tenantId = RestActionUtils.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request);
         try {
             Template template;
             Map<String, String> useCaseDefaultsMap = Collections.emptyMap();
@@ -221,6 +221,9 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             if (updateFields) {
                 params = Map.of(UPDATE_WORKFLOW_FIELDS, "true");
             }
+            if (tenantId != null) {
+                template.setTenantId(tenantId);
+            }
 
             WorkflowRequest workflowRequest = new WorkflowRequest(
                 workflowId,

src/main/java/org/opensearch/flowframework/rest/RestDeleteWorkflowAction.java

@@ -33,6 +33,7 @@
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI;
 import static org.opensearch.flowframework.common.FlowFrameworkSettings.FLOW_FRAMEWORK_ENABLED;
+import static org.opensearch.flowframework.model.Template.createEmptyTemplateWithTenantId;
 
 /**
  * Rest Action to facilitate requests to delete a stored template
@@ -82,7 +83,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request
             if (workflowId == null) {
                 throw new FlowFrameworkException("workflow_id cannot be null", RestStatus.BAD_REQUEST);
             }
-            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null, request.params());
+            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, createEmptyTemplateWithTenantId(tenantId), request.params());
             return channel -> client.execute(DeleteWorkflowAction.INSTANCE, workflowRequest, ActionListener.wrap(response -> {
                 XContentBuilder builder = response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS);
                 channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder));

src/main/java/org/opensearch/flowframework/rest/RestDeprovisionWorkflowAction.java

@@ -35,6 +35,7 @@
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI;
 import static org.opensearch.flowframework.common.FlowFrameworkSettings.FLOW_FRAMEWORK_ENABLED;
+import static org.opensearch.flowframework.model.Template.createEmptyTemplateWithTenantId;
 
 /**
  * Rest Action to facilitate requests to de-provision a workflow
@@ -81,7 +82,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request
             }
             WorkflowRequest workflowRequest = new WorkflowRequest(
                 workflowId,
-                null,
+                createEmptyTemplateWithTenantId(tenantId),
                 allowDelete == null ? Collections.emptyMap() : Map.of(ALLOW_DELETE, allowDelete)
             );
 

src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowAction.java

@@ -32,6 +32,7 @@
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI;
 import static org.opensearch.flowframework.common.FlowFrameworkSettings.FLOW_FRAMEWORK_ENABLED;
+import static org.opensearch.flowframework.model.Template.createEmptyTemplateWithTenantId;
 
 /**
  * Rest Action to facilitate requests to get a stored template
@@ -71,7 +72,6 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request
                 );
             }
             String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
-
             // Always consume content to silently ignore it
             // https://github.com/opensearch-project/flow-framework/issues/578
             request.content();
@@ -81,7 +81,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request
                 throw new FlowFrameworkException("workflow_id cannot be null", RestStatus.BAD_REQUEST);
             }
 
-            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null);
+            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, createEmptyTemplateWithTenantId(tenantId));
             return channel -> client.execute(GetWorkflowAction.INSTANCE, workflowRequest, ActionListener.wrap(response -> {
                 XContentBuilder builder = response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS);
                 channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder));

src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java

@@ -37,6 +37,7 @@
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI;
 import static org.opensearch.flowframework.common.FlowFrameworkSettings.FLOW_FRAMEWORK_ENABLED;
+import static org.opensearch.flowframework.model.Template.createEmptyTemplateWithTenantId;
 
 /**
  * Rest action to facilitate requests to provision a workflow from an inline defined or stored use case template
@@ -82,13 +83,13 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
                     RestStatus.FORBIDDEN
                 );
             }
+            String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
             // Validate params
             if (workflowId == null) {
                 throw new FlowFrameworkException("workflow_id cannot be null", RestStatus.BAD_REQUEST);
             }
-            String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
             // Create request and provision
-            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, null, params);
+            WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, createEmptyTemplateWithTenantId(tenantId), params);
             return channel -> client.execute(ProvisionWorkflowAction.INSTANCE, workflowRequest, ActionListener.wrap(response -> {
                 XContentBuilder builder = response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS);
                 channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder));

src/test/java/org/opensearch/flowframework/model/TemplateTests.java

@@ -19,6 +19,7 @@
 
 import java.io.IOException;
 import java.time.Instant;
+import java.time.temporal.ChronoUnit;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -47,7 +48,7 @@ public void testTemplate() throws IOException {
         Workflow workflow = new Workflow(Map.of("key", "value"), nodes, edges);
         Map<String, Object> uiMetadata = null;
 
-        Instant now = Instant.now();
+        Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS);
         Template template = new Template(
             "test",
             "a test template",
@@ -74,6 +75,9 @@ public void testTemplate() throws IOException {
         assertEquals(now, template.lastUpdatedTime());
         assertNull(template.lastProvisionedTime());
         assertEquals("Workflow [userParams={key=value}, nodes=[A, B], edges=[A->B]]", wf.toString());
+        assertNull(template.getTenantId());
+        template.setTenantId("tenant-id");
+        assertEquals("tenant-id", template.getTenantId());
 
         String json = TemplateTestJsonUtil.parseToJson(template);
 
@@ -86,10 +90,11 @@ public void testTemplate() throws IOException {
         assertEquals(uiMetadata, templateX.getUiMetadata());
         Workflow wfX = templateX.workflows().get("workflow");
         assertNotNull(wfX);
-        assertEquals(now, template.createdTime());
-        assertEquals(now, template.lastUpdatedTime());
-        assertNull(template.lastProvisionedTime());
+        assertEquals(now, templateX.createdTime());
+        assertEquals(now, templateX.lastUpdatedTime());
+        assertNull(templateX.lastProvisionedTime());
         assertEquals("Workflow [userParams={key=value}, nodes=[A, B], edges=[A->B]]", wfX.toString());
+        assertEquals("tenant-id", templateX.getTenantId());
 
         // Test invalid field if updating
         XContentParser parser = JsonXContent.jsonXContent.createParser(