Skip to content

Commit d59fced

Browse files
b4sjoob4sjoo
authored
Add feature enable setting for controller index (opensearch-project#2971)
* Add feature enable setting for controller index Signed-off-by: b4sjoo <sicheng.song@outlook.com> * Fix IT Signed-off-by: b4sjoo <sicheng.song@outlook.com> * Fix UT Signed-off-by: b4sjoo <sicheng.song@outlook.com> --------- Signed-off-by: b4sjoo <sicheng.song@outlook.com>
1 parent 66d8e2b commit d59fced

21 files changed

+263
-28
lines changed

‎plugin/src/main/java/org/opensearch/ml/action/controller/CreateControllerTransportAction.java

+9-1
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import static org.opensearch.ml.common.FunctionName.REMOTE;
1010
import static org.opensearch.ml.common.FunctionName.TEXT_EMBEDDING;
1111
import static org.opensearch.ml.common.utils.StringUtils.getErrorMessage;
12+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1213

1314
import java.util.ArrayList;
1415
import java.util.Arrays;
@@ -51,6 +52,7 @@
5152
import org.opensearch.ml.helper.ModelAccessControlHelper;
5253
import org.opensearch.ml.model.MLModelCacheHelper;
5354
import org.opensearch.ml.model.MLModelManager;
55+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
5456
import org.opensearch.ml.utils.RestActionUtils;
5557
import org.opensearch.tasks.Task;
5658
import org.opensearch.transport.TransportService;
@@ -68,6 +70,7 @@ public class CreateControllerTransportAction extends HandledTransportAction<Acti
6870
ClusterService clusterService;
6971
MLModelCacheHelper mlModelCacheHelper;
7072
ModelAccessControlHelper modelAccessControlHelper;
73+
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
7174

7275
@Inject
7376
public CreateControllerTransportAction(
@@ -78,7 +81,8 @@ public CreateControllerTransportAction(
7881
ClusterService clusterService,
7982
ModelAccessControlHelper modelAccessControlHelper,
8083
MLModelCacheHelper mlModelCacheHelper,
81-
MLModelManager mlModelManager
84+
MLModelManager mlModelManager,
85+
MLFeatureEnabledSetting mlFeatureEnabledSetting
8286
) {
8387
super(MLCreateControllerAction.NAME, transportService, actionFilters, MLCreateControllerRequest::new);
8488
this.mlIndicesHandler = mlIndicesHandler;
@@ -87,6 +91,7 @@ public CreateControllerTransportAction(
8791
this.clusterService = clusterService;
8892
this.mlModelCacheHelper = mlModelCacheHelper;
8993
this.modelAccessControlHelper = modelAccessControlHelper;
94+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
9095
}
9196

9297
@Override
@@ -98,6 +103,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLCrea
98103
String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
99104

100105
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
106+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
107+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
108+
}
101109
ActionListener<MLCreateControllerResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
102110
// TODO: Add support for multi tenancy
103111
mlModelManager.getModel(modelId, null, null, excludes, ActionListener.wrap(mlModel -> {

‎plugin/src/main/java/org/opensearch/ml/action/controller/DeleteControllerTransportAction.java

+9-1
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
99
import static org.opensearch.ml.common.utils.StringUtils.getErrorMessage;
10+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1011

1112
import java.util.ArrayList;
1213
import java.util.Arrays;
@@ -40,6 +41,7 @@
4041
import org.opensearch.ml.helper.ModelAccessControlHelper;
4142
import org.opensearch.ml.model.MLModelCacheHelper;
4243
import org.opensearch.ml.model.MLModelManager;
44+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
4345
import org.opensearch.ml.utils.RestActionUtils;
4446
import org.opensearch.tasks.Task;
4547
import org.opensearch.transport.TransportService;
@@ -57,6 +59,7 @@ public class DeleteControllerTransportAction extends HandledTransportAction<Acti
5759
MLModelManager mlModelManager;
5860
MLModelCacheHelper mlModelCacheHelper;
5961
ModelAccessControlHelper modelAccessControlHelper;
62+
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
6063

6164
@Inject
6265
public DeleteControllerTransportAction(
@@ -67,7 +70,8 @@ public DeleteControllerTransportAction(
6770
ClusterService clusterService,
6871
MLModelManager mlModelManager,
6972
MLModelCacheHelper mlModelCacheHelper,
70-
ModelAccessControlHelper modelAccessControlHelper
73+
ModelAccessControlHelper modelAccessControlHelper,
74+
MLFeatureEnabledSetting mlFeatureEnabledSetting
7175
) {
7276
super(MLControllerDeleteAction.NAME, transportService, actionFilters, MLControllerDeleteRequest::new);
7377
this.client = client;
@@ -76,6 +80,7 @@ public DeleteControllerTransportAction(
7680
this.mlModelManager = mlModelManager;
7781
this.mlModelCacheHelper = mlModelCacheHelper;
7882
this.modelAccessControlHelper = modelAccessControlHelper;
83+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
7984
}
8085

8186
@Override
@@ -85,6 +90,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
8590
User user = RestActionUtils.getUserContext(client);
8691
String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
8792
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
93+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
94+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
95+
}
8896
ActionListener<DeleteResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
8997
// TODO: Add support for multi tenancy
9098
mlModelManager.getModel(modelId, null, null, excludes, ActionListener.wrap(mlModel -> {

‎plugin/src/main/java/org/opensearch/ml/action/controller/GetControllerTransportAction.java

+9-1
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
99
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
1010
import static org.opensearch.ml.common.utils.StringUtils.getErrorMessage;
11+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1112
import static org.opensearch.ml.utils.MLNodeUtils.createXContentParserFromRegistry;
1213
import static org.opensearch.ml.utils.RestActionUtils.getFetchSourceContext;
1314

@@ -33,6 +34,7 @@
3334
import org.opensearch.ml.common.transport.controller.MLControllerGetResponse;
3435
import org.opensearch.ml.helper.ModelAccessControlHelper;
3536
import org.opensearch.ml.model.MLModelManager;
37+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
3638
import org.opensearch.ml.utils.RestActionUtils;
3739
import org.opensearch.search.fetch.subphase.FetchSourceContext;
3840
import org.opensearch.tasks.Task;
@@ -50,6 +52,7 @@ public class GetControllerTransportAction extends HandledTransportAction<ActionR
5052
ClusterService clusterService;
5153
MLModelManager mlModelManager;
5254
ModelAccessControlHelper modelAccessControlHelper;
55+
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
5356

5457
@Inject
5558
public GetControllerTransportAction(
@@ -59,14 +62,16 @@ public GetControllerTransportAction(
5962
NamedXContentRegistry xContentRegistry,
6063
ClusterService clusterService,
6164
MLModelManager mlModelManager,
62-
ModelAccessControlHelper modelAccessControlHelper
65+
ModelAccessControlHelper modelAccessControlHelper,
66+
MLFeatureEnabledSetting mlFeatureEnabledSetting
6367
) {
6468
super(MLControllerGetAction.NAME, transportService, actionFilters, MLControllerGetRequest::new);
6569
this.client = client;
6670
this.xContentRegistry = xContentRegistry;
6771
this.clusterService = clusterService;
6872
this.mlModelManager = mlModelManager;
6973
this.modelAccessControlHelper = modelAccessControlHelper;
74+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
7075
}
7176

7277
@Override
@@ -79,6 +84,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLCont
7984
String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
8085

8186
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
87+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
88+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
89+
}
8290
ActionListener<MLControllerGetResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
8391
client.get(getRequest, ActionListener.wrap(r -> {
8492
if (r != null && r.isExists()) {

‎plugin/src/main/java/org/opensearch/ml/action/controller/UpdateControllerTransportAction.java

+9-1
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import static org.opensearch.ml.common.FunctionName.REMOTE;
1010
import static org.opensearch.ml.common.FunctionName.TEXT_EMBEDDING;
1111
import static org.opensearch.ml.common.utils.StringUtils.getErrorMessage;
12+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1213

1314
import java.util.ArrayList;
1415
import java.util.Arrays;
@@ -46,6 +47,7 @@
4647
import org.opensearch.ml.helper.ModelAccessControlHelper;
4748
import org.opensearch.ml.model.MLModelCacheHelper;
4849
import org.opensearch.ml.model.MLModelManager;
50+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
4951
import org.opensearch.ml.utils.RestActionUtils;
5052
import org.opensearch.tasks.Task;
5153
import org.opensearch.transport.TransportService;
@@ -62,6 +64,7 @@ public class UpdateControllerTransportAction extends HandledTransportAction<Acti
6264
MLModelCacheHelper mlModelCacheHelper;
6365
ClusterService clusterService;
6466
ModelAccessControlHelper modelAccessControlHelper;
67+
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
6568

6669
@Inject
6770
public UpdateControllerTransportAction(
@@ -71,14 +74,16 @@ public UpdateControllerTransportAction(
7174
ClusterService clusterService,
7275
ModelAccessControlHelper modelAccessControlHelper,
7376
MLModelCacheHelper mlModelCacheHelper,
74-
MLModelManager mlModelManager
77+
MLModelManager mlModelManager,
78+
MLFeatureEnabledSetting mlFeatureEnabledSetting
7579
) {
7680
super(MLUpdateControllerAction.NAME, transportService, actionFilters, MLUpdateControllerRequest::new);
7781
this.client = client;
7882
this.mlModelManager = mlModelManager;
7983
this.clusterService = clusterService;
8084
this.mlModelCacheHelper = mlModelCacheHelper;
8185
this.modelAccessControlHelper = modelAccessControlHelper;
86+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
8287
}
8388

8489
@Override
@@ -90,6 +95,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Update
9095
String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
9196

9297
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
98+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
99+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
100+
}
93101
ActionListener<UpdateResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
94102
// TODO: Add support for multi tenancy
95103
mlModelManager.getModel(modelId, null, null, excludes, ActionListener.wrap(mlModel -> {

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

+5-4
Original file line numberDiff line numberDiff line change
@@ -750,10 +750,10 @@ public List<RestHandler> getRestHandlers(
750750
RestMemorySearchInteractionsAction restSearchInteractionsAction = new RestMemorySearchInteractionsAction();
751751
RestMemoryGetConversationAction restGetConversationAction = new RestMemoryGetConversationAction();
752752
RestMemoryGetInteractionAction restGetInteractionAction = new RestMemoryGetInteractionAction();
753-
RestMLCreateControllerAction restMLCreateControllerAction = new RestMLCreateControllerAction();
754-
RestMLGetControllerAction restMLGetControllerAction = new RestMLGetControllerAction();
755-
RestMLUpdateControllerAction restMLUpdateControllerAction = new RestMLUpdateControllerAction();
756-
RestMLDeleteControllerAction restMLDeleteControllerAction = new RestMLDeleteControllerAction();
753+
RestMLCreateControllerAction restMLCreateControllerAction = new RestMLCreateControllerAction(mlFeatureEnabledSetting);
754+
RestMLGetControllerAction restMLGetControllerAction = new RestMLGetControllerAction(mlFeatureEnabledSetting);
755+
RestMLUpdateControllerAction restMLUpdateControllerAction = new RestMLUpdateControllerAction(mlFeatureEnabledSetting);
756+
RestMLDeleteControllerAction restMLDeleteControllerAction = new RestMLDeleteControllerAction(mlFeatureEnabledSetting);
757757
RestMLGetAgentAction restMLGetAgentAction = new RestMLGetAgentAction(mlFeatureEnabledSetting);
758758
RestMLDeleteAgentAction restMLDeleteAgentAction = new RestMLDeleteAgentAction(mlFeatureEnabledSetting);
759759
RestMemoryUpdateConversationAction restMemoryUpdateConversationAction = new RestMemoryUpdateConversationAction();
@@ -946,6 +946,7 @@ public List<Setting<?>> getSettings() {
946946
MLCommonsSettings.ML_COMMONS_AGENT_FRAMEWORK_ENABLED,
947947
MLCommonsSettings.ML_COMMONS_MODEL_AUTO_DEPLOY_ENABLE,
948948
MLCommonsSettings.ML_COMMONS_MULTI_TENANCY_ENABLED,
949+
MLCommonsSettings.ML_COMMONS_CONTROLLER_ENABLED,
949950
// Settings for SdkClient
950951
SdkClientSettings.REMOTE_METADATA_TYPE,
951952
SdkClientSettings.REMOTE_METADATA_ENDPOINT,

‎plugin/src/main/java/org/opensearch/ml/rest/RestMLCreateControllerAction.java

+10-1
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
99
import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
10+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1011
import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_MODEL_ID;
1112
import static org.opensearch.ml.utils.RestActionUtils.getParameterId;
1213

@@ -20,6 +21,7 @@
2021
import org.opensearch.ml.common.controller.MLController;
2122
import org.opensearch.ml.common.transport.controller.MLCreateControllerAction;
2223
import org.opensearch.ml.common.transport.controller.MLCreateControllerRequest;
24+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
2325
import org.opensearch.rest.BaseRestHandler;
2426
import org.opensearch.rest.RestRequest;
2527
import org.opensearch.rest.action.RestToXContentListener;
@@ -29,11 +31,14 @@
2931
public class RestMLCreateControllerAction extends BaseRestHandler {
3032

3133
public final static String ML_CREATE_CONTROLLER_ACTION = "ml_create_controller_action";
34+
private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
3235

3336
/**
3437
* Constructor
3538
*/
36-
public RestMLCreateControllerAction() {}
39+
public RestMLCreateControllerAction(MLFeatureEnabledSetting mlFeatureEnabledSetting) {
40+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
41+
}
3742

3843
@Override
3944
public String getName() {
@@ -61,6 +66,10 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
6166
* @return MLCreateControllerRequest
6267
*/
6368
private MLCreateControllerRequest getRequest(RestRequest request) throws IOException {
69+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
70+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
71+
}
72+
6473
if (!request.hasContent()) {
6574
throw new OpenSearchParseException("Create model controller request has empty body");
6675
}

‎plugin/src/main/java/org/opensearch/ml/rest/RestMLDeleteControllerAction.java

+10-1
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
package org.opensearch.ml.rest;
77

88
import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
9+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
910
import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_MODEL_ID;
1011

1112
import java.io.IOException;
@@ -15,6 +16,7 @@
1516
import org.opensearch.client.node.NodeClient;
1617
import org.opensearch.ml.common.transport.controller.MLControllerDeleteAction;
1718
import org.opensearch.ml.common.transport.controller.MLControllerDeleteRequest;
19+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
1820
import org.opensearch.rest.BaseRestHandler;
1921
import org.opensearch.rest.RestRequest;
2022
import org.opensearch.rest.action.RestToXContentListener;
@@ -25,9 +27,13 @@
2527
* This class consists of the REST handler to delete ML Model.
2628
*/
2729
public class RestMLDeleteControllerAction extends BaseRestHandler {
30+
2831
private static final String ML_DELETE_CONTROLLER_ACTION = "ml_delete_controller_action";
32+
private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
2933

30-
public void RestMLDeleteControllerAction() {}
34+
public RestMLDeleteControllerAction(MLFeatureEnabledSetting mlFeatureEnabledSetting) {
35+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
36+
}
3137

3238
@Override
3339
public String getName() {
@@ -42,6 +48,9 @@ public List<Route> routes() {
4248

4349
@Override
4450
protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
51+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
52+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
53+
}
4554
String modelId = request.param(PARAMETER_MODEL_ID);
4655

4756
MLControllerDeleteRequest mlControllerDeleteRequest = new MLControllerDeleteRequest(modelId);

‎plugin/src/main/java/org/opensearch/ml/rest/RestMLGetControllerAction.java

+11-1
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
package org.opensearch.ml.rest;
77

88
import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
9+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
910
import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_MODEL_ID;
1011
import static org.opensearch.ml.utils.RestActionUtils.getParameterId;
1112
import static org.opensearch.ml.utils.RestActionUtils.returnContent;
@@ -17,6 +18,7 @@
1718
import org.opensearch.client.node.NodeClient;
1819
import org.opensearch.ml.common.transport.controller.MLControllerGetAction;
1920
import org.opensearch.ml.common.transport.controller.MLControllerGetRequest;
21+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
2022
import org.opensearch.rest.BaseRestHandler;
2123
import org.opensearch.rest.RestRequest;
2224
import org.opensearch.rest.action.RestToXContentListener;
@@ -25,12 +27,16 @@
2527
import com.google.common.collect.ImmutableList;
2628

2729
public class RestMLGetControllerAction extends BaseRestHandler {
30+
2831
private static final String ML_GET_CONTROLLER_ACTION = "ml_get_controller_action";
32+
private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
2933

3034
/**
3135
* Constructor
3236
*/
33-
public RestMLGetControllerAction() {}
37+
public RestMLGetControllerAction(MLFeatureEnabledSetting mlFeatureEnabledSetting) {
38+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
39+
}
3440

3541
@Override
3642
public String getName() {
@@ -57,6 +63,10 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
5763
*/
5864
@VisibleForTesting
5965
MLControllerGetRequest getRequest(RestRequest request) throws IOException {
66+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
67+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
68+
}
69+
6070
String modelId = getParameterId(request, PARAMETER_MODEL_ID);
6171
boolean returnContent = returnContent(request);
6272

0 commit comments

Comments
 (0)