More code review comments and refactoring

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Jan 24, 2025 · 3550786 · 3550786
1 parent bbed331
commit 3550786
Show file tree

Hide file tree

Showing 17 changed files with 156 additions and 199 deletions.
diff --git a/build.gradle b/build.gradle
@@ -192,9 +192,6 @@ dependencies {
     // Multi-tenant SDK Client
     implementation ("org.opensearch:opensearch-remote-metadata-sdk:${opensearch_build}") {
         exclude group: "org.apache.httpcomponents.client5", module: "httpclient5"
-        // Temporary pending merge of https://github.com/opensearch-project/opensearch-remote-metadata-sdk/pull/56
-        exclude group: "org.apache.logging.log4j", module: "log4j-api"
-        exclude group: "org.apache.logging.log4j", module: "log4j-core"
     }
     testImplementation 'org.junit.jupiter:junit-jupiter:5.11.4'
     testImplementation "com.fasterxml.jackson.core:jackson-annotations:${versions.jackson}"
@@ -281,7 +278,7 @@ integTest {
     if (System.getProperty("tests.rest.cluster") != null && System.getProperty("tests.rest.tenantaware") == null) {
         filter {
             includeTestsMatching "org.opensearch.flowframework.rest.*IT"
-            excludeTestsMatching "org.opensearch.ml.rest.*TenantAwareIT"
+            excludeTestsMatching "org.opensearch.flowframework.rest.*TenantAwareIT"
         }
     }
 
@@ -304,6 +301,7 @@ integTest {
         filter {
             includeTestsMatching "org.opensearch.flowframework.rest.FlowFrameworkSecureRestApiIT"
             excludeTestsMatching "org.opensearch.flowframework.rest.FlowFrameworkRestApiIT"
+            excludeTestsMatching "org.opensearch.flowframework.rest.*TenantAwareIT"
         }
     }
 

diff --git a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java
@@ -475,7 +475,7 @@ public void updateTemplateInGlobalContext(
                 documentId
             ).getFormattedMessage();
             logger.error(errorMessage);
-            listener.onFailure(new FlowFrameworkException(errorMessage, RestStatus.BAD_REQUEST));
+            listener.onFailure(new FlowFrameworkException(errorMessage, RestStatus.INTERNAL_SERVER_ERROR));
             return;
         }
         doesTemplateExist(documentId, tenantId, templateExists -> {
@@ -792,15 +792,14 @@ public void updateFlowFrameworkSystemIndexDoc(
                 sdkClient.updateDataObjectAsync(updateRequest).whenComplete((r, throwable) -> {
                     context.restore();
                     if (throwable == null) {
-                        UpdateResponse response;
                         try {
-                            response = UpdateResponse.fromXContent(r.parser());
-                            logger.info("Deleted workflow state doc: {}", documentId);
+                            UpdateResponse response = UpdateResponse.fromXContent(r.parser());
+                            logger.info("Updated workflow state doc: {}", documentId);
                             listener.onResponse(response);
                         } catch (Exception e) {
-                            logger.error("Failed to parse delete response", e);
+                            logger.error("Failed to parse update response", e);
                             listener.onFailure(
-                                new FlowFrameworkException("Failed to parse delete response", RestStatus.INTERNAL_SERVER_ERROR)
+                                new FlowFrameworkException("Failed to parse update response", RestStatus.INTERNAL_SERVER_ERROR)
                             );
                         }
                     } else {

diff --git a/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java
@@ -17,7 +17,7 @@
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.flowframework.common.FlowFrameworkSettings;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
-import org.opensearch.flowframework.util.RestActionUtils;
+import org.opensearch.flowframework.util.TenantAwareHelper;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestChannel;
@@ -84,7 +84,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);
+            String tenantId = TenantAwareHelper.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request);
             SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
             searchSourceBuilder.parseXContent(request.contentOrSourceParamParser());
             searchSourceBuilder.seqNoAndPrimaryTerm(true).version(true);

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java
@@ -25,7 +25,7 @@
 import org.opensearch.flowframework.transport.CreateWorkflowAction;
 import org.opensearch.flowframework.transport.WorkflowRequest;
 import org.opensearch.flowframework.util.ParseUtils;
-import org.opensearch.flowframework.util.RestActionUtils;
+import org.opensearch.flowframework.util.TenantAwareHelper;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestRequest;
@@ -115,7 +115,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);
+            String tenantId = TenantAwareHelper.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.",

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestDeleteWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestDeleteWorkflowAction.java
@@ -20,7 +20,7 @@
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.transport.DeleteWorkflowAction;
 import org.opensearch.flowframework.transport.WorkflowRequest;
-import org.opensearch.flowframework.util.RestActionUtils;
+import org.opensearch.flowframework.util.TenantAwareHelper;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestRequest;
@@ -73,7 +73,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request
                     RestStatus.FORBIDDEN
                 );
             }
-            String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
+            String tenantId = TenantAwareHelper.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
 
             // Always consume content to silently ignore it
             // https://github.com/opensearch-project/flow-framework/issues/578

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestDeprovisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestDeprovisionWorkflowAction.java
@@ -20,7 +20,7 @@
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.transport.DeprovisionWorkflowAction;
 import org.opensearch.flowframework.transport.WorkflowRequest;
-import org.opensearch.flowframework.util.RestActionUtils;
+import org.opensearch.flowframework.util.TenantAwareHelper;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestRequest;
@@ -70,7 +70,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request
                     RestStatus.FORBIDDEN
                 );
             }
