From ecdfa33dbbab614e18f8e8694856482258143469 Mon Sep 17 00:00:00 2001
From: Louis Chu <clingzhi@amazon.com>
Date: Thu, 28 Mar 2024 15:12:45 -0700
Subject: [PATCH] Update superAdminAWSCredentialsProvider for bwc

Signed-off-by: Louis Chu <clingzhi@amazon.com>
---
 .../core/storage/FlintOpenSearchClient.java   | 43 +++++++++++--------
 1 file changed, 25 insertions(+), 18 deletions(-)

diff --git a/flint-core/src/main/scala/org/opensearch/flint/core/storage/FlintOpenSearchClient.java b/flint-core/src/main/scala/org/opensearch/flint/core/storage/FlintOpenSearchClient.java
index e5f32d2fa..8e16e9ec7 100644
--- a/flint-core/src/main/scala/org/opensearch/flint/core/storage/FlintOpenSearchClient.java
+++ b/flint-core/src/main/scala/org/opensearch/flint/core/storage/FlintOpenSearchClient.java
@@ -261,11 +261,24 @@ public IRestHighLevelClient createClient() {
       signer.setServiceName("es");
       signer.setRegionName(options.getRegion());
 
-      // Initialize and attempt to instantiate custom AWSCredentialsProviders.
+      // Use DefaultAWSCredentialsProviderChain by default.
       final AtomicReference<AWSCredentialsProvider> customAWSCredentialsProvider =
-              initializeAndInstantiateProvider(options.getCustomAwsCredentialsProvider());
+              new AtomicReference<>(new DefaultAWSCredentialsProviderChain());
+      String customProviderClass = options.getCustomAwsCredentialsProvider();
+      if (!Strings.isNullOrEmpty(customProviderClass)) {
+        instantiateProvider(customProviderClass, customAWSCredentialsProvider);
+      }
+
+      // Use the customAWSCredentialsProvider for superAdminAWSCredentialsProvider by default,
+      // unless a separate superAdmin provider class name is specified
+      String superAdminProviderClass = options.getSuperAdminAwsCredentialsProvider();
       final AtomicReference<AWSCredentialsProvider> superAdminAWSCredentialsProvider =
-              initializeAndInstantiateProvider(options.getSuperAdminAwsCredentialsProvider());
+              new AtomicReference<>(new DefaultAWSCredentialsProviderChain());
+      if (Strings.isNullOrEmpty(superAdminProviderClass)) {
+        superAdminAWSCredentialsProvider.set(customAWSCredentialsProvider.get());
+      } else {
+        instantiateProvider(superAdminProviderClass, superAdminAWSCredentialsProvider);
+      }
 
       restClientBuilder.setHttpClientConfigCallback(builder -> {
                 HttpAsyncClientBuilder delegate = builder.addInterceptorLast(
@@ -295,23 +308,17 @@ public IRestHighLevelClient createClient() {
   }
 
   /**
-   * Initializes and possibly instantiates an AWS credentials provider. If a provider class name is provided,
-   * this method attempts to instantiate the provider using reflection. Otherwise, it defaults to using the
-   * {@link DefaultAWSCredentialsProviderChain}.
+   * Attempts to instantiate the AWS credential provider using reflection.
    */
-  private AtomicReference<AWSCredentialsProvider> initializeAndInstantiateProvider(String providerClass) {
-    AWSCredentialsProvider provider = new DefaultAWSCredentialsProviderChain();
-    if (!Strings.isNullOrEmpty(providerClass)) {
-      try {
-        Class<?> clazz = Class.forName(providerClass);
-        Constructor<?> constructor = clazz.getDeclaredConstructor();
-        constructor.setAccessible(true);
-        provider = (AWSCredentialsProvider) constructor.newInstance();
-      } catch (Exception e) {
-        throw new RuntimeException("Failed to instantiate AWSCredentialsProvider: " + providerClass, e);
-      }
+  private void instantiateProvider(String providerClass, AtomicReference<AWSCredentialsProvider> provider) {
+    try {
+      Class<?> awsCredentialsProviderClass = Class.forName(providerClass);
+      Constructor<?> ctor = awsCredentialsProviderClass.getDeclaredConstructor();
+      ctor.setAccessible(true);
+      provider.set((AWSCredentialsProvider) ctor.newInstance());
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to instantiate AWSCredentialsProvider: " + providerClass, e);
     }
-    return new AtomicReference<>(provider);
   }
 
   /*