Make AWS client choice configurable for cluster or serverless

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Sep 13, 2024 · ae8a4fa · ae8a4fa
1 parent 697f126
commit ae8a4fa
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 18 deletions.
diff --git a/common/src/main/java/org/opensearch/sdk/SdkClientSettings.java b/common/src/main/java/org/opensearch/sdk/SdkClientSettings.java
@@ -28,8 +28,11 @@ public class SdkClientSettings {
     public static final String REMOTE_METADATA_ENDPOINT_KEY = "plugins.ml_commons.remote_metadata_endpoint";
     /** The key for remote metadata region, applicable to AWS remote clusters or DynamoDB. */
     public static final String REMOTE_METADATA_REGION_KEY = "plugins.ml_commons.remote_metadata_region";
+    /** The key for remote metadata AWS service, either "es" for cluster or "aoss" for OpenSearch Serverless */
+    public static final String REMOTE_METADATA_SERVICE_NAME_KEY = "plugins.ml_commons.remote_metadata_service_name";
 
     public static final Setting<String> REMOTE_METADATA_TYPE = Setting.simpleString(REMOTE_METADATA_TYPE_KEY, Property.NodeScope, Property.Final);
     public static final Setting<String> REMOTE_METADATA_ENDPOINT = Setting.simpleString(REMOTE_METADATA_ENDPOINT_KEY, Property.NodeScope, Property.Final);
     public static final Setting<String> REMOTE_METADATA_REGION = Setting.simpleString(REMOTE_METADATA_REGION_KEY, Property.NodeScope, Property.Final);
+    public static final Setting<String> REMOTE_METADATA_SERVICE_NAME = Setting.simpleString(REMOTE_METADATA_SERVICE_NAME_KEY, Property.NodeScope, Property.Final);
 }
diff --git a/plugin/build.gradle b/plugin/build.gradle
@@ -357,9 +357,7 @@ List<String> jacocoExclusions = [
         'org.opensearch.ml.helper.ModelAccessControlHelper',
         'org.opensearch.ml.action.models.DeleteModelTransportAction.2',
         'org.opensearch.ml.model.MLModelCacheHelper',
-        'org.opensearch.ml.model.MLModelCacheHelper.1',
-        'org.opensearch.ml.sdkclient.SdkClientFactory',
-        'org.opensearch.ml.sdkclient.SdkClientFactory.SocketAccess'
+        'org.opensearch.ml.model.MLModelCacheHelper.1'
 ]
 
 jacocoTestCoverageVerification {

diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
@@ -949,7 +949,8 @@ public List<Setting<?>> getSettings() {
                 // Settings for SdkClient
                 SdkClientSettings.REMOTE_METADATA_TYPE,
                 SdkClientSettings.REMOTE_METADATA_ENDPOINT,
-                SdkClientSettings.REMOTE_METADATA_REGION
+                SdkClientSettings.REMOTE_METADATA_REGION,
+                SdkClientSettings.REMOTE_METADATA_SERVICE_NAME
             );
     }
 

diff --git a/plugin/src/main/java/org/opensearch/ml/sdkclient/SdkClientFactory.java b/plugin/src/main/java/org/opensearch/ml/sdkclient/SdkClientFactory.java
@@ -75,6 +75,7 @@ public static SdkClient createSdkClient(Client client, NamedXContentRegistry xCo
         String remoteMetadataType = SdkClientSettings.REMOTE_METADATA_TYPE.get(settings);
         String remoteMetadataEndpoint = SdkClientSettings.REMOTE_METADATA_ENDPOINT.get(settings);
         String region = SdkClientSettings.REMOTE_METADATA_REGION.get(settings);
+        String serviceName = SdkClientSettings.REMOTE_METADATA_SERVICE_NAME.get(settings);
 
         switch (remoteMetadataType) {
             case SdkClientSettings.REMOTE_OPENSEARCH:
@@ -84,24 +85,18 @@ public static SdkClient createSdkClient(Client client, NamedXContentRegistry xCo
                 log.info("Using remote opensearch cluster as metadata store");
                 return new SdkClient(new RemoteClusterIndicesClient(createOpenSearchClient(remoteMetadataEndpoint)));
             case SdkClientSettings.AWS_OPENSEARCH_SERVICE:
-                if (Strings.isBlank(remoteMetadataEndpoint) || Strings.isBlank(region)) {
-                    throw new OpenSearchException("AWS Opensearch Service client requires a metadata endpoint and region.");
-                }
+                validateAwsParams(remoteMetadataType, remoteMetadataEndpoint, region, serviceName);
                 log.info("Using remote AWS Opensearch Service cluster as metadata store");
                 return new SdkClient(
-                    new RemoteClusterIndicesClient(createAwsOpenSearchServiceClient(remoteMetadataEndpoint, region, "es"))
+                    new RemoteClusterIndicesClient(createAwsOpenSearchServiceClient(remoteMetadataEndpoint, region, serviceName))
                 );
             case SdkClientSettings.AWS_DYNAMO_DB:
-                if (Strings.isBlank(remoteMetadataEndpoint) || Strings.isBlank(region)) {
-                    throw new OpenSearchException(
-                        "AWS DynamoDB and Opensearch Serverless Service client requires a metadata endpoint and region."
-                    );
-                }
+                validateAwsParams(remoteMetadataType, remoteMetadataEndpoint, region, serviceName);
                 log.info("Using dynamo DB as metadata store");
                 return new SdkClient(
                     new DDBOpenSearchClient(
                         createDynamoDbClient(region),
-                        new RemoteClusterIndicesClient(createAwsOpenSearchServiceClient(remoteMetadataEndpoint, region, "aoss"))
+                        new RemoteClusterIndicesClient(createAwsOpenSearchServiceClient(remoteMetadataEndpoint, region, serviceName))
                     )
                 );
             default:
@@ -110,6 +105,15 @@ public static SdkClient createSdkClient(Client client, NamedXContentRegistry xCo
         }
     }
 
+    private static void validateAwsParams(String clientType, String remoteMetadataEndpoint, String region, String serviceName) {
+        if (Strings.isBlank(remoteMetadataEndpoint) || Strings.isBlank(region)) {
+            throw new OpenSearchException(clientType + " client requires a metadata endpoint and region.");
+        }
+        if (!"es".equals(serviceName) && !"aoss".equals(serviceName)) {
+            throw new OpenSearchException(clientType + " client requires a signing service of 'es' or 'aoss'.");
+        }
+    }
+
     // Package private for testing
     static SdkClient wrapSdkClientDelegate(SdkClientDelegate delegate) {
         return new SdkClient(delegate);

diff --git a/plugin/src/main/plugin-metadata/plugin-security.policy b/plugin/src/main/plugin-metadata/plugin-security.policy
@@ -24,9 +24,7 @@ grant {
 
     // Circuit Breaker
     permission java.lang.RuntimePermission "getFileSystemAttributes";
-    
+
     // AWS credentials needed for clients
-    permission java.io.FilePermission "<<ALL FILES>>", "read";
-
-    permission java.security.AllPermission;
+    permission java.io.FilePermission "${user.home}/.aws/-", "read";
 };