Make AWS client choice configurable for cluster or serverless

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

Hide file tree

Showing 6 changed files with 48 additions and 20 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
@@ -179,7 +179,7 @@ integTest {
     dependsOn "bundlePlugin"
     systemProperty 'tests.security.manager', 'false'
     systemProperty 'java.io.tmpdir', opensearch_tmp_dir.absolutePath
-    
+
     systemProperty "https", System.getProperty("https")
     systemProperty "user", System.getProperty("user")
     systemProperty "password", System.getProperty("password")
@@ -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";
 };
diff --git a/plugin/src/test/java/org/opensearch/ml/sdkclient/SdkClientFactoryTests.java b/plugin/src/test/java/org/opensearch/ml/sdkclient/SdkClientFactoryTests.java
@@ -47,6 +47,7 @@ public void testAwsOpenSearchServiceBinding() {
             .put(SdkClientSettings.REMOTE_METADATA_TYPE_KEY, SdkClientSettings.AWS_OPENSEARCH_SERVICE)
             .put(SdkClientSettings.REMOTE_METADATA_ENDPOINT_KEY, "example.org")
             .put(SdkClientSettings.REMOTE_METADATA_REGION_KEY, "eu-west-3")
+            .put(SdkClientSettings.REMOTE_METADATA_SERVICE_NAME_KEY, "es")
             .build();
         SdkClient sdkClient = SdkClientFactory.createSdkClient(mock(Client.class), NamedXContentRegistry.EMPTY, settings);
         assertTrue(sdkClient.getDelegate() instanceof RemoteClusterIndicesClient);
@@ -56,8 +57,9 @@ public void testDDBBinding() {
         Settings settings = Settings
             .builder()
             .put(SdkClientSettings.REMOTE_METADATA_TYPE_KEY, SdkClientSettings.AWS_DYNAMO_DB)
-            .put(SdkClientSettings.REMOTE_METADATA_ENDPOINT_KEY, "http://example.org")
+            .put(SdkClientSettings.REMOTE_METADATA_ENDPOINT_KEY, "example.org")
             .put(SdkClientSettings.REMOTE_METADATA_REGION_KEY, "eu-west-3")
+            .put(SdkClientSettings.REMOTE_METADATA_SERVICE_NAME_KEY, "aoss")
             .build();
         SdkClient sdkClient = SdkClientFactory.createSdkClient(mock(Client.class), NamedXContentRegistry.EMPTY, settings);
         assertTrue(sdkClient.getDelegate() instanceof DDBOpenSearchClient);
@@ -80,6 +82,17 @@ public void testAwsOpenSearchServiceBindingException() {
             OpenSearchException.class,
             () -> SdkClientFactory.createSdkClient(mock(Client.class), NamedXContentRegistry.EMPTY, settings)
         );
+        Settings settings2 = Settings
+            .builder()
+            .put(SdkClientSettings.REMOTE_METADATA_TYPE_KEY, SdkClientSettings.AWS_OPENSEARCH_SERVICE)
+            .put(SdkClientSettings.REMOTE_METADATA_ENDPOINT_KEY, "example.org")
+            .put(SdkClientSettings.REMOTE_METADATA_REGION_KEY, "eu-west-3")
+            .put(SdkClientSettings.REMOTE_METADATA_SERVICE_NAME_KEY, "invalid")
+            .build();
+        assertThrows(
+            OpenSearchException.class,
+            () -> SdkClientFactory.createSdkClient(mock(Client.class), NamedXContentRegistry.EMPTY, settings2)
+        );
     }
 
     public void testDDBBindingException() {
@@ -88,5 +101,16 @@ public void testDDBBindingException() {
             OpenSearchException.class,
             () -> SdkClientFactory.createSdkClient(mock(Client.class), NamedXContentRegistry.EMPTY, settings)
         );
+        Settings settingss = Settings
+            .builder()
+            .put(SdkClientSettings.REMOTE_METADATA_TYPE_KEY, SdkClientSettings.AWS_DYNAMO_DB)
+            .put(SdkClientSettings.REMOTE_METADATA_ENDPOINT_KEY, "example.org")
+            .put(SdkClientSettings.REMOTE_METADATA_REGION_KEY, "eu-west-3")
+            .put(SdkClientSettings.REMOTE_METADATA_SERVICE_NAME_KEY, "invalid")
+            .build();
+        assertThrows(
+            OpenSearchException.class,
+            () -> SdkClientFactory.createSdkClient(mock(Client.class), NamedXContentRegistry.EMPTY, settingss)
+        );
     }
 }