Skip to content

Commit 2d3b021

Browse files
Adds existing feature-flag in conjunction with new feature flage
Signed-off-by: Darshit Chanpura <[email protected]>
1 parent 6906ff0 commit 2d3b021

File tree

4 files changed

+27
-19
lines changed

4 files changed

+27
-19
lines changed

plugin/src/main/java/org/opensearch/ml/action/model_group/SearchModelGroupTransportAction.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
package org.opensearch.ml.action.model_group;
77

88
import static org.opensearch.ml.action.handler.MLSearchHandler.wrapRestActionListener;
9+
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MODEL_ACCESS_CONTROL_ENABLED;
910
import static org.opensearch.ml.utils.RestActionUtils.wrapListenerToHandleSearchIndexNotFound;
1011
import static org.opensearch.security.spi.resources.FeatureConfigConstants.OPENSEARCH_RESOURCE_SHARING_ENABLED;
1112
import static org.opensearch.security.spi.resources.FeatureConfigConstants.OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT;
@@ -82,8 +83,8 @@ private void preProcessRoleAndPerformSearch(
8283
User user,
8384
ActionListener<SearchResponse> listener
8485
) {
85-
boolean isResourceSharingFeatureEnabled = this.settings
86-
.getAsBoolean(OPENSEARCH_RESOURCE_SHARING_ENABLED, OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT);
86+
boolean isResourceSharingFeatureEnabled = ML_COMMONS_MODEL_ACCESS_CONTROL_ENABLED.get(settings)
87+
&& this.settings.getAsBoolean(OPENSEARCH_RESOURCE_SHARING_ENABLED, OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT);
8788
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
8889
ActionListener<SearchResponse> wrappedListener = ActionListener.runBefore(listener, context::restore);
8990

plugin/src/main/java/org/opensearch/ml/action/model_group/TransportUpdateModelGroupAction.java

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static org.opensearch.common.xcontent.json.JsonXContent.jsonXContent;
99
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
1010
import static org.opensearch.ml.common.CommonValue.ML_MODEL_GROUP_INDEX;
11+
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MODEL_ACCESS_CONTROL_ENABLED;
1112
import static org.opensearch.ml.utils.MLExceptionUtils.logException;
1213
import static org.opensearch.security.spi.resources.FeatureConfigConstants.OPENSEARCH_RESOURCE_SHARING_ENABLED;
1314
import static org.opensearch.security.spi.resources.FeatureConfigConstants.OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT;
@@ -115,8 +116,8 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLUpda
115116
return;
116117
}
117118
User user = RestActionUtils.getUserContext(client);
118-
boolean isResourceSharingFeatureEnabled = this.settings
119-
.getAsBoolean(OPENSEARCH_RESOURCE_SHARING_ENABLED, OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT);
119+
boolean isResourceSharingFeatureEnabled = ML_COMMONS_MODEL_ACCESS_CONTROL_ENABLED.get(settings)
120+
&& this.settings.getAsBoolean(OPENSEARCH_RESOURCE_SHARING_ENABLED, OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT);
120121
FetchSourceContext fetchSourceContext = new FetchSourceContext(true, Strings.EMPTY_ARRAY, Strings.EMPTY_ARRAY);
121122
GetDataObjectRequest getDataObjectRequest = GetDataObjectRequest
122123
.builder()
@@ -162,7 +163,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLUpda
162163
.getInstance()
163164
.getResourceSharingClient();
164165
resourceSharingClient
165-
.verifyResourceAccess(modelGroupId, ML_MODEL_GROUP_INDEX, ActionListener.wrap(isAuthorized -> {
166+
.verifyAccess(modelGroupId, ML_MODEL_GROUP_INDEX, ActionListener.wrap(isAuthorized -> {
166167
if (!isAuthorized) {
167168
listener
168169
.onFailure(

plugin/src/main/java/org/opensearch/ml/helper/ModelAccessControlHelper.java

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,11 @@ public ModelAccessControlHelper(ClusterService clusterService, Settings settings
9090
RangeQueryBuilder.class
9191
);
9292

93+
private boolean isResourceSharingFeatureEnabled(Settings settings) {
94+
return isModelAccessControlEnabled()
95+
&& settings.getAsBoolean(OPENSEARCH_RESOURCE_SHARING_ENABLED, OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT);
96+
}
97+
9398
// TODO Eventually remove this when all usages of it have been migrated to the SdkClient version
9499
public void validateModelGroupAccess(
95100
User user,
@@ -102,11 +107,10 @@ public void validateModelGroupAccess(
102107
listener.onResponse(true);
103108
return;
104109
}
105-
boolean isResourceSharingFeatureEnabled = settings
106-
.getAsBoolean(OPENSEARCH_RESOURCE_SHARING_ENABLED, OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT);
110+
boolean isResourceSharingFeatureEnabled = isResourceSharingFeatureEnabled(settings);
107111
if (isResourceSharingFeatureEnabled) {
108112
ResourceSharingClient resourceSharingClient = ResourceSharingClientAccessor.getInstance().getResourceSharingClient();
109-
resourceSharingClient.verifyResourceAccess(modelGroupId, ML_MODEL_GROUP_INDEX, ActionListener.wrap(isAuthorized -> {
113+
resourceSharingClient.verifyAccess(modelGroupId, ML_MODEL_GROUP_INDEX, ActionListener.wrap(isAuthorized -> {
110114
if (!isAuthorized) {
111115
listener
112116
.onFailure(
@@ -174,11 +178,10 @@ public void validateModelGroupAccess(
174178
listener.onResponse(true);
175179
return;
176180
}
177-
boolean isResourceSharingFeatureEnabled = settings
178-
.getAsBoolean(OPENSEARCH_RESOURCE_SHARING_ENABLED, OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT);
181+
boolean isResourceSharingFeatureEnabled = isResourceSharingFeatureEnabled(settings);
179182
if (isResourceSharingFeatureEnabled) {
180183
ResourceSharingClient resourceSharingClient = ResourceSharingClientAccessor.getInstance().getResourceSharingClient();
181-
resourceSharingClient.verifyResourceAccess(modelGroupId, ML_MODEL_GROUP_INDEX, ActionListener.wrap(isAuthorized -> {
184+
resourceSharingClient.verifyAccess(modelGroupId, ML_MODEL_GROUP_INDEX, ActionListener.wrap(isAuthorized -> {
182185
if (!isAuthorized) {
183186
listener
184187
.onFailure(
@@ -371,8 +374,7 @@ public SearchSourceBuilder addUserBackendRolesFilter(User user, SearchSourceBuil
371374
}
372375

373376
public SearchSourceBuilder createSearchSourceBuilder(User user, Settings settings) {
374-
boolean isResourceSharingFeatureEnabled = settings
375-
.getAsBoolean(OPENSEARCH_RESOURCE_SHARING_ENABLED, OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT);
377+
boolean isResourceSharingFeatureEnabled = isResourceSharingFeatureEnabled(settings);
376378
// TODO: Remove this feature flag check once feature is GA, as it will be enabled by default
377379
if (isResourceSharingFeatureEnabled) {
378380
return addAccessibleModelGroupsFilter(new SearchSourceBuilder());

plugin/src/main/java/org/opensearch/ml/model/MLModelGroupManager.java

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@
77

88
import static org.opensearch.common.xcontent.json.JsonXContent.jsonXContent;
99
import static org.opensearch.ml.common.CommonValue.ML_MODEL_GROUP_INDEX;
10+
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MODEL_ACCESS_CONTROL_ENABLED;
1011
import static org.opensearch.security.spi.resources.FeatureConfigConstants.OPENSEARCH_RESOURCE_SHARING_ENABLED;
1112
import static org.opensearch.security.spi.resources.FeatureConfigConstants.OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT;
13+
import static org.opensearch.security.spi.resources.ResourceAccessLevels.PLACE_HOLDER;
1214

1315
import java.time.Instant;
1416
import java.util.HashSet;
@@ -57,7 +59,8 @@
5759
import org.opensearch.search.builder.SearchSourceBuilder;
5860
import org.opensearch.security.spi.resources.client.ResourceSharingClient;
5961
import org.opensearch.security.spi.resources.sharing.Recipient;
60-
import org.opensearch.security.spi.resources.sharing.SharedWithActionGroup;
62+
import org.opensearch.security.spi.resources.sharing.Recipients;
63+
import org.opensearch.security.spi.resources.sharing.ShareWith;
6164
import org.opensearch.transport.client.Client;
6265

6366
import lombok.extern.log4j.Log4j2;
@@ -98,8 +101,8 @@ public void createModelGroup(MLRegisterModelGroupInput input, ActionListener<Str
98101
User user = RestActionUtils.getUserContext(client);
99102
// Create a recipient sharing list
100103
AtomicReference<Map<Recipient, Set<String>>> recipientMap = new AtomicReference<>();
101-
boolean isResourceSharingFeatureEnabled = this.settings
102-
.getAsBoolean(OPENSEARCH_RESOURCE_SHARING_ENABLED, OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT);
104+
boolean isResourceSharingFeatureEnabled = ML_COMMONS_MODEL_ACCESS_CONTROL_ENABLED.get(settings)
105+
&& this.settings.getAsBoolean(OPENSEARCH_RESOURCE_SHARING_ENABLED, OPENSEARCH_RESOURCE_SHARING_ENABLED_DEFAULT);
103106

104107
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
105108
ActionListener<String> wrappedListener = ActionListener.runBefore(listener, context::restore);
@@ -191,8 +194,9 @@ public void createModelGroup(MLRegisterModelGroupInput input, ActionListener<Str
191194
// Create an entry in resource-sharing index
192195
String modelGroupId = indexResponse.getId();
193196
String modelGroupIndex = indexResponse.getIndex();
194-
SharedWithActionGroup.ActionGroupRecipients recipients =
195-
new SharedWithActionGroup.ActionGroupRecipients(recipientMap.get());
197+
ShareWith shareWith = new ShareWith(
198+
Map.of(PLACE_HOLDER, new Recipients(recipientMap.get()))
199+
);
196200

197201
ResourceSharingClient resourceSharingClient = ResourceSharingClientAccessor
198202
.getInstance()
@@ -202,7 +206,7 @@ public void createModelGroup(MLRegisterModelGroupInput input, ActionListener<Str
202206
.share(
203207
modelGroupId,
204208
modelGroupIndex,
205-
recipients,
209+
shareWith,
206210
ActionListener.wrap(resourceSharing -> {
207211
log
208212
.debug(

0 commit comments

Comments
 (0)