refactor to add feature flag for s3 export approach

Author: Trianz-Akshay

athena-snowflake/athena-snowflake-connection.yaml

@@ -42,6 +42,13 @@ Parameters:
     Description: "(Optional) A custom role to be used by the Connector lambda"
     Type: String
     Default: ""
+  EnableS3Export:
+    Description: '(Optional) Enable S3 export functionality for data transfer. Set to true to use S3 export path, false for direct query path.'
+    Type: String
+    Default: 'false'
+    AllowedValues:
+      - 'true'
+      - 'false'
 
 Conditions:
   HasSecurityGroups: !Not [ !Equals [ !Join ["", !Ref SecurityGroupIds], "" ] ]
@@ -59,6 +66,7 @@ Resources:
       Environment:
         Variables:
           glue_connection: !Ref GlueConnection
+          SNOWFLAKE_ENABLE_S3_EXPORT: !Ref EnableS3Export
       FunctionName: !Ref LambdaFunctionName
       PackageType: "Image"
       ImageUri: !Sub
@@ -87,7 +95,6 @@ Resources:
             Principal:
               Service:
                 - lambda.amazonaws.com
-              AWS: '*'
             Action:
               - "sts:AssumeRole"
 

athena-snowflake/athena-snowflake.yaml

@@ -58,6 +58,13 @@ Parameters:
     Description: "(Optional) A custom role to be used by the Connector lambda"
     Type: String
     Default: ""
+  EnableS3Export:
+    Description: '(Optional) Enable S3 export functionality for data transfer. Set to true to use S3 export path, false for direct query path.'
+    Type: String
+    Default: 'false'
+    AllowedValues:
+      - 'true'
+      - 'false'
 Conditions:
   HasPermissionsBoundary: !Not [ !Equals [ !Ref PermissionsBoundaryARN, "" ] ]
   HasSecurityGroups: !Not [ !Equals [ !Join ["", !Ref SecurityGroupIds], "" ] ]
@@ -75,13 +82,14 @@ Resources:
           spill_bucket: !Ref SpillBucket
           spill_prefix: !Ref SpillPrefix
           default: !Ref DefaultConnectionString
+          SNOWFLAKE_ENABLE_S3_EXPORT: !Ref EnableS3Export
       FunctionName: !Ref LambdaFunctionName
       PackageType: "Image"
       ImageUri: !Sub
         - '${Account}.dkr.ecr.${AWS::Region}.amazonaws.com/athena-federation-repository-snowflake:2022.47.1'
         - Account: !If [IsRegionBAH, 084828588479, !If [IsRegionHKG, 183295418215, 292517598671]]
       ImageConfig:
-        Command: [ "com.amazonaws.athena.connectors.snowflake.SnowflakeCompositeHandler" ]
+        Command: [ "com.amazonaws.athena.connectors.snowflake.SnowflakeMuxCompositeHandler" ]
       Description: "Enables Amazon Athena to communicate with Snowflake using JDBC"
       Timeout: !Ref LambdaTimeout
       MemorySize: !Ref LambdaMemory
@@ -104,7 +112,6 @@ Resources:
             Principal:
               Service:
                 - lambda.amazonaws.com
-              AWS: '*'
             Action:
               - "sts:AssumeRole"
 
@@ -137,6 +144,13 @@ Resources:
               - athena:GetQueryExecution
             Effect: Allow
             Resource: '*'
+          - Action:
+              - ec2:CreateNetworkInterface
+              - ec2:DeleteNetworkInterface
+              - ec2:DescribeNetworkInterfaces
+              - ec2:DetachNetworkInterface
+            Effect: Allow
+            Resource: '*'
         #S3CrudPolicy allows our connector to spill large responses to S3. You can optionally replace this pre-made policy
         #with one that is more restrictive and can only 'put' but not read,delete, or overwrite files.
           - Action:

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

@@ -42,7 +42,7 @@ public class SnowflakeCompositeHandler
 {
     public SnowflakeCompositeHandler() throws CertificateEncodingException, IOException, NoSuchAlgorithmException, KeyStoreException
     {
-        super(new SnowflakeMetadataHandler(new SnowflakeEnvironmentProperties().createEnvironment()), new SnowflakeRecordHandler(new SnowflakeEnvironmentProperties().createEnvironment()));
+        super(new SnowflakeMetadataHandler(new SnowflakeEnvironmentProperties(System.getenv()).createEnvironment()), new SnowflakeRecordHandler(new SnowflakeEnvironmentProperties(System.getenv()).createEnvironment()));
         installCaCertificate();
         setupNativeEnvironmentVariables();
     }

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

@@ -25,9 +25,15 @@ public final class SnowflakeConstants
     public static final String SNOWFLAKE_NAME = "snowflake";
     public static final String SNOWFLAKE_DRIVER_CLASS = "com.snowflake.client.jdbc.SnowflakeDriver";
     public static final int SNOWFLAKE_DEFAULT_PORT = 1025;
+    /**
+     * This constant limits the number of partitions. The default set to 50. A large number may cause a timeout issue.
+     * We arrived at this number after performance testing with datasets of different size
+     */
+    public static final int MAX_PARTITION_COUNT = 50;
     /**
      * This constant limits the number of records to be returned in a single split.
      */
+    public static final int SINGLE_SPLIT_LIMIT_COUNT = 10000;
     public static final String SNOWFLAKE_QUOTE_CHARACTER = "\"";
     /**
      * A ssl file location constant to store the SSL certificate

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

@@ -43,6 +43,14 @@ public class SnowflakeEnvironmentProperties extends JdbcEnvironmentProperties
     private static final String DB_PROPERTY_KEY = "db";
     private static final String SCHEMA_PROPERTY_KEY = "schema";
     private static final String SNOWFLAKE_ESCAPE_CHARACTER = "\"";
+    public static final String ENABLE_S3_EXPORT = "SNOWFLAKE_ENABLE_S3_EXPORT";
+
+    private final boolean enableS3Export;
+
+    public SnowflakeEnvironmentProperties(Map<String, String> properties)
+    {
+        this.enableS3Export = Boolean.parseBoolean(properties.getOrDefault(ENABLE_S3_EXPORT, "false"));
+    }
 
     @Override
     public Map<String, String> connectionPropertiesToEnvironment(Map<String, String> connectionProperties)
@@ -134,4 +142,9 @@ public static Map<String, String> getSnowFlakeParameter(Map<String, String> base
 
         return parameters;
     }
+
+    public boolean isS3ExportEnabled()
+    {
+        return enableS3Export;
+    }
 }