Abstract common OAuth handling and add OAuth support to Athena DataLake Gen2 Connector (#2932)

Co-authored-by: burhan94 <[email protected]>

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/athena-datalakegen2-connection.yaml

@@ -104,6 +104,7 @@ Resources:
         Statement:
           - Action:
               - secretsmanager:GetSecretValue
+              - secretsmanager:PutSecretValue
             Effect: Allow
             Resource: !Sub 'arn:${AWS::Partition}:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${SecretName}*'
           - Action:

athena-datalakegen2/athena-datalakegen2-package.yaml

@@ -82,6 +82,7 @@ Resources:
         - Statement:
             - Action:
                 - secretsmanager:GetSecretValue
+                - secretsmanager:PutSecretValue
               Effect: Allow
               Resource: !Sub 'arn:${AWS::Partition}:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${SecretNamePrefix}*'
           Version: '2012-10-17'

athena-datalakegen2/athena-datalakegen2.yaml

@@ -87,6 +87,7 @@ Resources:
         - Statement:
             - Action:
                 - secretsmanager:GetSecretValue
+                - secretsmanager:PutSecretValue
               Effect: Allow
               Resource: !Sub 'arn:${AWS::Partition}:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${SecretNamePrefix}*'
           Version: '2012-10-17'

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

@@ -19,6 +19,8 @@
  */
 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;
@@ -42,6 +44,7 @@
 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;
@@ -96,7 +99,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);
     }
@@ -283,4 +286,14 @@ protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schem
             return schemaBuilder.build();
         }
     }
+
+    @Override
+    protected CredentialsProvider getCredentialProvider()
+    {
+        return CredentialsProviderFactory.createCredentialProvider(
+            getDatabaseConnectionConfig().getSecret(),
+            getCachableSecretsManager(),
+            new DataLakeGen2OAuthCredentialsProvider()
+        );
+    }
 }

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

@@ -0,0 +1,80 @@
+/*-
+ * #%L
+ * athena-datalakegen2
+ * %%
+ * 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.connectors.datalakegen2;
+
+import com.amazonaws.athena.connector.credentials.CredentialsConstants;
+import com.amazonaws.athena.connector.credentials.OAuthCredentialsProvider;
+import com.google.common.annotations.VisibleForTesting;
+
+import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.util.Map;
+
+/**
+ * OAuth credentials provider for Azure Data Lake Gen2.
+ */
+public class DataLakeGen2OAuthCredentialsProvider extends OAuthCredentialsProvider
+{
+    private static final String TOKEN_ENDPOINT_FORMAT = "https://login.microsoftonline.com/%s/oauth2/v2.0/token";
+    private static final String SCOPE = "https://sql.azuresynapse.net/.default";
+    private static final String TENANT_ID = "tenant_id";
+
+    public DataLakeGen2OAuthCredentialsProvider()
+    {
+        super();
+    }
+
+    @VisibleForTesting
+    protected DataLakeGen2OAuthCredentialsProvider(HttpClient httpClient)
+    {
+        super(httpClient);
+    }
+
+    @Override
+    protected boolean isOAuthConfigured(Map<String, String> secretMap)
+    {
+        return secretMap.containsKey(CredentialsConstants.CLIENT_ID) &&
+               !secretMap.get(CredentialsConstants.CLIENT_ID).isEmpty() &&
+               secretMap.containsKey(CredentialsConstants.CLIENT_SECRET) &&
+               !secretMap.get(CredentialsConstants.CLIENT_SECRET).isEmpty() &&
+               secretMap.containsKey(TENANT_ID) &&
+               !secretMap.get(TENANT_ID).isEmpty();
+    }
+
+    @Override
+    protected HttpRequest buildTokenRequest(Map<String, String> secretMap)
+    {
+        String clientId = secretMap.get(CredentialsConstants.CLIENT_ID);
+        String clientSecret = secretMap.get(CredentialsConstants.CLIENT_SECRET);
+        String tenantId = secretMap.get(TENANT_ID);
+        String tokenEndpoint = String.format(TOKEN_ENDPOINT_FORMAT, tenantId);
+
+        String formData = String.format(
+                "grant_type=client_credentials&scope=%s&client_id=%s&client_secret=%s",
+                SCOPE, clientId, clientSecret);
+
+        return HttpRequest.newBuilder()
+            .uri(URI.create(tokenEndpoint))
+            .header("Content-Type", "application/x-www-form-urlencoded")
+            .POST(HttpRequest.BodyPublishers.ofString(formData))
+            .build();
+    }
+}

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

@@ -18,6 +18,8 @@
  * #L%
  */
 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;
@@ -76,4 +78,14 @@ public PreparedStatement buildSplitSql(Connection jdbcConnection, String catalog
         preparedStatement.setFetchSize(FETCH_SIZE);
         return preparedStatement;
     }
+
+    @Override
+    protected CredentialsProvider getCredentialProvider()
+    {
+        return CredentialsProviderFactory.createCredentialProvider(
+            getDatabaseConnectionConfig().getSecret(),
+            getCachableSecretsManager(),
+            new DataLakeGen2OAuthCredentialsProvider()
+        );
+    }
 }