Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions sdks/java/io/snowflake/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -1404,6 +1405,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();

Expand Down Expand Up @@ -1456,6 +1460,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);
Expand Down Expand Up @@ -1801,7 +1807,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.
*/
Expand All @@ -1818,59 +1824,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();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading