Athena-Vertica: Substrait Plan Implementation

.gitignore

@@ -11,7 +11,7 @@ dependency-reduced-pom.xml
 */.settings/
 */.DS_Store
 .settings/org.eclipse.m2e.core.prefs
-.vscode/settings.json
+.vscode/*
 */packaged.yaml
 .DS_Store
 */.jqwik-database

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

@@ -56,6 +56,7 @@
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.services.athena.AthenaClient;
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
 
@@ -298,7 +299,7 @@ private String encodeContinuationToken(int partition)
      * @throws Exception An Exception should be thrown for database connection failures , query syntax errors and so on.
      */
     @Override
-    protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schema partitionSchema) throws Exception
+    protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schema partitionSchema, AwsRequestOverrideConfiguration requestOverrideConfiguration) throws Exception
     {
         SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
         try (ResultSet resultSet = getColumns(jdbcConnection.getCatalog(), tableName, jdbcConnection.getMetaData());

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

@@ -55,6 +55,7 @@
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.services.athena.AthenaClient;
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
 
@@ -287,7 +288,7 @@ private String encodeContinuationToken(int partition)
      * @throws Exception An Exception should be thrown for database connection failures , query syntax errors and so on.
      */
     @Override
-    protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schema partitionSchema) throws Exception
+    protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schema partitionSchema, AwsRequestOverrideConfiguration requestOverrideConfiguration) throws Exception
     {
         SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
         try (ResultSet resultSet = getColumns(jdbcConnection.getCatalog(), tableName, jdbcConnection.getMetaData());

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

@@ -7,9 +7,9 @@
  * 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.
@@ -58,6 +58,7 @@
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.services.athena.AthenaClient;
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
 
@@ -216,7 +217,7 @@ protected Optional<ArrowType> convertDatasourceTypeToArrow(int columnIndex, int
      * @throws Exception
      */
     @Override
-    protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schema partitionSchema)
+    protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schema partitionSchema, AwsRequestOverrideConfiguration requestOverrideConfiguration)
             throws Exception
     {
         LOGGER.info("Inside getSchema");
@@ -251,7 +252,7 @@ protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schem
         }
 
         String environment = DataLakeGen2Util.checkEnvironment(jdbcConnection.getMetaData().getURL());
-        
+
         if (DataLakeGen2Constants.SQL_POOL.equalsIgnoreCase(environment)) {
             // getColumns() method from SQL Server driver is causing an exception in case of Azure Serverless environment.
             // so doing explicit data type conversion

athena-db2-as400/src/main/java/com/amazonaws/athena/connectors/db2as400/Db2As400MetadataHandler.java

@@ -55,6 +55,7 @@
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.services.athena.AthenaClient;
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
 
@@ -396,7 +397,7 @@ private List<String> getTableList(final Connection connection, String query, Str
      * @throws Exception
      */
     @Override
-    protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schema partitionSchema)
+    protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schema partitionSchema, AwsRequestOverrideConfiguration requestOverrideConfiguration)
             throws Exception
     {
         SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();

athena-db2/src/main/java/com/amazonaws/athena/connectors/db2/Db2MetadataHandler.java

@@ -65,6 +65,7 @@
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.services.athena.AthenaClient;
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
 
@@ -461,7 +462,7 @@ private List<String> getTableList(final Connection connection, String schemaName
      * @throws Exception
      */
     @Override
-    protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schema partitionSchema)
+    protected Schema getSchema(Connection jdbcConnection, TableName tableName, Schema partitionSchema, AwsRequestOverrideConfiguration requestOverrideConfiguration)
             throws Exception
     {
         String typeName;

athena-federation-sdk/pom.xml

@@ -240,6 +240,11 @@
             <artifactId>junit-jupiter-api</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-params</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>junit</groupId>
             <artifactId>junit</artifactId>

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/FederationRequestHandler.java

@@ -7,9 +7,9 @@
  * 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.
@@ -19,8 +19,16 @@
  */
 package com.amazonaws.athena.connector.lambda.handlers;
 
+import com.amazonaws.athena.connector.credentials.CredentialsProvider;
+import com.amazonaws.athena.connector.credentials.DefaultCredentialsProvider;
+import com.amazonaws.athena.connector.lambda.request.FederationRequest;
+import com.amazonaws.athena.connector.lambda.security.CachableSecretsManager;
+import com.amazonaws.athena.connector.lambda.security.FederatedIdentity;
 import com.amazonaws.athena.connector.lambda.security.KmsEncryptionProvider;
 import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.auth.credentials.AwsCredentials;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
@@ -36,13 +44,113 @@
 
 public interface FederationRequestHandler extends RequestStreamHandler
 {
+    /**
+     * Gets the CachableSecretsManager instance used by this handler.
+     * Implementations must provide access to their secrets manager instance.
+     *
+     * @return The CachableSecretsManager instance
+     */
+    CachableSecretsManager getCachableSecretsManager();
+
+    /**
+     * Gets the KmsEncryptionProvider instance used by this handler.
+     * Implementations must provide access to their KMS encryption provider instance.
+     *
+     * @return The KmsEncryptionProvider instance
+     */
+    KmsEncryptionProvider getKmsEncryptionProvider();
+
+    /**
+     * Resolves any secrets found in the supplied string, for example: MyString${WithSecret} would have ${WithSecret}
+     * replaced by the corresponding value of the secret in AWS Secrets Manager with that name. If no such secret is found
+     * the function throws.
+     *
+     * @param rawString The string in which you'd like to replace SecretsManager placeholders.
+     * (e.g. ThisIsA${Secret}Here - The ${Secret} would be replaced with the contents of a SecretsManager
+     * secret called Secret. If no such secret is found, the function throws. If no ${} are found in
+     * the input string, nothing is replaced and the original string is returned.
+     * @return The processed string with secrets resolved
+     */
+    default String resolveSecrets(String rawString)
+    {
+        return getCachableSecretsManager().resolveSecrets(rawString);
+    }
+
+    /**
+     * Resolves secrets with default credentials format (username:password).
+     *
+     * @param rawString The string containing secret placeholders to resolve
+     * @return The processed string with secrets resolved in default credentials format
+     */
+    default String resolveWithDefaultCredentials(String rawString)
+    {
+        return getCachableSecretsManager().resolveWithDefaultCredentials(rawString);
+    }
+
+    /**
+     * Retrieves a secret from AWS Secrets Manager.
+     *
+     * @param secretName The name of the secret to retrieve
+     * @return The secret value
+     */
+    default String getSecret(String secretName)
+    {
+        return getCachableSecretsManager().getSecret(secretName);
+    }
+
+    /**
+     * Retrieves a secret from AWS Secrets Manager with request override configuration.
+     *
+     * @param secretName The name of the secret to retrieve
+     * @param requestOverrideConfiguration AWS request override configuration for federated requests
+     * @return The secret value
+     */
+    default String getSecret(String secretName, AwsRequestOverrideConfiguration requestOverrideConfiguration)
+    {
+        return getCachableSecretsManager().getSecret(secretName, requestOverrideConfiguration);
+    }
+
     default AwsCredentials getSessionCredentials(String kmsKeyId,
                                                  String tokenString,
                                                  KmsEncryptionProvider kmsEncryptionProvider)
     {
         return kmsEncryptionProvider.getFasCredentials(kmsKeyId, tokenString);
     }
 
+    /**
+     * Gets the AWS request override configuration for a FederationRequest.
+     * This method extracts the configuration options from the federated identity and delegates
+     * to the Map-based overload.
+     *
+     * @param request The federation request
+     * @return The AWS request override configuration, or null if not a federated request
+     */
+    default AwsRequestOverrideConfiguration getRequestOverrideConfig(FederationRequest request)
+    {
+        if (isRequestFederated(request)) {
+            FederatedIdentity federatedIdentity = request.getIdentity();
+            Map<String, String> connectorRequestOptions = federatedIdentity != null ? federatedIdentity.getConfigOptions() : null;
+
+            if (connectorRequestOptions != null && connectorRequestOptions.get(FAS_TOKEN) != null) {
+                return getRequestOverrideConfig(connectorRequestOptions);
+            }
+        }
+        return null;
+    }
+
+    /**
+     * Gets the AWS request override configuration for the given config options.
+     * This is a convenience method that delegates to the full overload using the handler's
+     * KMS encryption provider.
+     *
+     * @param configOptions The configuration options map
+     * @return The AWS request override configuration, or null if not applicable
+     */
+    default AwsRequestOverrideConfiguration getRequestOverrideConfig(Map<String, String> configOptions)
+    {
+        return getRequestOverrideConfig(configOptions, getKmsEncryptionProvider());
+    }
+
     default AwsRequestOverrideConfiguration getRequestOverrideConfig(Map<String, String> configOptions,
                                                                      KmsEncryptionProvider kmsEncryptionProvider)
     {
@@ -85,4 +193,54 @@ default AthenaClient getAthenaClient(AwsRequestOverrideConfiguration awsRequestO
             return defaultAthena;
         }
     }
+
+    default boolean isRequestFederated(FederationRequest req)
+    {
+        FederatedIdentity federatedIdentity = req.getIdentity();
+        Map<String, String> connectorRequestOptions = federatedIdentity != null ? federatedIdentity.getConfigOptions() : null;
+        return (connectorRequestOptions != null && connectorRequestOptions.get(FAS_TOKEN) != null);
+    }
+
+    /**
+     * Gets a credentials provider for database connections with optional request override configuration.
+     * This method checks if a secret name is configured and creates a credentials provider if available.
+     * Subclasses can override createCredentialsProvider() to provide custom credential provider implementations.
+     *
+     * @param requestOverrideConfiguration Optional AWS request override configuration for federated requests
+     * @return CredentialsProvider instance or null if no secret is configured
+     */
+    default CredentialsProvider getCredentialProvider(AwsRequestOverrideConfiguration requestOverrideConfiguration)
+    {
+        final String secretName = getDatabaseConnectionSecret();
+        if (StringUtils.isNotBlank(secretName)) {
+            Logger logger = LoggerFactory.getLogger(this.getClass());
+            logger.info("Using Secrets Manager.");
+            return createCredentialsProvider(secretName, requestOverrideConfiguration);
+        }
+        return null;
+    }
+
+    /**
+     * Factory method to create CredentialsProvider. Subclasses can override this to provide
+     * custom credential provider implementations (e.g., SnowflakeCredentialsProvider).
+     *
+     * @param secretName The secret name to retrieve credentials from
+     * @param requestOverrideConfiguration Optional AWS request override configuration
+     * @return CredentialsProvider instance
+     */
+    default CredentialsProvider createCredentialsProvider(String secretName, AwsRequestOverrideConfiguration requestOverrideConfiguration)
+    {
+        return new DefaultCredentialsProvider(getSecret(secretName, requestOverrideConfiguration));
+    }
+
+    /**
+     * Gets the database connection secret name. Subclasses that use database credentials
+     * should override this method to provide the secret name from their configuration.
+     *
+     * @return The secret name, or null if not applicable
+     */
+    default String getDatabaseConnectionSecret()
+    {
+        return null;
+    }
 }