Skip to content

Commit 57a59fc

Browse files
brianf-awsbrianf-aws
committed
pass isMultiTenancyEnabled across classes to early return index search
Signed-off-by: Brian Flores <[email protected]>
1 parent 937f3fe commit 57a59fc

File tree

10 files changed

+105
-15
lines changed

10 files changed

+105
-15
lines changed

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
import org.opensearch.ml.common.spi.tools.Tool;
4949
import org.opensearch.ml.engine.Executable;
5050
import org.opensearch.ml.engine.annotation.Function;
51+
import org.opensearch.ml.engine.indices.MLIndicesHandler;
5152
import org.opensearch.ml.engine.memory.ConversationIndexMemory;
5253
import org.opensearch.ml.engine.memory.ConversationIndexMessage;
5354
import org.opensearch.ml.memory.action.conversation.CreateInteractionResponse;
@@ -142,7 +143,7 @@ public void execute(Input input, ActionListener<Output> listener) {
142143
.fetchSourceContext(fetchSourceContext)
143144
.build();
144145

145-
if (clusterService.state().metadata().hasIndex(ML_AGENT_INDEX)) {
146+
if (MLIndicesHandler.doesMultiTenantIndexExist(clusterService, isMultiTenancyEnabled, ML_AGENT_INDEX)) {
146147
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
147148
sdkClient
148149
.getDataObjectAsync(getDataObjectRequest, client.threadPool().executor("opensearch_ml_general"))

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/indices/MLIndicesHandler.java‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
3232
import org.opensearch.ml.common.MLIndex;
3333
import org.opensearch.ml.common.exception.MLException;
3434

35+
import com.google.common.annotations.VisibleForTesting;
36+
3537
import lombok.AccessLevel;
3638
import lombok.RequiredArgsConstructor;
3739
import lombok.experimental.FieldDefaults;
@@ -41,9 +43,9 @@
4143
@RequiredArgsConstructor
4244
@Log4j2
4345
public class MLIndicesHandler {
44-
4546
ClusterService clusterService;
4647
Client client;
48+
boolean isMultiTenancyEnabled;
4749
private static final Map<String, AtomicBoolean> indexMappingUpdated = new HashMap<>();
4850

4951
static {
@@ -93,7 +95,7 @@ public void initMLIndexIfAbsent(MLIndex index, ActionListener<Boolean> listener)
9395
String mapping = index.getMapping();
9496
try (ThreadContext.StoredContext threadContext = client.threadPool().getThreadContext().stashContext()) {
9597
ActionListener<Boolean> internalListener = ActionListener.runBefore(listener, () -> threadContext.restore());
96-
if (!clusterService.state().metadata().hasIndex(indexName)) {
98+
if (!MLIndicesHandler.doesMultiTenantIndexExist(clusterService, isMultiTenancyEnabled, indexName)) {
9799
ActionListener<CreateIndexResponse> actionListener = ActionListener.wrap(r -> {
98100
if (r.isAcknowledged()) {
99101
log.info("create index:{}", indexName);
@@ -208,4 +210,23 @@ public void shouldUpdateIndex(String indexName, Integer newVersion, ActionListen
208210
listener.onResponse(newVersion > oldVersion);
209211
}
210212

213+
/**
214+
* Determines whether an index exists on non-multi tenancy enabled environments. Otherwise,
215+
* returns true when multiTenancy is Enabled
216+
*
217+
* @param clusterService the cluster service
218+
* @param isMultiTenancyEnabled whether multi-tenancy is enabled
219+
* @param indexName - the index to search
220+
* @return boolean indicating the existence of an index. Returns true if multitenancy is enabled.
221+
* @implNote This method assumes if your environment enables multi tenancy, then your plugin indices are
222+
* pre-populated. If this is incorrect, it will result in unwanted early returns without checking the clusterService.
223+
*/
224+
public static boolean doesMultiTenantIndexExist(ClusterService clusterService, boolean isMultiTenancyEnabled, String indexName) {
225+
return isMultiTenancyEnabled || clusterService.state().metadata().hasIndex(indexName);
226+
}
227+
228+
@VisibleForTesting
229+
public boolean doesIndexExist(String indexName) {
230+
return MLIndicesHandler.doesMultiTenantIndexExist(clusterService, isMultiTenancyEnabled, indexName);
231+
}
211232
}

‎ml-algorithms/src/test/java/org/opensearch/ml/engine/indices/MLIndicesHandlerTest.java‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package org.opensearch.ml.engine.indices;
22

33
import static org.junit.Assert.assertEquals;
4+
import static org.junit.Assert.assertFalse;
5+
import static org.junit.Assert.assertTrue;
46
import static org.mockito.ArgumentMatchers.any;
57
import static org.mockito.ArgumentMatchers.anyString;
68
import static org.mockito.ArgumentMatchers.isA;
@@ -11,6 +13,7 @@
1113
import static org.mockito.Mockito.when;
1214
import static org.opensearch.ml.common.CommonValue.META;
1315
import static org.opensearch.ml.common.CommonValue.ML_AGENT_INDEX;
16+
import static org.opensearch.ml.common.CommonValue.ML_CONFIG_INDEX;
1417
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX;
1518
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_META_INDEX;
1619
import static org.opensearch.ml.common.CommonValue.SCHEMA_VERSION_FIELD;
@@ -77,6 +80,8 @@ public class MLIndicesHandlerTest {
7780
ThreadContext threadContext;
7881
MLIndicesHandler indicesHandler;
7982

83+
private boolean isMultiTenancyEnabled = false;
84+
8085
@Before
8186
public void setUp() {
8287
MockitoAnnotations.openMocks(this);
@@ -101,7 +106,7 @@ public void setUp() {
101106
threadContext = new ThreadContext(settings);
102107
when(client.threadPool()).thenReturn(threadPool);
103108
when(threadPool.getThreadContext()).thenReturn(threadContext);
104-
indicesHandler = new MLIndicesHandler(clusterService, client);
109+
indicesHandler = new MLIndicesHandler(clusterService, client, isMultiTenancyEnabled);
105110
}
106111

107112
@Test
@@ -216,4 +221,26 @@ public void initMLConnectorIndex_ResourceAlreadyExistsException_RaceCondition()
216221
verify(listener).onResponse(argumentCaptor.capture());
217222
assertEquals(true, argumentCaptor.getValue());
218223
}
224+
225+
@Test
226+
public void doesMultiTenantIndexExist_multiTenancyEnabled_returnsTrue() {
227+
assertTrue(MLIndicesHandler.doesMultiTenantIndexExist(null, true, null));
228+
MLIndicesHandler mlIndicesHandler = new MLIndicesHandler(clusterService, client, true);
229+
assertTrue(mlIndicesHandler.doesIndexExist(ML_CONFIG_INDEX));
230+
}
231+
232+
@Test
233+
public void doesMultiTenantIndexExist_multiTenancyDisabledSearchesClusterService_returnsValidSearchResult() {
234+
assertFalse(MLIndicesHandler.doesMultiTenantIndexExist(clusterService, false, null));
235+
236+
String sampleIndexName = "test-index";
237+
MLIndicesHandler mlIndicesHandler = new MLIndicesHandler(clusterService, client, false);
238+
239+
when(clusterService.state().metadata().hasIndex(sampleIndexName)).thenReturn(true);
240+
assertTrue(mlIndicesHandler.doesIndexExist(sampleIndexName));
241+
242+
when(clusterService.state().metadata().hasIndex(sampleIndexName)).thenReturn(false);
243+
assertFalse(mlIndicesHandler.doesIndexExist(sampleIndexName));
244+
}
245+
219246
}

‎plugin/src/main/java/org/opensearch/ml/action/connector/ExecuteConnectorTransportAction.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,9 @@
2626
import org.opensearch.ml.engine.MLEngineClassLoader;
2727
import org.opensearch.ml.engine.algorithms.remote.RemoteConnectorExecutor;
2828
import org.opensearch.ml.engine.encryptor.EncryptorImpl;
29+
import org.opensearch.ml.engine.indices.MLIndicesHandler;
2930
import org.opensearch.ml.helper.ConnectorAccessControlHelper;
31+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
3032
import org.opensearch.script.ScriptService;
3133
import org.opensearch.tasks.Task;
3234
import org.opensearch.transport.TransportService;
@@ -43,6 +45,7 @@ public class ExecuteConnectorTransportAction extends HandledTransportAction<Acti
4345

4446
ConnectorAccessControlHelper connectorAccessControlHelper;
4547
EncryptorImpl encryptor;
48+
MLFeatureEnabledSetting mlFeatureEnabledSetting;
4649

4750
@Inject
4851
public ExecuteConnectorTransportAction(
@@ -53,7 +56,8 @@ public ExecuteConnectorTransportAction(
5356
ScriptService scriptService,
5457
NamedXContentRegistry xContentRegistry,
5558
ConnectorAccessControlHelper connectorAccessControlHelper,
56-
EncryptorImpl encryptor
59+
EncryptorImpl encryptor,
60+
MLFeatureEnabledSetting mlFeatureEnabledSetting
5761
) {
5862
super(MLExecuteConnectorAction.NAME, transportService, actionFilters, MLConnectorDeleteRequest::new);
5963
this.client = client;
@@ -62,6 +66,7 @@ public ExecuteConnectorTransportAction(
6266
this.xContentRegistry = xContentRegistry;
6367
this.connectorAccessControlHelper = connectorAccessControlHelper;
6468
this.encryptor = encryptor;
69+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
6570
}
6671

6772
@Override
@@ -70,7 +75,8 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLTask
7075
String connectorId = executeConnectorRequest.getConnectorId();
7176
String connectorAction = ConnectorAction.ActionType.EXECUTE.name();
7277

73-
if (clusterService.state().metadata().hasIndex(ML_CONNECTOR_INDEX)) {
78+
if (MLIndicesHandler
79+
.doesMultiTenantIndexExist(clusterService, mlFeatureEnabledSetting.isMultiTenancyEnabled(), ML_CONNECTOR_INDEX)) {
7480
ActionListener<Connector> listener = ActionListener.wrap(connector -> {
7581
if (connectorAccessControlHelper.validateConnectorAccess(client, connector)) {
7682
// adding tenantID as null, because we are not implement multi-tenancy for this feature yet.

‎plugin/src/main/java/org/opensearch/ml/action/handler/MLSearchHandler.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import org.opensearch.ml.common.connector.HttpConnector;
4141
import org.opensearch.ml.common.exception.MLException;
4242
import org.opensearch.ml.common.exception.MLResourceNotFoundException;
43+
import org.opensearch.ml.engine.indices.MLIndicesHandler;
4344
import org.opensearch.ml.helper.ModelAccessControlHelper;
4445
import org.opensearch.ml.utils.RestActionUtils;
4546
import org.opensearch.remote.metadata.client.SdkClient;
@@ -59,6 +60,7 @@
5960
@Log4j2
6061
public class MLSearchHandler {
6162
private final Client client;
63+
private final boolean isMultiTenancyEnabled;
6264
private NamedXContentRegistry xContentRegistry;
6365

6466
private ModelAccessControlHelper modelAccessControlHelper;
@@ -69,12 +71,14 @@ public MLSearchHandler(
6971
Client client,
7072
NamedXContentRegistry xContentRegistry,
7173
ModelAccessControlHelper modelAccessControlHelper,
72-
ClusterService clusterService
74+
ClusterService clusterService,
75+
boolean isMultiTenancyEnabled
7376
) {
7477
this.modelAccessControlHelper = modelAccessControlHelper;
7578
this.client = client;
7679
this.xContentRegistry = xContentRegistry;
7780
this.clusterService = clusterService;
81+
this.isMultiTenancyEnabled = isMultiTenancyEnabled;
7882
}
7983

8084
/**
@@ -131,7 +135,7 @@ public void search(SdkClient sdkClient, SearchRequest request, String tenantId,
131135
final ActionListener<SearchResponse> doubleWrapperListener = ActionListener
132136
.wrap(wrappedListener::onResponse, e -> wrapListenerToHandleSearchIndexNotFound(e, wrappedListener));
133137
if (modelAccessControlHelper.skipModelAccessControl(user)
134-
|| !clusterService.state().metadata().hasIndex(CommonValue.ML_MODEL_GROUP_INDEX)) {
138+
|| !MLIndicesHandler.doesMultiTenantIndexExist(clusterService, isMultiTenancyEnabled, CommonValue.ML_MODEL_GROUP_INDEX)) {
135139

136140
SearchDataObjectRequest searchDataObjectRequest = SearchDataObjectRequest
137141
.builder()

‎plugin/src/main/java/org/opensearch/ml/action/tasks/CancelBatchJobTransportAction.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
import org.opensearch.ml.engine.algorithms.remote.ConnectorUtils;
5656
import org.opensearch.ml.engine.algorithms.remote.RemoteConnectorExecutor;
5757
import org.opensearch.ml.engine.encryptor.EncryptorImpl;
58+
import org.opensearch.ml.engine.indices.MLIndicesHandler;
5859
import org.opensearch.ml.helper.ConnectorAccessControlHelper;
5960
import org.opensearch.ml.helper.ModelAccessControlHelper;
6061
import org.opensearch.ml.model.MLModelManager;
@@ -199,7 +200,12 @@ private void processRemoteBatchPrediction(MLTask mlTask, ActionListener<MLCancel
199200
if (model.getConnector() != null) {
200201
Connector connector = model.getConnector();
201202
executeConnector(connector, mlInput, actionListener);
202-
} else if (clusterService.state().metadata().hasIndex(ML_CONNECTOR_INDEX)) {
203+
} else if (MLIndicesHandler
204+
.doesMultiTenantIndexExist(
205+
clusterService,
206+
mlFeatureEnabledSetting.isMultiTenancyEnabled(),
207+
ML_CONNECTOR_INDEX
208+
)) {
203209
ActionListener<Connector> listener = ActionListener
204210
.wrap(connector -> { executeConnector(connector, mlInput, actionListener); }, e -> {
205211
log.error("Failed to get connector {}", model.getConnectorId(), e);

‎plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@
8282
import org.opensearch.ml.engine.algorithms.remote.ConnectorUtils;
8383
import org.opensearch.ml.engine.algorithms.remote.RemoteConnectorExecutor;
8484
import org.opensearch.ml.engine.encryptor.EncryptorImpl;
85+
import org.opensearch.ml.engine.indices.MLIndicesHandler;
8586
import org.opensearch.ml.engine.utils.S3Utils;
8687
import org.opensearch.ml.helper.ConnectorAccessControlHelper;
8788
import org.opensearch.ml.helper.ModelAccessControlHelper;
@@ -390,7 +391,12 @@ private void processRemoteBatchPrediction(
390391
remoteJob,
391392
actionListener
392393
);
393-
} else if (clusterService.state().metadata().hasIndex(ML_CONNECTOR_INDEX)) {
394+
} else if (MLIndicesHandler
395+
.doesMultiTenantIndexExist(
396+
clusterService,
397+
mlFeatureEnabledSetting.isMultiTenancyEnabled(),
398+
ML_CONNECTOR_INDEX
399+
)) {
394400
ActionListener<Connector> listener = ActionListener.wrap(connector -> {
395401
executeConnector(
396402
connector,

‎plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -497,7 +497,9 @@ public Collection<Object> createComponents(
497497
Path dataPath = environment.dataFiles()[0];
498498
Path configFile = environment.configFile();
499499

500-
mlIndicesHandler = new MLIndicesHandler(clusterService, client);
500+
mlFeatureEnabledSetting = new MLFeatureEnabledSetting(clusterService, settings);
501+
502+
mlIndicesHandler = new MLIndicesHandler(clusterService, client, mlFeatureEnabledSetting.isMultiTenancyEnabled());
501503

502504
SdkClient sdkClient = SdkClientFactory
503505
.createSdkClient(
@@ -562,7 +564,6 @@ public Collection<Object> createComponents(
562564
mlInputDatasetHandler = new MLInputDatasetHandler(client);
563565
modelAccessControlHelper = new ModelAccessControlHelper(clusterService, settings);
564566
connectorAccessControlHelper = new ConnectorAccessControlHelper(clusterService, settings);
565-
mlFeatureEnabledSetting = new MLFeatureEnabledSetting(clusterService, settings);
566567
mlModelManager = new MLModelManager(
567568
clusterService,
568569
scriptService,
@@ -689,7 +690,13 @@ public Collection<Object> createComponents(
689690

690691
MetricsCorrelation metricsCorrelation = new MetricsCorrelation(client, settings, clusterService);
691692
MLEngineClassLoader.register(FunctionName.METRICS_CORRELATION, metricsCorrelation);
692-
MLSearchHandler mlSearchHandler = new MLSearchHandler(client, xContentRegistry, modelAccessControlHelper, clusterService);
693+
MLSearchHandler mlSearchHandler = new MLSearchHandler(
694+
client,
695+
xContentRegistry,
696+
modelAccessControlHelper,
697+
clusterService,
698+
mlFeatureEnabledSetting.isMultiTenancyEnabled()
699+
);
693700
MLModelAutoReDeployer mlModelAutoRedeployer = new MLModelAutoReDeployer(
694701
clusterService,
695702
client,

‎plugin/src/test/java/org/opensearch/ml/action/connector/ExecuteConnectorTransportActionTests.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
import org.opensearch.ml.common.transport.connector.MLExecuteConnectorRequest;
3636
import org.opensearch.ml.engine.encryptor.EncryptorImpl;
3737
import org.opensearch.ml.helper.ConnectorAccessControlHelper;
38+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
3839
import org.opensearch.script.ScriptService;
3940
import org.opensearch.tasks.Task;
4041
import org.opensearch.test.OpenSearchTestCase;
@@ -75,6 +76,8 @@ public class ExecuteConnectorTransportActionTests extends OpenSearchTestCase {
7576
private Task task;
7677
@Mock
7778
ThreadPool threadPool;
79+
@Mock
80+
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
7881

7982
@Before
8083
public void setup() {
@@ -109,7 +112,8 @@ public void setup() {
109112
scriptService,
110113
xContentRegistry,
111114
connectorAccessControlHelper,
112-
encryptor
115+
encryptor,
116+
mlFeatureEnabledSetting
113117
);
114118
}
115119

‎plugin/src/test/java/org/opensearch/ml/action/models/SearchModelTransportActionTests.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,15 @@ public class SearchModelTransportActionTests extends OpenSearchTestCase {
114114
public void setup() {
115115
MockitoAnnotations.openMocks(this);
116116
sdkClient = SdkClientFactory.createSdkClient(client, NamedXContentRegistry.EMPTY, Collections.emptyMap());
117-
mlSearchHandler = spy(new MLSearchHandler(client, namedXContentRegistry, modelAccessControlHelper, clusterService));
117+
mlSearchHandler = spy(
118+
new MLSearchHandler(
119+
client,
120+
namedXContentRegistry,
121+
modelAccessControlHelper,
122+
clusterService,
123+
mlFeatureEnabledSetting.isMultiTenancyEnabled()
124+
)
125+
);
118126
searchModelTransportAction = new SearchModelTransportAction(
119127
transportService,
120128
actionFilters,

0 commit comments

Comments
 (0)