Extract CredentialsProvider instance creation logic to SDK

Author: Jithendar12

athena-datalakegen2/src/main/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2CredentialProviderUtils.java

@@ -20,6 +20,7 @@
 package com.amazonaws.athena.connectors.datalakegen2;
 
 import com.amazonaws.athena.connector.credentials.CredentialsProvider;
+import com.amazonaws.athena.connector.credentials.CredentialsProviderFactory;
 import com.amazonaws.athena.connector.lambda.QueryStatusChecker;
 import com.amazonaws.athena.connector.lambda.data.Block;
 import com.amazonaws.athena.connector.lambda.data.BlockAllocator;
@@ -289,9 +290,10 @@ protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schem
     @Override
     protected CredentialsProvider getCredentialProvider()
     {
-        return DataLakeGen2CredentialProviderUtils.getCredentialProvider(
+        return CredentialsProviderFactory.createCredentialProvider(
             getDatabaseConnectionConfig().getSecret(),
-            getCachableSecretsManager()
+            getCachableSecretsManager(),
+            DataLakeGen2OAuthCredentialsProvider.class
         );
     }
 }

athena-datalakegen2/src/main/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2MetadataHandler.java

@@ -49,7 +49,8 @@ public DataLakeGen2OAuthCredentialsProvider(String secretName, Map<String, Strin
         super(secretName, secretMap, secretsManager, httpClient);
     }
 
-    protected static boolean isOAuthConfigured(Map<String, String> secretMap)
+    @Override
+    protected boolean isOAuthConfigured(Map<String, String> secretMap)
     {
         return secretMap.containsKey(CredentialsConstants.CLIENT_ID) &&
                !secretMap.get(CredentialsConstants.CLIENT_ID).isEmpty() &&

athena-datalakegen2/src/main/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2OAuthCredentialsProvider.java

@@ -19,6 +19,7 @@
  */
 package com.amazonaws.athena.connectors.datalakegen2;
 import com.amazonaws.athena.connector.credentials.CredentialsProvider;
+import com.amazonaws.athena.connector.credentials.CredentialsProviderFactory;
 import com.amazonaws.athena.connector.lambda.domain.Split;
 import com.amazonaws.athena.connector.lambda.domain.TableName;
 import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints;
@@ -81,9 +82,10 @@ public PreparedStatement buildSplitSql(Connection jdbcConnection, String catalog
     @Override
     protected CredentialsProvider getCredentialProvider()
     {
-        return DataLakeGen2CredentialProviderUtils.getCredentialProvider(
+        return CredentialsProviderFactory.createCredentialProvider(
             getDatabaseConnectionConfig().getSecret(),
-            getCachableSecretsManager()
+            getCachableSecretsManager(),
+            DataLakeGen2OAuthCredentialsProvider.class
         );
     }
 }

athena-datalakegen2/src/main/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2RecordHandler.java

@@ -82,7 +82,7 @@ public void testIsOAuthConfigured_WithValidConfig()
         secretMap.put(CLIENT_SECRET, TEST_CLIENT_SECRET);
         secretMap.put(TENANT_ID, TEST_TENANT_ID);
 
-        assertTrue(DataLakeGen2OAuthCredentialsProvider.isOAuthConfigured(secretMap));
+        assertTrue(credentialsProvider.isOAuthConfigured(secretMap));
     }
 
     @Test
@@ -92,7 +92,7 @@ public void testIsOAuthConfigured_WithMissingConfig()
         secretMap.put(CLIENT_ID, TEST_CLIENT_ID);
         // Missing client_secret and tenant_id
 
-        assertFalse(DataLakeGen2OAuthCredentialsProvider.isOAuthConfigured(secretMap));
+        assertFalse(credentialsProvider.isOAuthConfigured(secretMap));
     }
 
     @Test

athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2CredentialProviderUtilsTest.java

@@ -0,0 +1,101 @@
+/*-
+ * #%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.exceptions.AthenaConnectorException;
+import com.amazonaws.athena.connector.lambda.security.CachableSecretsManager;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.commons.lang3.StringUtils;
+import software.amazon.awssdk.services.glue.model.ErrorDetails;
+import software.amazon.awssdk.services.glue.model.FederationSourceErrorCode;
+
+import java.io.IOException;
+import java.lang.reflect.Constructor;
+import java.util.Map;
+
+/**
+ * Factory class for handling credentials provider creation.
+ * This class can be used by any connector that needs to support both
+ * OAuth and username/password authentication.
+ */
+public final class CredentialsProviderFactory
+{
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+    private CredentialsProviderFactory()
+    {
+    }
+
+    /**
+     * Creates a credentials provider based on the secret configuration.
+     * If OAuth is configured (determined by the provider's isOAuthConfigured method), creates an instance
+     * of the specified OAuth provider. Otherwise, creates a default username/password credentials provider.
+     *
+     * @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)
+     * @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(
+            String secretName,
+            CachableSecretsManager secretsManager,
+            Class<T> oAuthProviderClass)
+    {
+        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);
+                }
+                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);
+            }
+            catch (IOException ioException) {
+                throw new AthenaConnectorException("Could not deserialize credentials into HashMap: ",
+                        ErrorDetails.builder()
+                                .errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
+                                .errorMessage(ioException.getMessage())
+                                .build());
+            }
+        }
+
+        return null;
+    }
+}

athena-datalakegen2/src/test/java/com/amazonaws/athena/connectors/datalakegen2/DataLakeGen2OAuthCredentialsProviderTest.java

@@ -106,6 +106,11 @@ public Credentials getCredential()
         }
     }
 
+    /**
+     * Checks if OAuth is configured by verifying required fields exist.
+     */
+    protected abstract boolean isOAuthConfigured(Map<String, String> secretMap);
+
     /**
      * Builds the token request for the specific OAuth provider.
      */