refactor checks to ml indices to return true when MultiTenancy enabled

brianf-aws · brianf-aws · commit 2d0957febc31 · 2025-08-11T12:58:09.000-07:00
Signed-off-by: Brian Flores &lt;iflorbri@amazon.com&gt;
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java
@@ -68,6 +68,7 @@
 import org.opensearch.ml.engine.Executable;
 import org.opensearch.ml.engine.annotation.Function;
 import org.opensearch.ml.engine.encryptor.Encryptor;
+import org.opensearch.ml.engine.indices.MLIndicesHandler;
 import org.opensearch.ml.engine.memory.ConversationIndexMemory;
 import org.opensearch.ml.engine.memory.ConversationIndexMessage;
 import org.opensearch.ml.engine.tools.QueryPlanningTool;
@@ -173,7 +174,7 @@ public void execute(Input input, ActionListener<Output> listener) {
             .fetchSourceContext(fetchSourceContext)
             .build();
 
-        if (clusterService.state().metadata().hasIndex(ML_AGENT_INDEX)) {
+        if (MLIndicesHandler.doesMultiTenantIndexExists(clusterService, mlFeatureEnabledSetting.isMultiTenancyEnabled(), ML_AGENT_INDEX)) {
             try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
                 sdkClient
                     .getDataObjectAsync(getDataObjectRequest, client.threadPool().executor("opensearch_ml_general"))
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/indices/MLIndicesHandler.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/indices/MLIndicesHandler.java
@@ -30,6 +30,7 @@
 import org.opensearch.ml.common.CommonValue;
 import org.opensearch.ml.common.MLIndex;
 import org.opensearch.ml.common.exception.MLException;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.transport.client.Client;
 
 import lombok.AccessLevel;
@@ -44,6 +45,7 @@ public class MLIndicesHandler {
 
     ClusterService clusterService;
     Client client;
+    MLFeatureEnabledSetting mlFeatureEnabledSetting;
     private static final Map<String, AtomicBoolean> indexMappingUpdated = new HashMap<>();
 
     static {
@@ -52,6 +54,10 @@ public class MLIndicesHandler {
         }
     }
 
+    public static boolean doesMultiTenantIndexExists(ClusterService clusterService, boolean isMultiTenancyEnabled, String indexName) {
+        return isMultiTenancyEnabled || clusterService.state().metadata().hasIndex(indexName);
+    }
+
     public void initModelGroupIndexIfAbsent(ActionListener<Boolean> listener) {
         initMLIndexIfAbsent(MLIndex.MODEL_GROUP, listener);
     }
@@ -105,7 +111,7 @@ public void initMLIndexIfAbsent(MLIndex index, ActionListener<Boolean> listener)
         String mapping = index.getMapping();
         try (ThreadContext.StoredContext threadContext = client.threadPool().getThreadContext().stashContext()) {
             ActionListener<Boolean> internalListener = ActionListener.runBefore(listener, () -> threadContext.restore());
-            if (!clusterService.state().metadata().hasIndex(indexName)) {
+            if (!MLIndicesHandler.doesMultiTenantIndexExists(clusterService, mlFeatureEnabledSetting.isMultiTenancyEnabled(), indexName)) {
                 ActionListener<CreateIndexResponse> actionListener = ActionListener.wrap(r -> {
                     if (r.isAcknowledged()) {
                         log.info("create index:{}", indexName);
diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/indices/MLIndicesHandlerTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/indices/MLIndicesHandlerTest.java
@@ -35,6 +35,7 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.client.AdminClient;
 import org.opensearch.transport.client.Client;
@@ -74,6 +75,9 @@ public class MLIndicesHandlerTest {
     @Mock
     private ThreadPool threadPool;
 
+    @Mock
+    private MLFeatureEnabledSetting mlFeatureEnabledSetting;
+
     Settings settings;
     ThreadContext threadContext;
     MLIndicesHandler indicesHandler;
@@ -102,7 +106,7 @@ public void setUp() {
         threadContext = new ThreadContext(settings);
         when(client.threadPool()).thenReturn(threadPool);
         when(threadPool.getThreadContext()).thenReturn(threadContext);
-        indicesHandler = new MLIndicesHandler(clusterService, client);
+        indicesHandler = new MLIndicesHandler(clusterService, client, mlFeatureEnabledSetting);
     }
 
     @Test
diff --git a/plugin/src/main/java/org/opensearch/ml/action/connector/ExecuteConnectorTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/connector/ExecuteConnectorTransportAction.java
@@ -26,6 +26,7 @@
 import org.opensearch.ml.engine.MLEngineClassLoader;
 import org.opensearch.ml.engine.algorithms.remote.RemoteConnectorExecutor;
 import org.opensearch.ml.engine.encryptor.EncryptorImpl;
+import org.opensearch.ml.engine.indices.MLIndicesHandler;
 import org.opensearch.ml.helper.ConnectorAccessControlHelper;
 import org.opensearch.script.ScriptService;
 import org.opensearch.tasks.Task;
@@ -74,7 +75,8 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLTask
         String connectorId = executeConnectorRequest.getConnectorId();
         String connectorAction = ConnectorAction.ActionType.EXECUTE.name();
 
-        if (clusterService.state().metadata().hasIndex(ML_CONNECTOR_INDEX)) {
+        if (MLIndicesHandler
+            .doesMultiTenantIndexExists(clusterService, mlFeatureEnabledSetting.isMultiTenancyEnabled(), ML_CONNECTOR_INDEX)) {
             ActionListener<Connector> listener = ActionListener.wrap(connector -> {
                 if (connectorAccessControlHelper.validateConnectorAccess(client, connector)) {
                     // adding tenantID as null, because we are not implement multi-tenancy for this feature yet.
diff --git a/plugin/src/main/java/org/opensearch/ml/action/handler/MLSearchHandler.java b/plugin/src/main/java/org/opensearch/ml/action/handler/MLSearchHandler.java
@@ -39,6 +39,8 @@
 import org.opensearch.ml.common.connector.HttpConnector;
 import org.opensearch.ml.common.exception.MLException;
 import org.opensearch.ml.common.exception.MLResourceNotFoundException;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
+import org.opensearch.ml.engine.indices.MLIndicesHandler;
 import org.opensearch.ml.helper.ModelAccessControlHelper;
 import org.opensearch.ml.utils.RestActionUtils;
 import org.opensearch.remote.metadata.client.SdkClient;
@@ -65,17 +67,20 @@ public class MLSearchHandler {
     private ModelAccessControlHelper modelAccessControlHelper;
 
     private ClusterService clusterService;
+    private MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
     public MLSearchHandler(
         Client client,
         NamedXContentRegistry xContentRegistry,
         ModelAccessControlHelper modelAccessControlHelper,
-        ClusterService clusterService
+        ClusterService clusterService,
+        MLFeatureEnabledSetting mlFeatureEnabledSetting
     ) {
         this.modelAccessControlHelper = modelAccessControlHelper;
         this.client = client;
         this.xContentRegistry = xContentRegistry;
         this.clusterService = clusterService;
+        this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
     }
 
     /**
@@ -132,7 +137,12 @@ public void search(SdkClient sdkClient, SearchRequest request, String tenantId,
             final ActionListener<SearchResponse> doubleWrapperListener = ActionListener
                 .wrap(wrappedListener::onResponse, e -> wrapListenerToHandleSearchIndexNotFound(e, wrappedListener));
             if (modelAccessControlHelper.skipModelAccessControl(user)
-                || !clusterService.state().metadata().hasIndex(CommonValue.ML_MODEL_GROUP_INDEX)) {
+                || !MLIndicesHandler
+                    .doesMultiTenantIndexExists(
+                        clusterService,
+                        mlFeatureEnabledSetting.isMultiTenancyEnabled(),
+                        CommonValue.ML_MODEL_GROUP_INDEX
+                    )) {
 
                 SearchDataObjectRequest searchDataObjectRequest = SearchDataObjectRequest
                     .builder()
diff --git a/plugin/src/main/java/org/opensearch/ml/action/tasks/CancelBatchJobTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/tasks/CancelBatchJobTransportAction.java
@@ -55,6 +55,7 @@
 import org.opensearch.ml.engine.algorithms.remote.ConnectorUtils;
 import org.opensearch.ml.engine.algorithms.remote.RemoteConnectorExecutor;
 import org.opensearch.ml.engine.encryptor.EncryptorImpl;
+import org.opensearch.ml.engine.indices.MLIndicesHandler;
 import org.opensearch.ml.helper.ConnectorAccessControlHelper;
 import org.opensearch.ml.helper.ModelAccessControlHelper;
 import org.opensearch.ml.model.MLModelManager;
@@ -199,7 +200,12 @@ private void processRemoteBatchPrediction(MLTask mlTask, ActionListener<MLCancel
                         if (model.getConnector() != null) {
                             Connector connector = model.getConnector();
                             executeConnector(connector, mlInput, actionListener);
-                        } else if (clusterService.state().metadata().hasIndex(ML_CONNECTOR_INDEX)) {
+                        } else if (MLIndicesHandler
+                            .doesMultiTenantIndexExists(
+                                clusterService,
+                                mlFeatureEnabledSetting.isMultiTenancyEnabled(),
+                                ML_CONNECTOR_INDEX
+                            )) {
                             ActionListener<Connector> listener = ActionListener
                                 .wrap(connector -> { executeConnector(connector, mlInput, actionListener); }, e -> {
                                     log.error("Failed to get connector {}", model.getConnectorId(), e);
diff --git a/plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java
@@ -82,6 +82,7 @@
 import org.opensearch.ml.engine.algorithms.remote.ConnectorUtils;
 import org.opensearch.ml.engine.algorithms.remote.RemoteConnectorExecutor;
 import org.opensearch.ml.engine.encryptor.EncryptorImpl;
+import org.opensearch.ml.engine.indices.MLIndicesHandler;
 import org.opensearch.ml.engine.utils.S3Utils;
 import org.opensearch.ml.helper.ConnectorAccessControlHelper;
 import org.opensearch.ml.helper.ModelAccessControlHelper;
@@ -390,7 +391,12 @@ private void processRemoteBatchPrediction(
                                             remoteJob,
                                             actionListener
                                         );
-                                    } else if (clusterService.state().metadata().hasIndex(ML_CONNECTOR_INDEX)) {
+                                    } else if (MLIndicesHandler
+                                        .doesMultiTenantIndexExists(
+                                            clusterService,
+                                            mlFeatureEnabledSetting.isMultiTenancyEnabled(),
+                                            ML_CONNECTOR_INDEX
+                                        )) {
                                         ActionListener<Connector> listener = ActionListener.wrap(connector -> {
                                             executeConnector(
                                                 connector,
diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
@@ -599,7 +599,7 @@ public Collection<Object> createComponents(
         Settings settings = environment.settings();
         Path dataPath = environment.dataFiles()[0];
 
-        mlIndicesHandler = new MLIndicesHandler(clusterService, client);
+        mlIndicesHandler = new MLIndicesHandler(clusterService, client, mlFeatureEnabledSetting);
 
         SdkClient sdkClient = SdkClientFactory
             .createSdkClient(
@@ -802,7 +802,13 @@ public Collection<Object> createComponents(
         MLToolExecutor toolExecutor = new MLToolExecutor(client, sdkClient, settings, clusterService, xContentRegistry, toolFactories);
         MLEngineClassLoader.register(FunctionName.TOOL, toolExecutor);
 
-        MLSearchHandler mlSearchHandler = new MLSearchHandler(client, xContentRegistry, modelAccessControlHelper, clusterService);
+        MLSearchHandler mlSearchHandler = new MLSearchHandler(
+            client,
+            xContentRegistry,
+            modelAccessControlHelper,
+            clusterService,
+            mlFeatureEnabledSetting
+        );
         MLModelAutoReDeployer mlModelAutoRedeployer = new MLModelAutoReDeployer(
             clusterService,
             client,
diff --git a/plugin/src/test/java/org/opensearch/ml/action/models/SearchModelTransportActionTests.java b/plugin/src/test/java/org/opensearch/ml/action/models/SearchModelTransportActionTests.java
@@ -114,7 +114,9 @@ public class SearchModelTransportActionTests extends OpenSearchTestCase {
     public void setup() {
         MockitoAnnotations.openMocks(this);
         sdkClient = SdkClientFactory.createSdkClient(client, NamedXContentRegistry.EMPTY, Collections.emptyMap());
-        mlSearchHandler = spy(new MLSearchHandler(client, namedXContentRegistry, modelAccessControlHelper, clusterService));
+        mlSearchHandler = spy(
+            new MLSearchHandler(client, namedXContentRegistry, modelAccessControlHelper, clusterService, mlFeatureEnabledSetting)
+        );
         searchModelTransportAction = new SearchModelTransportAction(
             transportService,
             actionFilters,
