Add snowflake JDBC Connector

athena-jdbc/pom.xml

@@ -28,6 +28,11 @@
                 </exclusion>
             </exclusions>
         </dependency>
+        <dependency>
+            <groupId>net.snowflake</groupId>
+            <artifactId>snowflake-jdbc</artifactId>
+            <version>3.13.4</version>
+        </dependency>
         <dependency>
             <groupId>com.amazonaws</groupId>
             <artifactId>athena-federation-integ-test</artifactId>

athena-jdbc/src/main/java/com/amazonaws/connectors/athena/jdbc/connection/GenericJdbcConnectionFactory.java

@@ -56,13 +56,17 @@ public class GenericJdbcConnectionFactory
     private static final String REDSHIFT_DRIVER_CLASS = "com.amazon.redshift.jdbc.Driver";
     private static final int REDSHIFT_DEFAULT_PORT = 5439;
 
+    private static final String SNOWFLAKE_DRIVER_CLASS = "net.snowflake.client.jdbc.SnowflakeDriver";
+    private static final int SNOWFLAKE_DEFAULT_PORT = 443;
+
     private static final String SECRET_NAME_PATTERN_STRING = "(\\$\\{[a-zA-Z0-9:/_+=.@-]+})";
     public static final Pattern SECRET_NAME_PATTERN = Pattern.compile(SECRET_NAME_PATTERN_STRING);
 
     private static final ImmutableMap<DatabaseEngine, DatabaseConnectionInfo> CONNECTION_INFO = ImmutableMap.of(
             DatabaseEngine.MYSQL, new DatabaseConnectionInfo(MYSQL_DRIVER_CLASS, MYSQL_DEFAULT_PORT),
             DatabaseEngine.POSTGRES, new DatabaseConnectionInfo(POSTGRESQL_DRIVER_CLASS, POSTGRESQL_DEFAULT_PORT),
-            DatabaseEngine.REDSHIFT, new DatabaseConnectionInfo(REDSHIFT_DRIVER_CLASS, REDSHIFT_DEFAULT_PORT));
+            DatabaseEngine.REDSHIFT, new DatabaseConnectionInfo(REDSHIFT_DRIVER_CLASS, REDSHIFT_DEFAULT_PORT),
+            DatabaseEngine.SNOWFLAKE, new DatabaseConnectionInfo(SNOWFLAKE_DRIVER_CLASS, SNOWFLAKE_DEFAULT_PORT));
 
     private final DatabaseConnectionConfig databaseConnectionConfig;
     private final Properties jdbcProperties;