-            String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
+            String tenantId = TenantAwareHelper.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
 
             // Always consume content to silently ignore it
             // https://github.com/opensearch-project/flow-framework/issues/578

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowAction.java
@@ -20,7 +20,7 @@
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.transport.GetWorkflowAction;
 import org.opensearch.flowframework.transport.WorkflowRequest;
-import org.opensearch.flowframework.util.RestActionUtils;
+import org.opensearch.flowframework.util.TenantAwareHelper;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestRequest;
@@ -71,7 +71,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request
                     RestStatus.FORBIDDEN
                 );
             }
-            String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
+            String tenantId = TenantAwareHelper.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
             // Always consume content to silently ignore it
             // https://github.com/opensearch-project/flow-framework/issues/578
             request.content();

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStateAction.java b/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStateAction.java
@@ -20,7 +20,7 @@
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.transport.GetWorkflowStateAction;
 import org.opensearch.flowframework.transport.GetWorkflowStateRequest;
-import org.opensearch.flowframework.util.RestActionUtils;
+import org.opensearch.flowframework.util.TenantAwareHelper;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestRequest;
@@ -66,7 +66,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request
                     RestStatus.FORBIDDEN
                 );
             }
-            String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
+            String tenantId = TenantAwareHelper.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
 
             // Always consume content to silently ignore it
             // https://github.com/opensearch-project/flow-framework/issues/578

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStepAction.java b/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStepAction.java
@@ -20,7 +20,7 @@
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.transport.GetWorkflowStepAction;
 import org.opensearch.flowframework.transport.WorkflowRequest;
-import org.opensearch.flowframework.util.RestActionUtils;
+import org.opensearch.flowframework.util.TenantAwareHelper;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestRequest;
@@ -71,7 +71,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
                     RestStatus.FORBIDDEN
                 );
             }
-            String tenantId = RestActionUtils.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request);
+            String tenantId = TenantAwareHelper.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request);
 
             // Always consume content to silently ignore it
             // https://github.com/opensearch-project/flow-framework/issues/578

diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java
@@ -22,7 +22,7 @@
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.transport.ProvisionWorkflowAction;
 import org.opensearch.flowframework.transport.WorkflowRequest;
-import org.opensearch.flowframework.util.RestActionUtils;
+import org.opensearch.flowframework.util.TenantAwareHelper;
 import org.opensearch.rest.BaseRestHandler;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestRequest;
@@ -86,7 +86,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
                     RestStatus.FORBIDDEN
                 );
             }
-            String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
+            String tenantId = TenantAwareHelper.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request);
             // Validate params
             if (workflowId == null) {
                 throw new FlowFrameworkException("workflow_id cannot be null", RestStatus.BAD_REQUEST);

diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java
@@ -76,7 +76,7 @@ public class EncryptorUtils {
     private static final String WRAPPING_ALGORITHM = "AES/GCM/NOPADDING";
 
     // concurrent map can't have null as a key. This key is to support single tenancy
-    public static final String DEFAULT_TENANT_ID = "03000200-0400-0500-0006-000700080009";
+    public static final String DEFAULT_TENANT_ID = "";
 
     private final ClusterService clusterService;
     private final Client client;

diff --git a/src/main/java/org/opensearch/flowframework/util/RestActionUtils.java b/src/main/java/org/opensearch/flowframework/util/RestActionUtils.java
diff --git a/src/main/java/org/opensearch/flowframework/util/TenantAwareHelper.java b/src/main/java/org/opensearch/flowframework/util/TenantAwareHelper.java
@@ -10,8 +10,12 @@
 
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.flowframework.common.CommonValue;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
+import org.opensearch.rest.RestRequest;
 
+import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -29,7 +33,7 @@ public class TenantAwareHelper {
      */
     public static boolean validateTenantId(boolean isMultiTenancyEnabled, String tenantId, ActionListener<?> listener) {
         if (isMultiTenancyEnabled && tenantId == null) {
-            listener.onFailure(new FlowFrameworkException("You don't have permission to access this resource", RestStatus.FORBIDDEN));
+            listener.onFailure(new FlowFrameworkException("No permission to access this resource", RestStatus.FORBIDDEN));
             return false;
         } else {
             return true;
@@ -52,8 +56,34 @@ public static boolean validateTenantResource(
         ActionListener<?> listener
     ) {
         if (isMultiTenancyEnabled && !Objects.equals(tenantIdFromRequest, tenantIdFromResource)) {
-            listener.onFailure(new FlowFrameworkException("You don't have permission to access this resource", RestStatus.FORBIDDEN));
+            listener.onFailure(new FlowFrameworkException("No permission to access this resource", RestStatus.FORBIDDEN));
             return false;
         } else return true;
     }
+
+    /**
+     * Finds the tenant id in the REST Headers
+     * @param isMultiTenancyEnabled whether multitenancy is enabled
+     * @param restRequest the RestRequest
+     * @return The tenant ID from the headers or null if multitenancy is not enabled
+     */
+    public static String getTenantID(Boolean isMultiTenancyEnabled, RestRequest restRequest) {
+        if (!isMultiTenancyEnabled) {
+            return null;
+        }
+
+        Map<String, List<String>> headers = restRequest.getHeaders();
+
+        List<String> tenantIdList = headers.get(CommonValue.TENANT_ID_HEADER);
+        if (tenantIdList == null || tenantIdList.isEmpty()) {
+            throw new FlowFrameworkException("Tenant ID header is missing or has no value", RestStatus.FORBIDDEN);
+        }
+
+        String tenantId = tenantIdList.get(0);
+        if (tenantId == null) {
+            throw new FlowFrameworkException("Tenant ID can't be null", RestStatus.FORBIDDEN);
+        }
+
+        return tenantId;
+    }
 }
diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkTenantAwareRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkTenantAwareRestTestCase.java
@@ -55,8 +55,8 @@ public abstract class FlowFrameworkTenantAwareRestTestCase extends FlowFramework
     protected static final String EMPTY_CONTENT = "{}";
 
     // REST Response error reasons
-    protected static final String MISSING_TENANT_REASON = "Tenant ID header is missing";
-    protected static final String NO_PERMISSION_REASON = "You don't have permission to access this resource";
+    protected static final String MISSING_TENANT_REASON = "Tenant ID header is missing or has no value";
+    protected static final String NO_PERMISSION_REASON = "No permission to access this resource";
 
     protected String tenantId = randomAlphaOfLength(5);
     protected String otherTenantId = randomAlphaOfLength(6);

diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
@@ -810,9 +810,7 @@ public void testAllDefaultUseCasesCreation() throws Exception {
         }
     }
 
-    // TODO Re-enable
-    // This test is currently failing due to ML Commons MLPredictionTaskRequest serialization
-    public void disableTemporarilytestSemanticSearchWithLocalModelEndToEnd() throws Exception {
+    public void testSemanticSearchWithLocalModelEndToEnd() throws Exception {
         // Checking if plugins are part of the integration test cluster so we can continue with this test
         List<String> plugins = catPlugins();
         if (!plugins.contains("opensearch-knn") && !plugins.contains("opensearch-neural-search")) {