Adapts changes to make ResourceExtension consumed via @Inject

DarshitChanpura · DarshitChanpura · commit 8611f2549147 · 2025-08-15T14:15:29.000-04:00
Signed-off-by: Darshit Chanpura &lt;dchanp@amazon.com&gt;
diff --git a/build.gradle b/build.gradle
@@ -11,7 +11,7 @@ buildscript {
     ext {
         opensearch_group = "org.opensearch"
         isSnapshot = "true" == System.getProperty("build.snapshot", "true")
-        opensearch_version = System.getProperty("opensearch.version", "3.2.0-SNAPSHOT")
+        opensearch_version = System.getProperty("opensearch.version", "3.3.0-SNAPSHOT")
         buildVersionQualifier = System.getProperty("build.version_qualifier", "")
         asm_version = "9.7"
 
diff --git a/common/src/main/java/org/opensearch/ml/common/ResourceSharingClientAccessor.java b/common/src/main/java/org/opensearch/ml/common/ResourceSharingClientAccessor.java
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
@@ -22,6 +22,7 @@
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.Nullable;
+import org.opensearch.common.inject.Inject;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.commons.authuser.User;
 import org.opensearch.core.action.ActionListener;
@@ -38,13 +39,13 @@
 import org.opensearch.ml.common.CommonValue;
 import org.opensearch.ml.common.MLModel;
 import org.opensearch.ml.common.MLModelGroup;
-import org.opensearch.ml.common.ResourceSharingClientAccessor;
 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.resources.MLResourceSharingExtension;
 import org.opensearch.ml.utils.RestActionUtils;
 import org.opensearch.remote.metadata.client.SdkClient;
 import org.opensearch.remote.metadata.client.SearchDataObjectRequest;
@@ -71,6 +72,9 @@ public class MLSearchHandler {
     private ClusterService clusterService;
     private MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
+    @Inject(optional = true)
+    public MLResourceSharingExtension mlResourceSharingExtension;
+
     public MLSearchHandler(
         Client client,
         NamedXContentRegistry xContentRegistry,
@@ -87,7 +91,8 @@ public MLSearchHandler(
 
     /**
      * Fetch all the models from the model group index, and then create a combined query to model version index.
-     * @param sdkClient sdkclient a wrapper of the client
+     *
+     * @param sdkClient                  sdkclient a wrapper of the client
      * @param request
      * @param actionListener
      */
@@ -145,7 +150,7 @@ public void search(SdkClient sdkClient, SearchRequest request, String tenantId,
                     mlFeatureEnabledSetting.isMultiTenancyEnabled(),
                     CommonValue.ML_MODEL_GROUP_INDEX
                 );
-            boolean rsClientPresent = ResourceSharingClientAccessor.getInstance().getResourceSharingClient() != null;
+            boolean rsClientPresent = mlResourceSharingExtension != null && mlResourceSharingExtension.getResourceSharingClient() != null;
 
             if (skip || !hasIndex) {
                 // No gating at all
diff --git a/plugin/src/main/java/org/opensearch/ml/action/model_group/DeleteModelGroupTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/model_group/DeleteModelGroupTransportAction.java
@@ -27,12 +27,12 @@
 import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.TermQueryBuilder;
-import org.opensearch.ml.common.ResourceSharingClientAccessor;
 import org.opensearch.ml.common.exception.MLValidationException;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.model_group.MLModelGroupDeleteAction;
 import org.opensearch.ml.common.transport.model_group.MLModelGroupDeleteRequest;
 import org.opensearch.ml.helper.ModelAccessControlHelper;
+import org.opensearch.ml.resources.MLResourceSharingExtension;
 import org.opensearch.ml.utils.RestActionUtils;
 import org.opensearch.ml.utils.TenantAwareHelper;
 import org.opensearch.remote.metadata.client.DeleteDataObjectRequest;
@@ -62,6 +62,9 @@ public class DeleteModelGroupTransportAction extends HandledTransportAction<Acti
     final ModelAccessControlHelper modelAccessControlHelper;
     private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
+    @Inject(optional = true)
+    public MLResourceSharingExtension mlResourceSharingExtension;
+
     @Inject
     public DeleteModelGroupTransportAction(
         TransportService transportService,
@@ -96,7 +99,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
             ActionListener<DeleteResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
 
             // if resource sharing feature is enabled, access will be automatically checked by security plugin, so no need to check again
-            if (ResourceSharingClientAccessor.getInstance().getResourceSharingClient() != null) {
+            if (mlResourceSharingExtension != null && mlResourceSharingExtension.getResourceSharingClient() != null) {
                 checkForAssociatedModels(modelGroupId, tenantId, wrappedListener);
             } else {
                 validateAndDeleteModelGroup(modelGroupId, tenantId, wrappedListener);
diff --git a/plugin/src/main/java/org/opensearch/ml/action/model_group/GetModelGroupTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/model_group/GetModelGroupTransportAction.java
@@ -27,12 +27,12 @@
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.ml.common.MLModelGroup;
-import org.opensearch.ml.common.ResourceSharingClientAccessor;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.model_group.MLModelGroupGetAction;
 import org.opensearch.ml.common.transport.model_group.MLModelGroupGetRequest;
 import org.opensearch.ml.common.transport.model_group.MLModelGroupGetResponse;
 import org.opensearch.ml.helper.ModelAccessControlHelper;
+import org.opensearch.ml.resources.MLResourceSharingExtension;
 import org.opensearch.ml.utils.RestActionUtils;
 import org.opensearch.ml.utils.TenantAwareHelper;
 import org.opensearch.remote.metadata.client.GetDataObjectRequest;
@@ -59,6 +59,9 @@ public class GetModelGroupTransportAction extends HandledTransportAction<ActionR
     final ModelAccessControlHelper modelAccessControlHelper;
     private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
+    @Inject(optional = true)
+    public MLResourceSharingExtension mlResourceSharingExtension;
+
     @Inject
     public GetModelGroupTransportAction(
         TransportService transportService,
@@ -186,7 +189,7 @@ private void validateModelGroupAccess(
     ) {
         // if resource sharing feature is enabled, security plugin will have automatically evaluated access to this model group, hence no
         // need to validate again
-        if (ResourceSharingClientAccessor.getInstance().getResourceSharingClient() != null) {
+        if (mlResourceSharingExtension != null && mlResourceSharingExtension.getResourceSharingClient() != null) {
             wrappedListener.onResponse(MLModelGroupGetResponse.builder().mlModelGroup(mlModelGroup).build());
             return;
         }
diff --git a/plugin/src/main/java/org/opensearch/ml/action/model_group/SearchModelGroupTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/model_group/SearchModelGroupTransportAction.java
@@ -20,11 +20,11 @@
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.commons.authuser.User;
 import org.opensearch.core.action.ActionListener;
-import org.opensearch.ml.common.ResourceSharingClientAccessor;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.model_group.MLModelGroupSearchAction;
 import org.opensearch.ml.common.transport.search.MLSearchActionRequest;
 import org.opensearch.ml.helper.ModelAccessControlHelper;
+import org.opensearch.ml.resources.MLResourceSharingExtension;
 import org.opensearch.ml.utils.RestActionUtils;
 import org.opensearch.ml.utils.TenantAwareHelper;
 import org.opensearch.remote.metadata.client.SdkClient;
@@ -47,6 +47,9 @@ public class SearchModelGroupTransportAction extends HandledTransportAction<MLSe
 
     ModelAccessControlHelper modelAccessControlHelper;
 
+    @Inject(optional = true)
+    public MLResourceSharingExtension mlResourceSharingExtension;
+
     @Inject
     public SearchModelGroupTransportAction(
         TransportService transportService,
@@ -89,7 +92,7 @@ private void preProcessRoleAndPerformSearch(
                 .wrap(wrappedListener::onResponse, e -> wrapListenerToHandleSearchIndexNotFound(e, wrappedListener));
 
             // If resource-sharing feature is enabled, we fetch accessible model-groups and restrict the search to those model-groups only.
-            if (ResourceSharingClientAccessor.getInstance().getResourceSharingClient() != null) {
+            if (mlResourceSharingExtension != null && mlResourceSharingExtension.getResourceSharingClient() != null) {
                 // If a model-group is shared, then it will have been shared at-least at read access, hence the final result is guaranteed
                 // to only contain model-groups that the user at-least has read access to.
                 addAccessibleModelGroupsFilterAndSearch(tenantId, request, doubleWrappedListener);
@@ -113,7 +116,7 @@ private void addAccessibleModelGroupsFilterAndSearch(
         ActionListener<SearchResponse> wrappedListener
     ) {
         SearchSourceBuilder sourceBuilder = request.source() != null ? request.source() : new SearchSourceBuilder();
-        ResourceSharingClient rsc = ResourceSharingClientAccessor.getInstance().getResourceSharingClient();
+        ResourceSharingClient rsc = mlResourceSharingExtension.getResourceSharingClient();
         // filter by accessible model-groups
         rsc.getAccessibleResourceIds(ML_MODEL_GROUP_INDEX, ActionListener.wrap(ids -> {
             sourceBuilder.query(modelAccessControlHelper.mergeWithAccessFilter(sourceBuilder.query(), ids));
diff --git a/plugin/src/main/java/org/opensearch/ml/action/model_group/TransportUpdateModelGroupAction.java b/plugin/src/main/java/org/opensearch/ml/action/model_group/TransportUpdateModelGroupAction.java
@@ -35,7 +35,6 @@
 import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.ml.common.AccessMode;
 import org.opensearch.ml.common.MLModelGroup;
-import org.opensearch.ml.common.ResourceSharingClientAccessor;
 import org.opensearch.ml.common.exception.MLValidationException;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.model_group.MLUpdateModelGroupAction;
@@ -44,6 +43,7 @@
 import org.opensearch.ml.common.transport.model_group.MLUpdateModelGroupResponse;
 import org.opensearch.ml.helper.ModelAccessControlHelper;
 import org.opensearch.ml.model.MLModelGroupManager;
+import org.opensearch.ml.resources.MLResourceSharingExtension;
 import org.opensearch.ml.utils.RestActionUtils;
 import org.opensearch.ml.utils.TenantAwareHelper;
 import org.opensearch.remote.metadata.client.GetDataObjectRequest;
@@ -74,6 +74,9 @@ public class TransportUpdateModelGroupAction extends HandledTransportAction<Acti
     MLModelGroupManager mlModelGroupManager;
     private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
+    @Inject(optional = true)
+    public MLResourceSharingExtension mlResourceSharingExtension;
+
     @Inject
     public TransportUpdateModelGroupAction(
         TransportService transportService,
@@ -149,7 +152,10 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLUpda
                                     )) {
                                     // NOTE all sharing and revoking must happen through share API exposed by security plugin
                                     // client == null -> feature is disabled, follow old route
-                                    if (ResourceSharingClientAccessor.getInstance().getResourceSharingClient() == null) {
+                                    // client != null -> feature is enabled, evaluation already happened in security plugin, move to update
+                                    // method
+                                    if (mlResourceSharingExtension == null
+                                        || mlResourceSharingExtension.getResourceSharingClient() == null) {
                                         // TODO: At some point, this call must be replaced by the one above, (i.e. no user info to
                                         // be stored in model-group index)
                                         if (modelAccessControlHelper.isSecurityEnabledAndModelAccessControlEnabled(user)) {
@@ -160,8 +166,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLUpda
 
                                     }
                                     // For backwards compatibility we still allow storing backend_roles
-                                    // data in ml_model_group
-                                    // index
+                                    // data in ml_model_group index
                                     updateModelGroup(modelGroupId, r.source(), updateModelGroupInput, wrappedListener, user);
 
                                 }
diff --git a/plugin/src/main/java/org/opensearch/ml/helper/ModelAccessControlHelper.java b/plugin/src/main/java/org/opensearch/ml/helper/ModelAccessControlHelper.java
@@ -25,6 +25,7 @@
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.inject.Inject;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.common.xcontent.LoggingDeprecationHandler;
@@ -49,10 +50,10 @@
 import org.opensearch.index.query.TermsQueryBuilder;
 import org.opensearch.ml.common.AccessMode;
 import org.opensearch.ml.common.MLModelGroup;
-import org.opensearch.ml.common.ResourceSharingClientAccessor;
 import org.opensearch.ml.common.exception.MLResourceNotFoundException;
 import org.opensearch.ml.common.exception.MLValidationException;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
+import org.opensearch.ml.resources.MLResourceSharingExtension;
 import org.opensearch.ml.utils.MLNodeUtils;
 import org.opensearch.ml.utils.TenantAwareHelper;
 import org.opensearch.remote.metadata.client.GetDataObjectRequest;
@@ -72,6 +73,9 @@ public class ModelAccessControlHelper {
 
     private volatile Boolean modelAccessControlEnabled;
 
+    @Inject
+    public MLResourceSharingExtension mlResourceSharingExtension;
+
     public ModelAccessControlHelper(ClusterService clusterService, Settings settings) {
         modelAccessControlEnabled = ML_COMMONS_MODEL_ACCESS_CONTROL_ENABLED.get(settings);
         clusterService
@@ -97,8 +101,9 @@ public void validateModelGroupAccess(User user, String modelGroupId, String acti
             listener.onResponse(true);
             return;
         }
-        if (ResourceSharingClientAccessor.getInstance().getResourceSharingClient() != null) {
-            ResourceSharingClient resourceSharingClient = ResourceSharingClientAccessor.getInstance().getResourceSharingClient();
+
+        if (mlResourceSharingExtension != null && mlResourceSharingExtension.getResourceSharingClient() != null) {
+            ResourceSharingClient resourceSharingClient = mlResourceSharingExtension.getResourceSharingClient();
             resourceSharingClient.verifyAccess(modelGroupId, ML_MODEL_GROUP_INDEX, action, ActionListener.wrap(isAuthorized -> {
                 if (!isAuthorized) {
                     listener
@@ -167,8 +172,8 @@ public void validateModelGroupAccess(
             listener.onResponse(true);
             return;
         }
-        if (ResourceSharingClientAccessor.getInstance().getResourceSharingClient() != null) {
-            ResourceSharingClient resourceSharingClient = ResourceSharingClientAccessor.getInstance().getResourceSharingClient();
+        if (mlResourceSharingExtension != null && mlResourceSharingExtension.getResourceSharingClient() != null) {
+            ResourceSharingClient resourceSharingClient = mlResourceSharingExtension.getResourceSharingClient();
             resourceSharingClient.verifyAccess(modelGroupId, ML_MODEL_GROUP_INDEX, action, ActionListener.wrap(isAuthorized -> {
                 if (!isAuthorized) {
                     listener
@@ -382,9 +387,9 @@ public void addAccessibleModelGroupsFilterAndSearch(
         Consumer<Set<String>> onSuccess,
         ActionListener<SearchResponse> wrappedListener
     ) {
-        ResourceSharingClient rsc = ResourceSharingClientAccessor.getInstance().getResourceSharingClient();
+        ResourceSharingClient resourceSharingClient = mlResourceSharingExtension.getResourceSharingClient();
         // filter by accessible model-groups
-        rsc.getAccessibleResourceIds(ML_MODEL_GROUP_INDEX, ActionListener.wrap(onSuccess::accept, e -> {
+        resourceSharingClient.getAccessibleResourceIds(ML_MODEL_GROUP_INDEX, ActionListener.wrap(onSuccess::accept, e -> {
             // Fail-safe: deny-all and still return a response
             SearchSourceBuilder reqSrc = request.source() != null ? request.source() : new SearchSourceBuilder();
             reqSrc.query(mergeWithAccessFilter(reqSrc.query(), Collections.emptySet()));
diff --git a/plugin/src/main/java/org/opensearch/ml/resources/MLResourceSharingExtension.java b/plugin/src/main/java/org/opensearch/ml/resources/MLResourceSharingExtension.java
@@ -10,20 +10,26 @@
 import java.util.Set;
 
 import org.opensearch.ml.common.MLModelGroup;
-import org.opensearch.ml.common.ResourceSharingClientAccessor;
 import org.opensearch.security.spi.resources.ResourceProvider;
 import org.opensearch.security.spi.resources.ResourceSharingExtension;
 import org.opensearch.security.spi.resources.client.ResourceSharingClient;
 
 public class MLResourceSharingExtension implements ResourceSharingExtension {
 
+    private ResourceSharingClient resourceSharingClient;
+
     @Override
     public Set<ResourceProvider> getResourceProviders() {
         return Set.of(new ResourceProvider(MLModelGroup.class.getCanonicalName(), ML_MODEL_GROUP_INDEX));
     }
 
     @Override
     public void assignResourceSharingClient(ResourceSharingClient resourceSharingClient) {
-        ResourceSharingClientAccessor.getInstance().setResourceSharingClient(resourceSharingClient);
+        this.resourceSharingClient = resourceSharingClient;
+    }
+
+    @Override
+    public ResourceSharingClient getResourceSharingClient() {
+        return resourceSharingClient;
     }
 }
diff --git a/plugin/src/test/java/org/opensearch/ml/action/model_group/SearchModelGroupTransportActionTests.java b/plugin/src/test/java/org/opensearch/ml/action/model_group/SearchModelGroupTransportActionTests.java
diff --git a/plugin/src/test/java/org/opensearch/ml/resources/MLResourceSharingExtensionTests.java b/plugin/src/test/java/org/opensearch/ml/resources/MLResourceSharingExtensionTests.java
