merge conflict changes

Author: Trianz-Akshay

athena-snowflake/athena-snowflake.yaml

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

athena-snowflake/src/main/java/com/amazonaws/athena/connectors/snowflake/SnowflakeConstants.java

@@ -45,6 +45,11 @@ public final class SnowflakeConstants
     public static final String SNOWFLAKE_SPLIT_QUERY_ID = "query_id";
     public static final String SNOWFLAKE_SPLIT_EXPORT_BUCKET = "exportBucket";
     public static final String SNOWFLAKE_SPLIT_OBJECT_KEY = "s3ObjectKey";
+    public static final String AUTH_CODE = "auth_code";
+    public static final String CLIENT_ID = "client_id";
+    public static final String TOKEN_URL = "token_url";
+    public static final String REDIRECT_URI = "redirect_uri";
+    public static final String CLIENT_SECRET = "client_secret";
 
     private SnowflakeConstants() {}
 }

athena-snowflake/src/main/java/com/amazonaws/athena/connectors/snowflake/SnowflakeMetadataHandler.java

@@ -20,6 +20,7 @@
 
 package com.amazonaws.athena.connectors.snowflake;
 
+import com.amazonaws.athena.connector.credentials.CredentialsProvider;
 import com.amazonaws.athena.connector.lambda.QueryStatusChecker;
 import com.amazonaws.athena.connector.lambda.data.Block;
 import com.amazonaws.athena.connector.lambda.data.BlockAllocator;
@@ -765,4 +766,15 @@ public String getRoleArn(Context context)
             throw new RuntimeException("Error fetching IAM role ARN: " + e.getMessage());
         }
     }
+
+    @Override
+    protected CredentialsProvider getCredentialProvider()
+    {
+        final String secretName = getDatabaseConnectionConfig().getSecret();
+        if (StringUtils.isNotBlank(secretName)) {
+            return new SnowflakeCredentialsProvider(secretName);
+        }
+
+        return null;
+    }
 }

athena-snowflake/src/main/java/com/amazonaws/athena/connectors/snowflake/SnowflakeRecordHandler.java

@@ -19,6 +19,7 @@
  */
 package com.amazonaws.athena.connectors.snowflake;
 
+import com.amazonaws.athena.connector.credentials.CredentialsProvider;
 import com.amazonaws.athena.connector.lambda.QueryStatusChecker;
 import com.amazonaws.athena.connector.lambda.data.Block;
 import com.amazonaws.athena.connector.lambda.data.BlockSpiller;
@@ -80,6 +81,7 @@
 import software.amazon.awssdk.services.glue.model.FederationSourceErrorCode;
 import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
+import software.amazon.awssdk.utils.StringUtils;
 import software.amazon.awssdk.utils.Validate;
 
 import java.math.BigDecimal;
@@ -445,4 +447,15 @@ public PreparedStatement buildSplitSql(Connection jdbcConnection, String catalog
         }
         return preparedStatement;
     }
+
+    @Override
+    protected CredentialsProvider getCredentialProvider()
+    {
+        final String secretName = getDatabaseConnectionConfig().getSecret();
+        if (StringUtils.isNotBlank(secretName)) {
+            return new SnowflakeCredentialsProvider(secretName);
+        }
+
+        return null;
+    }
 }

athena-snowflake/src/test/java/com/amazonaws/athena/connectors/snowflake/SnowflakeMetadataHandlerTest.java

@@ -290,7 +290,7 @@ public void getPartitions() throws Exception {
 
         GetTableLayoutRequest req = new GetTableLayoutRequest(this.federatedIdentity, "queryId", "default",
                 new TableName("schema1", "table1"),
-                new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT),
+                new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Map.of(), null),
                 tableSchema,
                 partitionCols);
 
@@ -390,7 +390,7 @@ public void doGetSplits() throws Exception {
                 new TableName("schema", "table_name"),
                 partitions,
                 partitionCols,
-                new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT),
+                new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap(), null),
                 null);
         GetSplitsRequest req = new GetSplitsRequest(originalReq, null);
 
@@ -549,7 +549,7 @@ public void testGetPartitionsForView() throws Exception {
 
         GetTableLayoutRequest req = new GetTableLayoutRequest(this.federatedIdentity, "queryId", "default",
                 new TableName("testSchema", "testView"),
-                new Constraints(new HashMap<>(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT),
+                new Constraints(new HashMap<>(), Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap(), null),
                 tableSchema,
                 partitionCols);
 

athena-snowflake/src/test/java/com/amazonaws/athena/connectors/snowflake/SnowflakeRecordHandlerTest.java

@@ -111,7 +111,7 @@ public class SnowflakeRecordHandlerTest
 {
     private static final Logger logger = LoggerFactory.getLogger(SnowflakeRecordHandlerTest.class);
     private SnowflakeRecordHandler handler;
-    private FederatedIdentity identity = new FederatedIdentity("arn", "account", Collections.emptyMap(), Collections.emptyList());
+    private FederatedIdentity identity = new FederatedIdentity("arn", "account", Collections.emptyMap(), Collections.emptyList(), Collections.emptyMap());
     private EncryptionKeyFactory keyFactory = new LocalKeyFactory();
     private Connection connection;
     private static final BufferAllocator bufferAllocator = new RootAllocator();
@@ -202,7 +202,7 @@ public void doReadRecordsNoSpill()
                     TABLE_NAME,
                     schemaRoot.getSchema(),
                     splitBuilder.build(),
-                    new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT),
+                    new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap(), null),
                     100_000_000_000L,
                     100_000_000_000L//100GB don't expect this to spill
             );
@@ -263,7 +263,7 @@ public void doReadRecordsSpill()
                     TABLE_NAME,
                     schemaRoot.getSchema(),
                     splitBuilder.build(),
-                    new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT),
+                    new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT, Collections.emptyMap(), null),
                     1_500_000L, //~1.5MB so we should see some spill
                     0L
             );