From a296eca561980fe3c58d2b5e1f8f1cbeaafe621f Mon Sep 17 00:00:00 2001 From: Daniel Widdis Date: Fri, 13 Dec 2024 00:30:54 -0800 Subject: [PATCH] Pass tenant id to transport actions in template Signed-off-by: Daniel Widdis --- .../flowframework/model/Template.java | 47 ++++++++++++++++++- .../rest/RestCreateWorkflowAction.java | 5 +- .../rest/RestDeleteWorkflowAction.java | 3 +- .../rest/RestDeprovisionWorkflowAction.java | 3 +- .../rest/RestGetWorkflowAction.java | 4 +- .../rest/RestProvisionWorkflowAction.java | 5 +- .../flowframework/model/TemplateTests.java | 13 +++-- 7 files changed, 68 insertions(+), 12 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/model/Template.java b/src/main/java/org/opensearch/flowframework/model/Template.java index 4c8c2c9ef..070e222d3 100644 --- a/src/main/java/org/opensearch/flowframework/model/Template.java +++ b/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=" diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java index d15c14bf3..68774884d 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java +++ b/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 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, diff --git a/src/main/java/org/opensearch/flowframework/rest/RestDeleteWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestDeleteWorkflowAction.java index 6db957be2..0ea4de8e0 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestDeleteWorkflowAction.java +++ b/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)); diff --git a/src/main/java/org/opensearch/flowframework/rest/RestDeprovisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestDeprovisionWorkflowAction.java index 90559bc92..2f89fb8e2 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestDeprovisionWorkflowAction.java +++ b/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) ); diff --git a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowAction.java index 50551e019..ec08f4a09 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowAction.java +++ b/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)); diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java index f6d8458b1..e30fe7b09 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java +++ b/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)); diff --git a/src/test/java/org/opensearch/flowframework/model/TemplateTests.java b/src/test/java/org/opensearch/flowframework/model/TemplateTests.java index d8550682b..03f1bdc35 100644 --- a/src/test/java/org/opensearch/flowframework/model/TemplateTests.java +++ b/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 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(