Address Review Comments

Author: Jithendar12

athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveJdbcConnectionFactory.java

@@ -20,6 +20,7 @@
 
 package com.amazonaws.athena.connectors.cloudera;
 
+import com.amazonaws.athena.connector.credentials.CredentialsConstants;
 import com.amazonaws.athena.connector.credentials.CredentialsProvider;
 import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig;
 import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionInfo;
@@ -62,7 +63,7 @@ public Connection getConnection(final CredentialsProvider credentialsProvider)
         if (null != credentialsProvider) {
             Matcher secretMatcher = SECRET_NAME_PATTERN.matcher(databaseConnectionConfig.getJdbcConnectionString());
             final String secretReplacement = String.format("UID=%s;PWD=%s",
-                    credentialsProvider.getCredential().getUser(), credentialsProvider.getCredential().getPassword());
+                    credentialsProvider.getCredentialMap().get(CredentialsConstants.USER), credentialsProvider.getCredentialMap().get(CredentialsConstants.PASSWORD));
             derivedJdbcString = secretMatcher.replaceAll(Matcher.quoteReplacement(secretReplacement));
         }
         else {

athena-cloudera-impala/src/main/java/com/amazonaws/athena/connectors/cloudera/ImpalaJdbcConnectionFactory.java

@@ -20,6 +20,7 @@
 
 package com.amazonaws.athena.connectors.cloudera;
 
+import com.amazonaws.athena.connector.credentials.CredentialsConstants;
 import com.amazonaws.athena.connector.credentials.CredentialsProvider;
 import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig;
 import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionInfo;
@@ -63,7 +64,7 @@ public Connection getConnection(final CredentialsProvider credentialsProvider)
             if (null != credentialsProvider) {
                 Matcher secretMatcher = SECRET_NAME_PATTERN.matcher(databaseConnectionConfig.getJdbcConnectionString());
                 final String secretReplacement = String.format("UID=%s;PWD=%s",
-                        credentialsProvider.getCredential().getUser(), credentialsProvider.getCredential().getPassword());
+                        credentialsProvider.getCredentialMap().get(CredentialsConstants.USER), credentialsProvider.getCredentialMap().get(CredentialsConstants.PASSWORD));
                 derivedJdbcString = secretMatcher.replaceAll(Matcher.quoteReplacement(secretReplacement));
             }
             else {

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

@@ -26,8 +26,5 @@ public class DataLakeGen2Constants
     public static final int DEFAULT_PORT = 1433;
     public static final String QUOTE_CHARACTER = "\"";
 
-    //Property name used to set OAuth access token in JDBC connection properties
-    public static final String ACCESS_TOKEN_PROPERTY = "accessToken";
-
     private DataLakeGen2Constants() {}
 }

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

@@ -20,6 +20,7 @@
 package com.amazonaws.athena.connectors.datalakegen2;
 
 import com.amazonaws.athena.connector.credentials.CredentialsProvider;
+import com.amazonaws.athena.connector.credentials.DefaultCredentialsProvider;
 import com.amazonaws.athena.connector.lambda.QueryStatusChecker;
 import com.amazonaws.athena.connector.lambda.data.Block;
 import com.amazonaws.athena.connector.lambda.data.BlockAllocator;
@@ -30,6 +31,7 @@
 import com.amazonaws.athena.connector.lambda.domain.Split;
 import com.amazonaws.athena.connector.lambda.domain.TableName;
 import com.amazonaws.athena.connector.lambda.domain.predicate.functions.StandardFunctions;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connector.lambda.metadata.GetDataSourceCapabilitiesRequest;
 import com.amazonaws.athena.connector.lambda.metadata.GetDataSourceCapabilitiesResponse;
 import com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest;
@@ -43,11 +45,13 @@
 import com.amazonaws.athena.connectors.datalakegen2.resolver.DataLakeGen2CaseResolver;
 import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig;
 import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionInfo;
+import com.amazonaws.athena.connectors.jdbc.connection.GenericJdbcConnectionFactory;
 import com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory;
 import com.amazonaws.athena.connectors.jdbc.manager.JDBCUtil;
 import com.amazonaws.athena.connectors.jdbc.manager.JdbcArrowTypeConverter;
 import com.amazonaws.athena.connectors.jdbc.manager.JdbcMetadataHandler;
 import com.amazonaws.athena.connectors.jdbc.resolver.JDBCCaseResolver;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -58,8 +62,11 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.services.athena.AthenaClient;
+import software.amazon.awssdk.services.glue.model.ErrorDetails;
+import software.amazon.awssdk.services.glue.model.FederationSourceErrorCode;
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
 
+import java.io.IOException;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
@@ -78,6 +85,7 @@
 public class DataLakeGen2MetadataHandler extends JdbcMetadataHandler
 {
     private static final Logger LOGGER = LoggerFactory.getLogger(DataLakeGen2MetadataHandler.class);
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
 
     static final Map<String, String> JDBC_PROPERTIES = ImmutableMap.of("databaseTerm", "SCHEMA");
     static final String PARTITION_NUMBER = "partition_number";
@@ -98,7 +106,7 @@ public DataLakeGen2MetadataHandler(java.util.Map<String, String> configOptions)
     public DataLakeGen2MetadataHandler(DatabaseConnectionConfig databaseConnectionConfig, java.util.Map<String, String> configOptions)
     {
         this(databaseConnectionConfig,
-                new DataLakeGen2JdbcConnectionFactory(databaseConnectionConfig, JDBC_PROPERTIES,
+                new GenericJdbcConnectionFactory(databaseConnectionConfig, JDBC_PROPERTIES,
                 new DatabaseConnectionInfo(DataLakeGen2Constants.DRIVER_CLASS, DataLakeGen2Constants.DEFAULT_PORT)),
                 configOptions);
     }
@@ -291,7 +299,22 @@ protected CredentialsProvider getCredentialProvider()
     {
         final String secretName = getDatabaseConnectionConfig().getSecret();
         if (StringUtils.isNotBlank(secretName)) {
-            return new DataLakeGen2CredentialsProvider(secretName);
+            try {
+                String secretString = getCachableSecretsManager().getSecret(secretName);
+                Map<String, String> secretMap = OBJECT_MAPPER.readValue(secretString, Map.class);
+
+                // Check if OAuth is configured
+                if (DataLakeGen2OAuthCredentialsProvider.isOAuthConfigured(secretMap)) {
+                    return new DataLakeGen2OAuthCredentialsProvider(secretName, secretMap, getCachableSecretsManager());
+                }
+
+                // Fall back to default credentials if OAuth is not configured
+                return new DefaultCredentialsProvider(secretString);
+            }
+            catch (IOException ioException) {
+                throw new AthenaConnectorException("Could not deserialize RDS credentials into HashMap: ",
+                        ErrorDetails.builder().errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString()).errorMessage(ioException.getMessage()).build());
+            }
         }
 
         return null;