Use interface-based approach instead of Reflection for Credentials Object Creation

Jithendar12 · Jithendar12 · commit f538ea667ac8 · 2025-09-04T17:54:36.000+05:30
diff --git a/athena-datalakegen2/src/main/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2MetadataHandler.java b/athena-datalakegen2/src/main/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2MetadataHandler.java
@@ -293,7 +293,7 @@ protected CredentialsProvider getCredentialProvider()
         return CredentialsProviderFactory.createCredentialProvider(
             getDatabaseConnectionConfig().getSecret(),
             getCachableSecretsManager(),
-            DataLakeGen2OAuthCredentialsProvider.class
+            new DataLakeGen2OAuthCredentialsProvider()
         );
     }
 }
diff --git a/athena-datalakegen2/src/main/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2OAuthCredentialsProvider.java b/athena-datalakegen2/src/main/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2OAuthCredentialsProvider.java
@@ -21,7 +21,6 @@
 
 import com.amazonaws.athena.connector.credentials.CredentialsConstants;
 import com.amazonaws.athena.connector.credentials.OAuthCredentialsProvider;
-import com.amazonaws.athena.connector.lambda.security.CachableSecretsManager;
 import com.google.common.annotations.VisibleForTesting;
 
 import java.net.URI;
@@ -38,15 +37,15 @@ public class DataLakeGen2OAuthCredentialsProvider extends OAuthCredentialsProvid
     private static final String SCOPE = "https://sql.azuresynapse.net/.default";
     private static final String TENANT_ID = "tenant_id";
 
-    public DataLakeGen2OAuthCredentialsProvider(String secretName, Map<String, String> secretMap, CachableSecretsManager secretsManager)
+    public DataLakeGen2OAuthCredentialsProvider()
     {
-        super(secretName, secretMap, secretsManager);
+        super();
     }
 
     @VisibleForTesting
-    public DataLakeGen2OAuthCredentialsProvider(String secretName, Map<String, String> secretMap, CachableSecretsManager secretsManager, HttpClient httpClient)
+    protected DataLakeGen2OAuthCredentialsProvider(HttpClient httpClient)
     {
-        super(secretName, secretMap, secretsManager, httpClient);
+        super(httpClient);
     }
 
     @Override
diff --git a/athena-datalakegen2/src/main/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2RecordHandler.java b/athena-datalakegen2/src/main/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2RecordHandler.java
@@ -85,7 +85,7 @@ protected CredentialsProvider getCredentialProvider()
         return CredentialsProviderFactory.createCredentialProvider(
             getDatabaseConnectionConfig().getSecret(),
             getCachableSecretsManager(),
-            DataLakeGen2OAuthCredentialsProvider.class
+            new DataLakeGen2OAuthCredentialsProvider()
         );
     }
 }
diff --git a/athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2OAuthCredentialsProviderTest.java b/athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2OAuthCredentialsProviderTest.java
@@ -70,7 +70,8 @@ public void setup()
     {
         MockitoAnnotations.openMocks(this);
         CachableSecretsManager cachableSecretsManager = new CachableSecretsManager(secretsManagerClient);
-        credentialsProvider = new DataLakeGen2OAuthCredentialsProvider(SECRET_NAME, new HashMap<>(), cachableSecretsManager, httpClient);
+        credentialsProvider = new DataLakeGen2OAuthCredentialsProvider(httpClient);
+        credentialsProvider.initialize(SECRET_NAME, new HashMap<>(), cachableSecretsManager);
         objectMapper = new ObjectMapper();
     }
 
diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/credentials/CredentialsProviderFactory.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/credentials/CredentialsProviderFactory.java
@@ -27,7 +27,6 @@
 import software.amazon.awssdk.services.glue.model.FederationSourceErrorCode;
 
 import java.io.IOException;
-import java.lang.reflect.Constructor;
 import java.util.Map;
 
 /**
@@ -50,40 +49,37 @@ private CredentialsProviderFactory()
      *
      * @param secretName The name of the secret in AWS Secrets Manager
      * @param secretsManager The secrets manager instance
-     * @param oAuthProviderClass The class of the OAuth provider to instantiate (must extend OAuthCredentialsProvider)
+     * @param provider The InitializableCredentialsProvider instance to check and initialize
      * @return A new CredentialsProvider instance based on the secret configuration
      * @throws AthenaConnectorException if there are errors deserializing the secret or creating the provider
      */