athena-jdbc/src/main/java/com/amazonaws/connectors/athena/jdbc/connection/JdbcConnectionFactory.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.
@@ -41,7 +41,8 @@ enum DatabaseEngine
     {
         MYSQL("mysql"),
         POSTGRES("postgres"),
-        REDSHIFT("redshift");
+        REDSHIFT("redshift"),
+        SNOWFLAKE("snowflake");
 
         private final String dbName;
 

athena-jdbc/src/main/java/com/amazonaws/connectors/athena/jdbc/manager/JDBCUtil.java

@@ -26,6 +26,8 @@
 import com.amazonaws.connectors.athena.jdbc.mysql.MySqlRecordHandler;
 import com.amazonaws.connectors.athena.jdbc.postgresql.PostGreSqlMetadataHandler;
 import com.amazonaws.connectors.athena.jdbc.postgresql.PostGreSqlRecordHandler;
+import com.amazonaws.connectors.athena.jdbc.snowflake.SnowflakeMetadataHandler;
+import com.amazonaws.connectors.athena.jdbc.snowflake.SnowflakeRecordHandler;
 import com.google.common.collect.ImmutableMap;
 import org.apache.commons.lang3.Validate;
 
@@ -100,6 +102,8 @@ private static JdbcMetadataHandler createJdbcMetadataHandler(final DatabaseConne
         switch (databaseConnectionConfig.getType()) {
             case MYSQL:
                 return new MySqlMetadataHandler(databaseConnectionConfig);
+            case SNOWFLAKE:
+                return new SnowflakeMetadataHandler(databaseConnectionConfig);
             case REDSHIFT:
             case POSTGRES:
                 return new PostGreSqlMetadataHandler(databaseConnectionConfig);
@@ -149,9 +153,11 @@ private static JdbcRecordHandler createJdbcRecordHandler(final DatabaseConnectio
         switch (databaseConnectionConfig.getType()) {
             case MYSQL:
                 return new MySqlRecordHandler(databaseConnectionConfig);
-            case REDSHIFT:
             case POSTGRES:
+            case REDSHIFT:
                 return new PostGreSqlRecordHandler(databaseConnectionConfig);
+            case SNOWFLAKE:
+                return new SnowflakeRecordHandler(databaseConnectionConfig);
             default:
                 throw new RuntimeException("Mux: Unhandled database engine " + databaseConnectionConfig.getType());
         }

athena-jdbc/src/main/java/com/amazonaws/connectors/athena/jdbc/manager/JdbcMetadataHandler.java

@@ -266,7 +266,7 @@ private Schema getSchema(Connection jdbcConnection, TableName tableName, Schema
         }
     }
 
-    private ResultSet getColumns(final String catalogName, final TableName tableHandle, final DatabaseMetaData metadata)
+    protected ResultSet getColumns(final String catalogName, final TableName tableHandle, final DatabaseMetaData metadata)
             throws SQLException
     {
         String escape = metadata.getSearchStringEscape();

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

@@ -0,0 +1,125 @@
+/*-
+ * #%L
+ * athena-jdbc
+ * %%
+ * Copyright (C) 2019 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.connectors.athena.jdbc.snowflake;
+
+import com.amazonaws.athena.connector.lambda.QueryStatusChecker;
+import com.amazonaws.athena.connector.lambda.data.BlockAllocator;
+import com.amazonaws.athena.connector.lambda.data.BlockWriter;
+import com.amazonaws.athena.connector.lambda.data.SchemaBuilder;
+import com.amazonaws.athena.connector.lambda.domain.Split;
+import com.amazonaws.athena.connector.lambda.domain.TableName;
+import com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest;
+import com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse;
+import com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest;
+import com.amazonaws.connectors.athena.jdbc.connection.DatabaseConnectionConfig;
+import com.amazonaws.connectors.athena.jdbc.connection.GenericJdbcConnectionFactory;
+import com.amazonaws.connectors.athena.jdbc.connection.JdbcConnectionFactory;
+import com.amazonaws.connectors.athena.jdbc.manager.JDBCUtil;
+import com.amazonaws.connectors.athena.jdbc.manager.JdbcMetadataHandler;
+import com.amazonaws.services.athena.AmazonAthena;
+import com.amazonaws.services.secretsmanager.AWSSecretsManager;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.sql.DatabaseMetaData;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Handles metadata for Snowflake.
+ */
+public class SnowflakeMetadataHandler
+        extends JdbcMetadataHandler
+{
+    static final Map<String, String> JDBC_PROPERTIES = ImmutableMap.of("databaseTerm", "SCHEMA");
+    private static final Logger LOGGER = LoggerFactory.getLogger(SnowflakeMetadataHandler.class);
+
+    /**
+     * Instantiates handler to be used by Lambda function directly.
+     *
+     * Recommend using {@link com.amazonaws.connectors.athena.jdbc.MultiplexingJdbcCompositeHandler} instead.
+     */
+    public SnowflakeMetadataHandler()
+    {
+        this(JDBCUtil.getSingleDatabaseConfigFromEnv(JdbcConnectionFactory.DatabaseEngine.SNOWFLAKE));
+    }
+
+    public SnowflakeMetadataHandler(final DatabaseConnectionConfig databaseConnectionConfig)
+    {
+        super(databaseConnectionConfig, new GenericJdbcConnectionFactory(databaseConnectionConfig, JDBC_PROPERTIES));
+    }
+
+    @VisibleForTesting
+    protected SnowflakeMetadataHandler(final DatabaseConnectionConfig databaseConnectionConfig, final AWSSecretsManager secretsManager,
+            final AmazonAthena athena, final JdbcConnectionFactory jdbcConnectionFactory)
+    {
+        super(databaseConnectionConfig, secretsManager, athena, jdbcConnectionFactory);
+    }
+
+    @Override
+    protected String escapeNamePattern(final String name, final String escape)
+    {
+        return super.escapeNamePattern(name, escape).toUpperCase();
+    }
+
+    @Override
+    protected ResultSet getColumns(final String catalogName, final TableName tableHandle, final DatabaseMetaData metadata)
+            throws SQLException
+    {
+        String escape = metadata.getSearchStringEscape();
+        return metadata.getColumns(
+                catalogName,
+                escapeNamePattern(tableHandle.getSchemaName(), escape),
+                escapeNamePattern(tableHandle.getTableName(), escape),
+                null);
+    }
+
+    @Override
+    public Schema getPartitionSchema(final String catalogName)
+    {
+        SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
+        return schemaBuilder.build();
+    }
+
+    @Override
+    public void getPartitions(final BlockWriter blockWriter, final GetTableLayoutRequest getTableLayoutRequest, QueryStatusChecker queryStatusChecker)
+    { }
+
+    @Override
+    public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest getSplitsRequest)
+    {
+        LOGGER.info("doGetSplits: enter - " + getSplitsRequest);
+
+        String catalogName = getSplitsRequest.getCatalogName();
+        Set<Split> splits = new HashSet<>();
+
+        Split split = Split.newBuilder(makeSpillLocation(getSplitsRequest), makeEncryptionKey()).build();
+        splits.add(split);
+
+        LOGGER.info("doGetSplits: exit - " + splits.size());
+        return new GetSplitsResponse(catalogName, splits);
+    }
+}

athena-jdbc/src/main/java/com/amazonaws/connectors/athena/jdbc/snowflake/SnowflakeQueryStringBuilder.java

@@ -0,0 +1,62 @@
+/*-
+ * #%L
+ * athena-jdbc
+ * %%
+ * Copyright (C) 2019 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.connectors.athena.jdbc.snowflake;
+
+import com.amazonaws.athena.connector.lambda.domain.Split;
+import com.amazonaws.connectors.athena.jdbc.manager.JdbcSplitQueryBuilder;
+import com.google.common.base.Strings;
+
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Extends {@link JdbcSplitQueryBuilder} and implements PostGreSql specific SQL clauses for split.
+ *
+ * PostGreSql partitions through child tables that can be used in a FROM clause.
+ */
+public class SnowflakeQueryStringBuilder
+        extends JdbcSplitQueryBuilder
+{
+    SnowflakeQueryStringBuilder(final String quoteCharacters)
+    {
+        super(quoteCharacters);
+    }
+
+    @Override
+    protected String getFromClauseWithSplit(String catalog, String schema, String table, Split split)
+    {
+        StringBuilder tableName = new StringBuilder();
+        if (!Strings.isNullOrEmpty(catalog)) {
+            tableName.append(quote(catalog.toUpperCase())).append('.');
+        }
+        if (!Strings.isNullOrEmpty(schema)) {
+            tableName.append(quote(schema.toUpperCase())).append('.');
+        }
+        tableName.append(quote(table.toUpperCase()));
+
+        return String.format(" FROM %s ", tableName);
+    }
+
+    @Override
+    protected List<String> getPartitionWhereClauses(final Split split)
+    {
+        return Collections.emptyList();
+    }
+}

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

@@ -0,0 +1,93 @@
+/*-
+ * #%L
+ * athena-jdbc
+ * %%
+ * Copyright (C) 2019 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.connectors.athena.jdbc.snowflake;
+
+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;
+import com.amazonaws.connectors.athena.jdbc.connection.DatabaseConnectionConfig;
+import com.amazonaws.connectors.athena.jdbc.connection.GenericJdbcConnectionFactory;
+import com.amazonaws.connectors.athena.jdbc.connection.JdbcConnectionFactory;
+import com.amazonaws.connectors.athena.jdbc.manager.JDBCUtil;
+import com.amazonaws.connectors.athena.jdbc.manager.JdbcRecordHandler;
+import com.amazonaws.connectors.athena.jdbc.manager.JdbcSplitQueryBuilder;
+import com.amazonaws.services.athena.AmazonAthena;
+import com.amazonaws.services.athena.AmazonAthenaClientBuilder;
+import com.amazonaws.services.s3.AmazonS3;
+import com.amazonaws.services.s3.AmazonS3ClientBuilder;
+import com.amazonaws.services.secretsmanager.AWSSecretsManager;
+import com.amazonaws.services.secretsmanager.AWSSecretsManagerClientBuilder;
+import com.google.common.annotations.VisibleForTesting;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.commons.lang3.Validate;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+
+public class SnowflakeRecordHandler
+        extends JdbcRecordHandler
+{
+    private static final Logger LOGGER = LoggerFactory.getLogger(SnowflakeRecordHandler.class);
+
+    private static final int FETCH_SIZE = 1000;
+
+    private final JdbcSplitQueryBuilder jdbcSplitQueryBuilder;
+
+    private static final String SNOWFLAKE_QUOTE_CHARACTER = "\"";
+
+    /**
+     * Instantiates handler to be used by Lambda function directly.
+     *
+     * Recommend using {@link com.amazonaws.connectors.athena.jdbc.MultiplexingJdbcCompositeHandler} instead.
+     */
+    public SnowflakeRecordHandler()
+    {
+        this(JDBCUtil.getSingleDatabaseConfigFromEnv(JdbcConnectionFactory.DatabaseEngine.SNOWFLAKE));
+    }
+
+    public SnowflakeRecordHandler(final DatabaseConnectionConfig databaseConnectionConfig)
+    {
+        this(databaseConnectionConfig, AmazonS3ClientBuilder.defaultClient(), AWSSecretsManagerClientBuilder.defaultClient(), AmazonAthenaClientBuilder.defaultClient(),
+                new GenericJdbcConnectionFactory(databaseConnectionConfig, SnowflakeMetadataHandler.JDBC_PROPERTIES), new SnowflakeQueryStringBuilder(SNOWFLAKE_QUOTE_CHARACTER));
+    }
+
+    @VisibleForTesting
+    SnowflakeRecordHandler(final DatabaseConnectionConfig databaseConnectionConfig, final AmazonS3 amazonS3, final AWSSecretsManager secretsManager,
+            final AmazonAthena athena, final JdbcConnectionFactory jdbcConnectionFactory, final JdbcSplitQueryBuilder jdbcSplitQueryBuilder)
+    {
+        super(amazonS3, secretsManager, athena, databaseConnectionConfig, jdbcConnectionFactory);
+        this.jdbcSplitQueryBuilder = Validate.notNull(jdbcSplitQueryBuilder, "query builder must not be null");
+    }
+
+    @Override
+    public PreparedStatement buildSplitSql(Connection jdbcConnection, String catalogName, TableName tableName, Schema schema, Constraints constraints, Split split)
+            throws SQLException
+    {
+        PreparedStatement preparedStatement = jdbcSplitQueryBuilder.buildSql(jdbcConnection, null, tableName.getSchemaName(), tableName.getTableName(), schema, constraints, split);
+
+        // Disable fetching all rows.
+        preparedStatement.setFetchSize(FETCH_SIZE);
+
+        return preparedStatement;
+    }
+}