Federation SDK: Query plan usage updates and fixes

athena-federation-sdk/pom.xml

@@ -290,6 +290,11 @@
             <artifactId>core</artifactId>
             <version>${io.substrait.version}</version>
         </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-params</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
     <build>
         <plugins>

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/data/Block.java

@@ -179,6 +179,34 @@ public boolean offerValue(String fieldName, int row, Object value)
         return false;
     }
 
+    /**
+     * Attempts to write the provided value to the specified field on the specified row. This method does _not_ update the
+     * row count on the underlying Apache Arrow VectorSchema. You must call setRowCount(...) to ensure the values
+     * your have written are considered 'valid rows' and thus available when you attempt to serialize this Block. This
+     * method replies on BlockUtils' field conversion/coercion logic to convert the provided value into a type that
+     * matches Apache Arrow's supported serialization format. For more details on coercion please see @BlockUtils
+     *
+     * @param fieldName The name of the field you wish to write to.
+     * @param row The row number to write to. Note that Apache Arrow Blocks begin with row 0 just like a typical array.
+     * @param value The value you wish to write.
+     * @param hasQueryPlan Whether the operation is running under a query plan, if true, bypasses constraint checks.
+     * @return True if the value was written to the Block (even if the field is missing from the Block),
+     * False if the value was not written due to failing a constraint or if query plan is present.
+     * @note This method will take no action if the provided fieldName is not a valid field in this Block's Schema.
+     * In such cases the method will return true.
+     */
+    public boolean offerValue(String fieldName, int row, Object value, boolean hasQueryPlan)
+    {
+        if (!hasQueryPlan && !constraintEvaluator.apply(fieldName, value)) {
+            return false;
+        }
+        FieldVector vector = getFieldVector(fieldName);
+        if (vector != null) {
+            BlockUtils.setValue(vector, row, value);
+        }
+        return true;
+    }
+
     /**
      * Attempts to set the provided value for the given field name and row. If the Block's schema does not
      * contain such a field, this method does nothing and returns false.

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;
+    }
 }