-    public static <T extends OAuthCredentialsProvider> CredentialsProvider createCredentialProvider(
+    public static CredentialsProvider createCredentialProvider(
             String secretName,
             CachableSecretsManager secretsManager,
-            Class<T> oAuthProviderClass)
+            InitializableCredentialsProvider provider)
     {
         if (StringUtils.isNotBlank(secretName)) {
             try {
                 String secretString = secretsManager.getSecret(secretName);
                 Map<String, String> secretMap = OBJECT_MAPPER.readValue(secretString, Map.class);
 
-                // Create an instance of the OAuth provider
-                T provider;
-                try {
-                    Constructor<T> constructor = oAuthProviderClass.getConstructor(
-                            String.class, Map.class, CachableSecretsManager.class);
-                    provider = constructor.newInstance(secretName, secretMap, secretsManager);
+                if (provider instanceof OAuthCredentialsProvider) {
+                    OAuthCredentialsProvider oauthProvider = (OAuthCredentialsProvider) provider;
+                    try {
+                        // Check if OAuth is configured
+                        if (oauthProvider.isOAuthConfigured(secretMap)) {
+                            oauthProvider.initialize(secretName, secretMap, secretsManager);
+                            return oauthProvider;
+                        }
+                    }
+                    catch (RuntimeException e) {
+                        throw new AthenaConnectorException("Failed to create OAuth provider: " + e.getMessage(),
+                                ErrorDetails.builder()
+                                        .errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
+                                        .errorMessage(e.getMessage())
+                                        .build());
+                    }
                 }
-                catch (ReflectiveOperationException e) {
-                    throw new AthenaConnectorException("Failed to create OAuth provider: " + e.getMessage(),
-                            ErrorDetails.builder()
-                                    .errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
-                                    .errorMessage(e.getMessage())
-                                    .build());
-                }
-
-                // Check if OAuth is configured
-                if (provider.isOAuthConfigured(secretMap)) {
-                    return provider;
-                }
-
                 // Fall back to default credentials if OAuth is not configured
                 return new DefaultCredentialsProvider(secretString);
             }
diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/credentials/InitializableCredentialsProvider.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/credentials/InitializableCredentialsProvider.java
@@ -0,0 +1,42 @@
+/*-
+ * #%L
+ * athena-federation-sdk
+ * %%
+ * Copyright (C) 2019 - 2025 Amazon Web Services
+ * %%
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * #L%
+ */
+package com.amazonaws.athena.connector.credentials;
+
+import com.amazonaws.athena.connector.lambda.security.CachableSecretsManager;
+
+import java.util.Map;
+
+/**
+ * Extended contract for credentials providers that require initialization with configuration from AWS Secrets Manager.
+ * Implementations must be initialized before use and may support various authentication methods including OAuth,
+ * username/password, or other configurable authentication schemes.
+ */
+public interface InitializableCredentialsProvider extends CredentialsProvider
+{
+    /**
+     * Initializes this credential provider with the given configuration.
+     * Must be called exactly once before any calls to getCredential().
+     *
+     * @param secretName The name of the secret in AWS Secrets Manager
+     * @param secretMap The secret configuration containing authentication parameters
+     * @param secretsManager The secrets manager instance for retrieving and updating secrets
+     */
+    void initialize(String secretName, Map<String, String> secretMap, CachableSecretsManager secretsManager);
+}
diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/credentials/OAuthCredentialsProvider.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/credentials/OAuthCredentialsProvider.java
@@ -24,6 +24,7 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.services.glue.model.ErrorDetails;
@@ -39,33 +40,42 @@
 
 import static com.amazonaws.athena.connector.credentials.CredentialsConstants.ACCESS_TOKEN;
 import static com.amazonaws.athena.connector.credentials.CredentialsConstants.EXPIRES_IN;
+import static com.amazonaws.athena.connector.credentials.CredentialsConstants.FETCHED_AT;
 
 /**
  * Base class for OAuth credential providers.
  * Handles OAuth token lifecycle management.
  */
-public abstract class OAuthCredentialsProvider implements CredentialsProvider
+public abstract class OAuthCredentialsProvider implements InitializableCredentialsProvider
 {
     private static final Logger LOGGER = LoggerFactory.getLogger(OAuthCredentialsProvider.class);
     private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
 
-    private final String secretName;
-    private final CachableSecretsManager secretsManager;
-    private final Map<String, String> secretMap;
+    private String secretName;
+    private CachableSecretsManager secretsManager;
+    private Map<String, String> secretMap;
+    private HttpClient httpClient;
 
-    private final HttpClient httpClient;
+    protected OAuthCredentialsProvider()
+    {
+        // Initialized later via initialize() method
+    }
 
-    protected OAuthCredentialsProvider(String secretName, Map<String, String> secretMap, CachableSecretsManager secretsManager)
+    @VisibleForTesting
+    protected OAuthCredentialsProvider(HttpClient httpClient)
     {
-        this(secretName, secretMap, secretsManager, HttpClient.newBuilder().build());
+        this.httpClient = httpClient;
     }
 
-    protected OAuthCredentialsProvider(String secretName, Map<String, String> secretMap, CachableSecretsManager secretsManager, HttpClient httpClient)
+    @Override
+    public void initialize(String secretName, Map<String, String> secretMap, CachableSecretsManager secretsManager)
     {
         this.secretName = secretName;
-        this.secretsManager = secretsManager;
         this.secretMap = secretMap;
-        this.httpClient = httpClient;
+        this.secretsManager = secretsManager;
+        this.httpClient = (this.httpClient != null)
+                ? this.httpClient
+                : HttpClient.newBuilder().build();
     }
 
     @Override
@@ -121,7 +131,7 @@ private String getOAuthAccessToken(Map<String, String> secretMap) throws IOExcep
         String cached = secretMap.get(ACCESS_TOKEN);
         if (cached != null) {
             long expiresIn = Long.parseLong(secretMap.getOrDefault(EXPIRES_IN, "0"));
-            long fetchedAt = Long.parseLong(secretMap.getOrDefault(CredentialsConstants.FETCHED_AT, "0"));
+            long fetchedAt = Long.parseLong(secretMap.getOrDefault(FETCHED_AT, "0"));
             long now = System.currentTimeMillis() / 1000;
 
             if ((now - fetchedAt) < expiresIn - 60) {
@@ -186,7 +196,7 @@ private String getOAuthAccessToken(Map<String, String> secretMap) throws IOExcep
 
         secretMap.put(ACCESS_TOKEN, accessToken);
         secretMap.put(EXPIRES_IN, String.valueOf(expiresIn));
-        secretMap.put(CredentialsConstants.FETCHED_AT, String.valueOf(fetchedAt));
+        secretMap.put(FETCHED_AT, String.valueOf(fetchedAt));
 
         try {
             String secretString = OBJECT_MAPPER.writeValueAsString(secretMap);
diff --git a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/credentials/CredentialsProviderFactoryTest.java b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/credentials/CredentialsProviderFactoryTest.java
@@ -70,11 +70,6 @@ public void setup()
     // Test OAuth provider class for testing
     public static class TestOAuthProvider extends OAuthCredentialsProvider
     {
-        public TestOAuthProvider(String secretName, Map<String, String> secretMap, CachableSecretsManager secretsManager)
-        {
-            super(secretName, secretMap, secretsManager);
-        }
-
         @Override
         protected boolean isOAuthConfigured(Map<String, String> secretMap)
         {
@@ -103,7 +98,7 @@ public void testCreateCredentialProvider_whenOAuthConfigured() throws JsonProces
         CredentialsProvider provider = CredentialsProviderFactory.createCredentialProvider(
                 SECRET_NAME, 
                 cachableSecretsManager,
-                TestOAuthProvider.class
+                new TestOAuthProvider()
         );
         assertTrue(provider instanceof TestOAuthProvider);
     }
@@ -123,7 +118,7 @@ public void testCreateCredentialProvider_whenUsernamePasswordConfigured() throws
         CredentialsProvider provider = CredentialsProviderFactory.createCredentialProvider(
                 SECRET_NAME,
                 cachableSecretsManager,
-                TestOAuthProvider.class
+                new TestOAuthProvider()
         );
         assertTrue(provider instanceof DefaultCredentialsProvider);
     }
@@ -134,7 +129,7 @@ public void testCreateCredentialProvider_whenNoSecret()
         CredentialsProvider provider = CredentialsProviderFactory.createCredentialProvider(
                 "",
                 cachableSecretsManager,
-                TestOAuthProvider.class
+                new TestOAuthProvider()
         );
         assertNull(provider);
     }
@@ -151,7 +146,7 @@ public void testCreateCredentialProvider_whenInvalidJson_throwsException()
             CredentialsProviderFactory.createCredentialProvider(
                     SECRET_NAME,
                     cachableSecretsManager,
-                    TestOAuthProvider.class
+                    new TestOAuthProvider()
             );
             fail("Expected AthenaConnectorException");
         }
@@ -168,11 +163,6 @@ public void testCreateCredentialProvider_whenInvalidOAuthClass_throwsException()
         // Create a class that doesn't properly implement the required methods
         class InvalidOAuthProvider extends OAuthCredentialsProvider
         {
-            public InvalidOAuthProvider(String secretName, Map<String, String> secretMap, CachableSecretsManager secretsManager)
-            {
-                super(secretName, secretMap, secretsManager);
-            }
-
             @Override
             protected boolean isOAuthConfigured(Map<String, String> secretMap)
             {
@@ -198,7 +188,7 @@ protected HttpRequest buildTokenRequest(Map<String, String> secretMap)
             CredentialsProviderFactory.createCredentialProvider(
                     SECRET_NAME,
                     cachableSecretsManager,
-                    InvalidOAuthProvider.class
+                    new InvalidOAuthProvider()
             );
             fail("Expected AthenaConnectorException");
         }
diff --git a/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/credentials/OAuthCredentialsProviderTest.java b/athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/credentials/OAuthCredentialsProviderTest.java
@@ -75,7 +75,8 @@ public class OAuthCredentialsProviderTest
     public void setUp()
     {
         secretMap = new HashMap<>();
-        credentialsProvider = new TestOAuthCredentialsProvider(TEST_SECRET_NAME, mockSecretsClient, mockHttpClient, secretMap);
+        credentialsProvider = new TestOAuthCredentialsProvider(mockHttpClient);
+        credentialsProvider.initialize(TEST_SECRET_NAME, secretMap, new CachableSecretsManager(mockSecretsClient));
     }
 
     @Test
@@ -377,9 +378,14 @@ private String createTokenResponse()
 
     private static class TestOAuthCredentialsProvider extends OAuthCredentialsProvider
     {
-        public TestOAuthCredentialsProvider(String secretName, SecretsManagerClient secretsClient, HttpClient httpClient, Map<String, String> secretMap)
+        public TestOAuthCredentialsProvider()
         {
-            super(secretName, secretMap, new CachableSecretsManager(secretsClient), httpClient);
+            super();
+        }
+        
+        public TestOAuthCredentialsProvider(HttpClient httpClient)
+        {
+            super(httpClient);
         }
 
         @Override

Original file line number	Diff line number	Diff line change
`@@ -293,7 +293,7 @@ protected CredentialsProvider getCredentialProvider()`
`293`	`293`	`return CredentialsProviderFactory.createCredentialProvider(`
`294`	`294`	`getDatabaseConnectionConfig().getSecret(),`
`295`	`295`	`getCachableSecretsManager(),`
`296`		`- DataLakeGen2OAuthCredentialsProvider.class`
	`296`	`+ new DataLakeGen2OAuthCredentialsProvider()`
`297`	`297`	`);`
`298`	`298`	`}`
`299`	`299`	`}`
Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,6 @@`
`21`	`21`
`22`	`22`	`import com.amazonaws.athena.connector.credentials.CredentialsConstants;`
`23`	`23`	`import com.amazonaws.athena.connector.credentials.OAuthCredentialsProvider;`
`24`		`-import com.amazonaws.athena.connector.lambda.security.CachableSecretsManager;`
`25`	`24`	`import com.google.common.annotations.VisibleForTesting;`
`26`	`25`
`27`	`26`	`import java.net.URI;`
`@@ -38,15 +37,15 @@ public class DataLakeGen2OAuthCredentialsProvider extends OAuthCredentialsProvid`
`38`	`37`	`private static final String SCOPE = "https://sql.azuresynapse.net/.default";`
`39`	`38`	`private static final String TENANT_ID = "tenant_id";`
`40`	`39`
`41`		`- public DataLakeGen2OAuthCredentialsProvider(String secretName, Map<String, String> secretMap, CachableSecretsManager secretsManager)`
	`40`	`+ public DataLakeGen2OAuthCredentialsProvider()`
`42`	`41`	`{`
`43`		`- super(secretName, secretMap, secretsManager);`
	`42`	`+ super();`
`44`	`43`	`}`
`45`	`44`
`46`	`45`	`@VisibleForTesting`
`47`		`- public DataLakeGen2OAuthCredentialsProvider(String secretName, Map<String, String> secretMap, CachableSecretsManager secretsManager, HttpClient httpClient)`
	`46`	`+ protected DataLakeGen2OAuthCredentialsProvider(HttpClient httpClient)`
`48`	`47`	`{`
`49`		`- super(secretName, secretMap, secretsManager, httpClient);`
	`48`	`+ super(httpClient);`
`50`	`49`	`}`
`51`	`50`
`52`	`51`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -85,7 +85,7 @@ protected CredentialsProvider getCredentialProvider()`
`85`	`85`	`return CredentialsProviderFactory.createCredentialProvider(`
`86`	`86`	`getDatabaseConnectionConfig().getSecret(),`
`87`	`87`	`getCachableSecretsManager(),`
`88`		`- DataLakeGen2OAuthCredentialsProvider.class`
	`88`	`+ new DataLakeGen2OAuthCredentialsProvider()`
`89`	`89`	`);`
`90`	`90`	`}`
`91`	`91`	`}`
Original file line number	Diff line number	Diff line change
`@@ -70,7 +70,8 @@ public void setup()`
`70`	`70`	`{`
`71`	`71`	`MockitoAnnotations.openMocks(this);`
`72`	`72`	`CachableSecretsManager cachableSecretsManager = new CachableSecretsManager(secretsManagerClient);`
`73`		`- credentialsProvider = new DataLakeGen2OAuthCredentialsProvider(SECRET_NAME, new HashMap<>(), cachableSecretsManager, httpClient);`
	`73`	`+ credentialsProvider = new DataLakeGen2OAuthCredentialsProvider(httpClient);`
	`74`	`+ credentialsProvider.initialize(SECRET_NAME, new HashMap<>(), cachableSecretsManager);`
`74`	`75`	`objectMapper = new ObjectMapper();`
`75`	`76`	`}`
`76`	`77`
Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,8 @@ public class OAuthCredentialsProviderTest`
`75`	`75`	`public void setUp()`
`76`	`76`	`{`
`77`	`77`	`secretMap = new HashMap<>();`
`78`		`- credentialsProvider = new TestOAuthCredentialsProvider(TEST_SECRET_NAME, mockSecretsClient, mockHttpClient, secretMap);`
	`78`	`+ credentialsProvider = new TestOAuthCredentialsProvider(mockHttpClient);`
	`79`	`+ credentialsProvider.initialize(TEST_SECRET_NAME, secretMap, new CachableSecretsManager(mockSecretsClient));`
`79`	`80`	`}`
`80`	`81`
`81`	`82`	`@Test`
`@@ -377,9 +378,14 @@ private String createTokenResponse()`
`377`	`378`
`378`	`379`	`private static class TestOAuthCredentialsProvider extends OAuthCredentialsProvider`
`379`	`380`	`{`
`380`		`- public TestOAuthCredentialsProvider(String secretName, SecretsManagerClient secretsClient, HttpClient httpClient, Map<String, String> secretMap)`
	`381`	`+ public TestOAuthCredentialsProvider()`
`381`	`382`	`{`
`382`		`- super(secretName, secretMap, new CachableSecretsManager(secretsClient), httpClient);`
	`383`	`+ super();`
	`384`	`+ }`
	`385`	`+`
	`386`	`+ public TestOAuthCredentialsProvider(HttpClient httpClient)`
	`387`	`+ {`
	`388`	`+ super(httpClient);`
`383`	`389`	`}`
`384`	`390`
`385`	`391`	`@Override`