Skip to content

Commit eaf0284

Browse files
turbturb
authored
Upgrade Snowflake JDBC client to v4 (#37885)
* Upgrade Snowflake driver to v4 * SnowflakeIO: add account parameter
1 parent c4638a3 commit eaf0284

9 files changed

Lines changed: 44 additions & 33 deletions

File tree

‎sdks/java/io/snowflake/build.gradle‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,16 +30,16 @@ dependencies {
3030
implementation project(path: ":sdks:java:extensions:google-cloud-platform-core")
3131
permitUnusedDeclared project(path: ":sdks:java:extensions:google-cloud-platform-core")
3232
implementation library.java.slf4j_api
33-
implementation group: 'net.snowflake', name: 'snowflake-jdbc', version: '3.20.0'
34-
implementation group: 'com.opencsv', name: 'opencsv', version: '5.0'
35-
implementation 'net.snowflake:snowflake-ingest-sdk:0.9.9'
33+
implementation group: 'net.snowflake', name: 'snowflake-jdbc', version: '4.0.2'
34+
implementation group: 'com.opencsv', name: 'opencsv', version: '5.12.0'
35+
implementation 'net.snowflake:snowflake-ingest-sdk:4.4.2'
3636
implementation "org.bouncycastle:bcprov-jdk15on:1.70"
3737
implementation library.java.joda_time
3838
testImplementation project(path: ":sdks:java:core", configuration: "shadowTest")
3939
testImplementation project(path: ":sdks:java:extensions:avro", configuration: "testRuntimeMigration")
4040
testImplementation project(path: ":sdks:java:io:common")
4141
testImplementation project(path: ":sdks:java:testing:test-utils")
42-
testImplementation 'com.google.cloud:google-cloud-storage:1.102.0'
42+
testImplementation 'com.google.cloud:google-cloud-storage:1.118.1'
4343
testImplementation library.java.avro
4444
testImplementation library.java.junit
4545
testImplementation library.java.slf4j_api

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

Lines changed: 30 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,8 @@
3636
import java.util.stream.Collectors;
3737
import javax.annotation.Nullable;
3838
import javax.sql.DataSource;
39-
import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
39+
import net.snowflake.client.api.datasource.SnowflakeDataSource;
40+
import net.snowflake.client.api.datasource.SnowflakeDataSourceFactory;
4041
import net.snowflake.ingest.SimpleIngestManager;
4142
import net.snowflake.ingest.connection.HistoryResponse;
4243
import org.apache.beam.sdk.coders.Coder;
@@ -1405,6 +1406,9 @@ public abstract static class DataSourceConfiguration implements Serializable {
14051406
@Nullable
14061407
public abstract ValueProvider<String> getPassword();
14071408

1409+
@Nullable
1410+
public abstract ValueProvider<String> getAccount();
1411+
14081412
@Nullable
14091413
public abstract PrivateKey getPrivateKey();
14101414

@@ -1457,6 +1461,8 @@ abstract static class Builder {
14571461

14581462
abstract Builder setPassword(ValueProvider<String> password);
14591463

1464+
abstract Builder setAccount(ValueProvider<String> account);
1465+
14601466
abstract Builder setPrivateKey(PrivateKey privateKey);
14611467

14621468
abstract Builder setRawPrivateKey(ValueProvider<String> rawPrivateKey);
@@ -1802,7 +1808,7 @@ public DataSourceConfiguration withAuthenticator(String authenticator) {
18021808
}
18031809

18041810
/**
1805-
* Sets loginTimeout that will be used in {@link SnowflakeBasicDataSource#setLoginTimeout}.
1811+
* Sets loginTimeout that will be used in {@link SnowflakeDataSource#setLoginTimeout}.
18061812
*
18071813
* @param loginTimeout Integer with timeout value.
18081814
*/
@@ -1819,59 +1825,62 @@ void populateDisplayData(DisplayData.Builder builder) {
18191825
}
18201826
}
18211827

1822-
/** Builds {@link SnowflakeBasicDataSource} based on the current configuration. */
1828+
/** Builds {@link SnowflakeDataSource} based on the current configuration. */
18231829
public DataSource buildDatasource() {
18241830
if (getDataSource() == null) {
1825-
SnowflakeBasicDataSource basicDataSource = new SnowflakeBasicDataSource();
1826-
basicDataSource.setUrl(buildUrl());
1831+
SnowflakeDataSource dataSource = SnowflakeDataSourceFactory.createDataSource();
1832+
dataSource.setUrl(buildUrl());
18271833

18281834
if (isNotEmpty(getOauthToken())) {
1829-
basicDataSource.setOauthToken(getOauthToken().get());
1835+
dataSource.setToken(getOauthToken().get());
18301836
} else if (isNotEmpty(getUsername()) && getPrivateKey() != null) {
1831-
basicDataSource.setUser(getUsername().get());
1832-
basicDataSource.setPrivateKey(getPrivateKey());
1837+
dataSource.setUser(getUsername().get());
1838+
dataSource.setPrivateKey(getPrivateKey());
18331839
} else if (isNotEmpty(getUsername()) && isNotEmpty(getRawPrivateKey())) {
18341840
PrivateKey privateKey =
18351841
KeyPairUtils.preparePrivateKey(
18361842
getRawPrivateKey().get(), getValueOrNull(getPrivateKeyPassphrase()));
1837-
basicDataSource.setPrivateKey(privateKey);
1838-
basicDataSource.setUser(getUsername().get());
1843+
dataSource.setPrivateKey(privateKey);
1844+
dataSource.setUser(getUsername().get());
18391845
} else if (isNotEmpty(getUsername()) && isNotEmpty(getPassword())) {
1840-
basicDataSource.setUser(getUsername().get());
1841-
basicDataSource.setPassword(getPassword().get());
1846+
dataSource.setUser(getUsername().get());
1847+
dataSource.setPassword(getPassword().get());
18421848
} else {
18431849
throw new RuntimeException("Missing credentials values. Please check your credentials");
18441850
}
18451851

1852+
if (isNotEmpty(getAccount())) {
1853+
dataSource.setAccount(getAccount().get());
1854+
}
18461855
if (isNotEmpty(getDatabase())) {
1847-
basicDataSource.setDatabaseName(getDatabase().get());
1856+
dataSource.setDatabaseName(getDatabase().get());
18481857
}
18491858
if (isNotEmpty(getWarehouse())) {
1850-
basicDataSource.setWarehouse(getWarehouse().get());
1859+
dataSource.setWarehouse(getWarehouse().get());
18511860
}
18521861
if (isNotEmpty(getSchema())) {
1853-
basicDataSource.setSchema(getSchema().get());
1862+
dataSource.setSchema(getSchema().get());
18541863
}
18551864
if (isNotEmpty(getServerName())) {
1856-
basicDataSource.setServerName(getServerName().get());
1865+
dataSource.setServerName(getServerName().get());
18571866
}
18581867
if (getPortNumber() != null) {
1859-
basicDataSource.setPortNumber(getPortNumber());
1868+
dataSource.setPortNumber(getPortNumber());
18601869
}
18611870
if (isNotEmpty(getRole())) {
1862-
basicDataSource.setRole(getRole().get());
1871+
dataSource.setRole(getRole().get());
18631872
}
18641873
if (getAuthenticator() != null) {
1865-
basicDataSource.setAuthenticator(getAuthenticator());
1874+
dataSource.setAuthenticator(getAuthenticator());
18661875
}
18671876
if (getLoginTimeout() != null) {
18681877
try {
1869-
basicDataSource.setLoginTimeout(getLoginTimeout());
1878+
dataSource.setLoginTimeout(getLoginTimeout());
18701879
} catch (SQLException e) {
18711880
throw new RuntimeException("Failed to setLoginTimeout");
18721881
}
18731882
}
1874-
return basicDataSource;
1883+
return dataSource;
18751884
}
18761885
return getDataSource();
18771886
}

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,12 @@
3636
import java.util.Map;
3737
import java.util.Properties;
3838
import java.util.concurrent.Executor;
39-
import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
39+
import net.snowflake.client.internal.api.implementation.datasource.SnowflakeBasicDataSource;
4040

4141
/**
42-
* Fake implementation of {@link net.snowflake.client.jdbc.SnowflakeBasicDataSource} used in tests.
42+
* Fake implementation of {@link
43+
* net.snowflake.client.internal.api.implementation.datasource.SnowflakeBasicDataSource} used in
44+
* tests.
4345
*/
4446
public class FakeSnowflakeBasicDataSource extends SnowflakeBasicDataSource implements Serializable {
4547
@Override

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import java.sql.SQLException;
2525
import java.util.ArrayList;
2626
import java.util.List;
27-
import net.snowflake.client.jdbc.SnowflakeSQLException;
27+
import net.snowflake.client.api.exception.SnowflakeSQLException;
2828
import org.apache.beam.sdk.io.snowflake.data.SnowflakeTableSchema;
2929
import org.apache.beam.sdk.io.snowflake.enums.CreateDisposition;
3030
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‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
import java.util.Map;
2424
import java.util.concurrent.ConcurrentHashMap;
2525
import java.util.stream.Collectors;
26-
import net.snowflake.client.jdbc.SnowflakeSQLException;
26+
import net.snowflake.client.api.exception.SnowflakeSQLException;
2727

2828
/** Fake implementation of Snowflake warehouse used in test code. */
2929
public class FakeSnowflakeDatabase implements Serializable {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
import java.util.Arrays;
2828
import java.util.stream.Collectors;
2929
import javax.sql.DataSource;
30-
import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
30+
import net.snowflake.client.internal.api.implementation.datasource.SnowflakeBasicDataSource;
3131
import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
3232
import org.apache.beam.sdk.io.snowflake.test.TestUtils;
3333
import org.junit.Before;

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
import java.util.List;
2424
import java.util.stream.Collectors;
2525
import java.util.stream.LongStream;
26-
import net.snowflake.client.jdbc.SnowflakeSQLException;
26+
import net.snowflake.client.api.exception.SnowflakeSQLException;
2727
import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
2828
import org.apache.beam.sdk.io.snowflake.data.SnowflakeColumn;
2929
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‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import java.util.List;
2525
import java.util.stream.Collectors;
2626
import java.util.stream.LongStream;
27-
import net.snowflake.client.jdbc.SnowflakeSQLException;
27+
import net.snowflake.client.api.exception.SnowflakeSQLException;
2828
import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
2929
import org.apache.beam.sdk.io.snowflake.services.SnowflakeServices;
3030
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‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
import java.util.Map;
3131
import java.util.stream.Collectors;
3232
import java.util.stream.LongStream;
33-
import net.snowflake.client.jdbc.SnowflakeSQLException;
33+
import net.snowflake.client.api.exception.SnowflakeSQLException;
3434
import org.apache.beam.sdk.coders.StringUtf8Coder;
3535
import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
3636
import org.apache.beam.sdk.io.snowflake.services.SnowflakeServices;

0 commit comments

Comments
 (0)