Upgrade Snowflake JDBC client to v4 (#37885)

* Upgrade Snowflake driver to v4

* SnowflakeIO: add account parameter

Author: turb

sdks/java/io/snowflake/build.gradle

@@ -30,16 +30,16 @@ dependencies {
   implementation project(path: ":sdks:java:extensions:google-cloud-platform-core")
   permitUnusedDeclared project(path: ":sdks:java:extensions:google-cloud-platform-core")
   implementation library.java.slf4j_api
-  implementation group: 'net.snowflake', name: 'snowflake-jdbc', version: '3.20.0'
-  implementation group: 'com.opencsv', name: 'opencsv', version: '5.0'
-  implementation 'net.snowflake:snowflake-ingest-sdk:0.9.9'
+  implementation group: 'net.snowflake', name: 'snowflake-jdbc', version: '4.0.2'
+  implementation group: 'com.opencsv', name: 'opencsv', version: '5.12.0'
+  implementation 'net.snowflake:snowflake-ingest-sdk:4.4.2'
   implementation "org.bouncycastle:bcprov-jdk15on:1.70"
   implementation library.java.joda_time
   testImplementation project(path: ":sdks:java:core", configuration: "shadowTest")
   testImplementation project(path: ":sdks:java:extensions:avro", configuration: "testRuntimeMigration")
   testImplementation project(path: ":sdks:java:io:common")
   testImplementation project(path: ":sdks:java:testing:test-utils")
-  testImplementation 'com.google.cloud:google-cloud-storage:1.102.0'
+  testImplementation 'com.google.cloud:google-cloud-storage:1.118.1'
   testImplementation library.java.avro
   testImplementation library.java.junit
   testImplementation library.java.slf4j_api

sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeIO.java

@@ -36,7 +36,8 @@
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import javax.sql.DataSource;
-import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
+import net.snowflake.client.api.datasource.SnowflakeDataSource;
+import net.snowflake.client.api.datasource.SnowflakeDataSourceFactory;
 import net.snowflake.ingest.SimpleIngestManager;
 import net.snowflake.ingest.connection.HistoryResponse;
 import org.apache.beam.sdk.coders.Coder;
@@ -1405,6 +1406,9 @@ public abstract static class DataSourceConfiguration implements Serializable {
     @Nullable
     public abstract ValueProvider<String> getPassword();
 
+    @Nullable
+    public abstract ValueProvider<String> getAccount();
+
     @Nullable
     public abstract PrivateKey getPrivateKey();
 
@@ -1457,6 +1461,8 @@ abstract static class Builder {
 
       abstract Builder setPassword(ValueProvider<String> password);
 
+      abstract Builder setAccount(ValueProvider<String> account);
+
       abstract Builder setPrivateKey(PrivateKey privateKey);
 
       abstract Builder setRawPrivateKey(ValueProvider<String> rawPrivateKey);
@@ -1802,7 +1808,7 @@ public DataSourceConfiguration withAuthenticator(String authenticator) {
     }
 
     /**
-     * Sets loginTimeout that will be used in {@link SnowflakeBasicDataSource#setLoginTimeout}.
+     * Sets loginTimeout that will be used in {@link SnowflakeDataSource#setLoginTimeout}.
      *
      * @param loginTimeout Integer with timeout value.
      */
@@ -1819,59 +1825,62 @@ void populateDisplayData(DisplayData.Builder builder) {
       }
     }
 
-    /** Builds {@link SnowflakeBasicDataSource} based on the current configuration. */
+    /** Builds {@link SnowflakeDataSource} based on the current configuration. */
     public DataSource buildDatasource() {
       if (getDataSource() == null) {
-        SnowflakeBasicDataSource basicDataSource = new SnowflakeBasicDataSource();
-        basicDataSource.setUrl(buildUrl());
+        SnowflakeDataSource dataSource = SnowflakeDataSourceFactory.createDataSource();
+        dataSource.setUrl(buildUrl());
 
         if (isNotEmpty(getOauthToken())) {
-          basicDataSource.setOauthToken(getOauthToken().get());
+          dataSource.setToken(getOauthToken().get());
         } else if (isNotEmpty(getUsername()) && getPrivateKey() != null) {
-          basicDataSource.setUser(getUsername().get());
-          basicDataSource.setPrivateKey(getPrivateKey());
+          dataSource.setUser(getUsername().get());
+          dataSource.setPrivateKey(getPrivateKey());
         } else if (isNotEmpty(getUsername()) && isNotEmpty(getRawPrivateKey())) {
           PrivateKey privateKey =
               KeyPairUtils.preparePrivateKey(
                   getRawPrivateKey().get(), getValueOrNull(getPrivateKeyPassphrase()));
-          basicDataSource.setPrivateKey(privateKey);
-          basicDataSource.setUser(getUsername().get());
+          dataSource.setPrivateKey(privateKey);
+          dataSource.setUser(getUsername().get());
         } else if (isNotEmpty(getUsername()) && isNotEmpty(getPassword())) {
-          basicDataSource.setUser(getUsername().get());
-          basicDataSource.setPassword(getPassword().get());
+          dataSource.setUser(getUsername().get());
+          dataSource.setPassword(getPassword().get());
         } else {
           throw new RuntimeException("Missing credentials values. Please check your credentials");
         }
 
+        if (isNotEmpty(getAccount())) {
+          dataSource.setAccount(getAccount().get());
+        }
         if (isNotEmpty(getDatabase())) {
-          basicDataSource.setDatabaseName(getDatabase().get());
+          dataSource.setDatabaseName(getDatabase().get());
         }
         if (isNotEmpty(getWarehouse())) {
-          basicDataSource.setWarehouse(getWarehouse().get());
+          dataSource.setWarehouse(getWarehouse().get());
         }
         if (isNotEmpty(getSchema())) {
-          basicDataSource.setSchema(getSchema().get());
+          dataSource.setSchema(getSchema().get());
         }
         if (isNotEmpty(getServerName())) {
-          basicDataSource.setServerName(getServerName().get());
+          dataSource.setServerName(getServerName().get());
         }
         if (getPortNumber() != null) {
-          basicDataSource.setPortNumber(getPortNumber());
+          dataSource.setPortNumber(getPortNumber());
         }
         if (isNotEmpty(getRole())) {
-          basicDataSource.setRole(getRole().get());
+          dataSource.setRole(getRole().get());
         }
         if (getAuthenticator() != null) {
-          basicDataSource.setAuthenticator(getAuthenticator());
+          dataSource.setAuthenticator(getAuthenticator());
         }
         if (getLoginTimeout() != null) {
           try {
-            basicDataSource.setLoginTimeout(getLoginTimeout());
+            dataSource.setLoginTimeout(getLoginTimeout());
           } catch (SQLException e) {
             throw new RuntimeException("Failed to setLoginTimeout");
           }
         }
-        return basicDataSource;
+        return dataSource;
       }
       return getDataSource();
     }

sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeBasicDataSource.java

@@ -36,10 +36,12 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.Executor;
-import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
+import net.snowflake.client.internal.api.implementation.datasource.SnowflakeBasicDataSource;
 
 /**
- * Fake implementation of {@link net.snowflake.client.jdbc.SnowflakeBasicDataSource} used in tests.
+ * Fake implementation of {@link
+ * net.snowflake.client.internal.api.implementation.datasource.SnowflakeBasicDataSource} used in
+ * tests.
  */
 public class FakeSnowflakeBasicDataSource extends SnowflakeBasicDataSource implements Serializable {
   @Override

sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeBatchServiceImpl.java

@@ -24,7 +24,7 @@
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
-import net.snowflake.client.jdbc.SnowflakeSQLException;
+import net.snowflake.client.api.exception.SnowflakeSQLException;
 import org.apache.beam.sdk.io.snowflake.data.SnowflakeTableSchema;
 import org.apache.beam.sdk.io.snowflake.enums.CreateDisposition;
 import org.apache.beam.sdk.io.snowflake.enums.WriteDisposition;

sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeDatabase.java

@@ -23,7 +23,7 @@
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
-import net.snowflake.client.jdbc.SnowflakeSQLException;
+import net.snowflake.client.api.exception.SnowflakeSQLException;
 
 /** Fake implementation of Snowflake warehouse used in test code. */
 public class FakeSnowflakeDatabase implements Serializable {

sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/DataSourceConfigurationTest.java

@@ -27,7 +27,7 @@
 import java.util.Arrays;
 import java.util.stream.Collectors;
 import javax.sql.DataSource;
-import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
+import net.snowflake.client.internal.api.implementation.datasource.SnowflakeBasicDataSource;
 import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
 import org.apache.beam.sdk.io.snowflake.test.TestUtils;
 import org.junit.Before;

sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/SchemaDispositionTest.java

@@ -23,7 +23,7 @@
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.LongStream;
-import net.snowflake.client.jdbc.SnowflakeSQLException;
+import net.snowflake.client.api.exception.SnowflakeSQLException;
 import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
 import org.apache.beam.sdk.io.snowflake.data.SnowflakeColumn;
 import org.apache.beam.sdk.io.snowflake.data.SnowflakeTableSchema;

sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/SnowflakeIOWriteTest.java

@@ -24,7 +24,7 @@
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.LongStream;
-import net.snowflake.client.jdbc.SnowflakeSQLException;
+import net.snowflake.client.api.exception.SnowflakeSQLException;
 import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
 import org.apache.beam.sdk.io.snowflake.services.SnowflakeServices;
 import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBasicDataSource;

sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/StreamingWriteTest.java

@@ -30,7 +30,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 import java.util.stream.LongStream;
-import net.snowflake.client.jdbc.SnowflakeSQLException;
+import net.snowflake.client.api.exception.SnowflakeSQLException;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
 import org.apache.beam.sdk.io.snowflake.services.SnowflakeServices;