purbanow commented on a change in pull request #12618:
URL: https://github.com/apache/beam/pull/12618#discussion_r481957844
##########
File path:
sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeIO.java
##########
@@ -1302,123 +1381,249 @@ public void finishBundle() throws Exception {
}
}
- private static String getValueOrNull(ValueProvider<String> valueProvider) {
- return valueProvider != null ? valueProvider.get() : null;
- }
-
/**
* A POJO describing a {@link DataSource}, providing all properties allowing
to create a {@link
* DataSource}.
*/
@AutoValue
public abstract static class DataSourceConfiguration implements Serializable
{
+ @Nullable
+ public abstract String getUrl();
- public abstract @Nullable String getUrl();
+ @Nullable
+ public abstract ValueProvider<String> getUsername();
- public abstract @Nullable String getUsername();
+ @Nullable
+ public abstract ValueProvider<String> getPassword();
- public abstract @Nullable String getPassword();
+ @Nullable
+ public abstract PrivateKey getPrivateKey();
- public abstract @Nullable PrivateKey getPrivateKey();
+ @Nullable
+ public abstract String getPrivateKeyPath();
- public abstract @Nullable String getOauthToken();
+ @Nullable
+ public abstract ValueProvider<String> getRawPrivateKey();
- public abstract @Nullable String getDatabase();
+ @Nullable
+ public abstract ValueProvider<String> getPrivateKeyPassphrase();
- public abstract @Nullable String getWarehouse();
+ @Nullable
+ public abstract ValueProvider<String> getOauthToken();
- public abstract @Nullable String getSchema();
+ @Nullable
+ public abstract ValueProvider<String> getDatabase();
- public abstract @Nullable String getServerName();
+ @Nullable
+ public abstract ValueProvider<String> getWarehouse();
- public abstract @Nullable Integer getPortNumber();
+ @Nullable
+ public abstract ValueProvider<String> getSchema();
- public abstract @Nullable String getRole();
+ @Nullable
+ public abstract ValueProvider<String> getServerName();
- public abstract @Nullable Integer getLoginTimeout();
+ @Nullable
+ public abstract Integer getPortNumber();
- public abstract @Nullable Boolean getSsl();
+ @Nullable
+ public abstract ValueProvider<String> getRole();
- public abstract @Nullable Boolean getValidate();
+ @Nullable
+ public abstract String getAuthenticator();
- public abstract @Nullable DataSource getDataSource();
+ @Nullable
+ public abstract Integer getLoginTimeout();
+
+ @Nullable
+ public abstract Boolean getSsl();
+
+ @Nullable
+ public abstract DataSource getDataSource();
abstract Builder builder();
@AutoValue.Builder
abstract static class Builder {
abstract Builder setUrl(String url);
- abstract Builder setUsername(String username);
+ abstract Builder setUsername(ValueProvider<String> username);
- abstract Builder setPassword(String password);
+ abstract Builder setPassword(ValueProvider<String> password);
abstract Builder setPrivateKey(PrivateKey privateKey);
- abstract Builder setOauthToken(String oauthToken);
+ abstract Builder setPrivateKeyPath(String privateKeyPath);
+
+ abstract Builder setRawPrivateKey(ValueProvider<String> rawPrivateKey);
+
+ abstract Builder setPrivateKeyPassphrase(ValueProvider<String>
privateKeyPassphrase);
- abstract Builder setDatabase(String database);
+ abstract Builder setOauthToken(ValueProvider<String> oauthToken);
- abstract Builder setWarehouse(String warehouse);
+ abstract Builder setDatabase(ValueProvider<String> database);
- abstract Builder setSchema(String schema);
+ abstract Builder setWarehouse(ValueProvider<String> warehouse);
- abstract Builder setServerName(String serverName);
+ abstract Builder setSchema(ValueProvider<String> schema);
+
+ abstract Builder setServerName(ValueProvider<String> serverName);
abstract Builder setPortNumber(Integer portNumber);
- abstract Builder setRole(String role);
+ abstract Builder setRole(ValueProvider<String> role);
+
+ abstract Builder setAuthenticator(String authenticator);
abstract Builder setLoginTimeout(Integer loginTimeout);
abstract Builder setSsl(Boolean ssl);
- abstract Builder setValidate(Boolean validate);
-
abstract Builder setDataSource(DataSource dataSource);
abstract DataSourceConfiguration build();
}
+ public static DataSourceConfiguration create() {
+ return new
AutoValue_SnowflakeIO_DataSourceConfiguration.Builder().build();
+ }
+
/**
* Creates {@link DataSourceConfiguration} from existing instance of
{@link DataSource}.
*
- * @param dataSource an instance of {@link DataSource}.
+ * @param dataSource - an instance of {@link DataSource}.
*/
public static DataSourceConfiguration create(DataSource dataSource) {
checkArgument(dataSource instanceof Serializable, "dataSource must be
Serializable");
return new AutoValue_SnowflakeIO_DataSourceConfiguration.Builder()
- .setValidate(true)
.setDataSource(dataSource)
.build();
}
/**
- * Creates {@link DataSourceConfiguration} from instance of {@link
SnowflakeCredentials}.
+ * Sets username/password authentication.
*
- * @param credentials an instance of {@link SnowflakeCredentials}.
+ * @param username - Snowflake username.
+ * @param password - Password for provided Snowflake username.
*/
- public static DataSourceConfiguration create(SnowflakeCredentials
credentials) {
- if (credentials instanceof UsernamePasswordSnowflakeCredentials) {
- return new AutoValue_SnowflakeIO_DataSourceConfiguration.Builder()
- .setValidate(true)
- .setUsername(((UsernamePasswordSnowflakeCredentials)
credentials).getUsername())
- .setPassword(((UsernamePasswordSnowflakeCredentials)
credentials).getPassword())
- .build();
- } else if (credentials instanceof OAuthTokenSnowflakeCredentials) {
- return new AutoValue_SnowflakeIO_DataSourceConfiguration.Builder()
- .setValidate(true)
- .setOauthToken(((OAuthTokenSnowflakeCredentials)
credentials).getToken())
- .build();
- } else if (credentials instanceof KeyPairSnowflakeCredentials) {
- return new AutoValue_SnowflakeIO_DataSourceConfiguration.Builder()
- .setValidate(true)
- .setUsername(((KeyPairSnowflakeCredentials)
credentials).getUsername())
- .setPrivateKey(((KeyPairSnowflakeCredentials)
credentials).getPrivateKey())
- .build();
- }
- throw new IllegalArgumentException(
- "Can't create DataSourceConfiguration from given credentials");
+ public DataSourceConfiguration withUsernamePasswordAuth(String username,
String password) {
+ return builder()
+ .setUsername(ValueProvider.StaticValueProvider.of(username))
+ .setPassword(ValueProvider.StaticValueProvider.of(password))
+ .build();
+ }
+
+ /**
+ * Sets username/password authentication.
+ *
+ * @param username - Snowflake username.
+ * @param password - Password for provided Snowflake username.
+ */
+ public DataSourceConfiguration withUsernamePasswordAuth(
+ ValueProvider<String> username, ValueProvider<String> password) {
+ return builder().setUsername(username).setPassword(password).build();
+ }
+
+ /**
+ * Sets OAuth authentication.
+ *
+ * @param token - OAuth token.
+ */
+ public DataSourceConfiguration withOAuth(String token) {
+ return
builder().setOauthToken(ValueProvider.StaticValueProvider.of(token)).build();
+ }
+
+ /**
+ * Sets OAuth authentication.
+ *
+ * @param token - OAuth token.
+ */
+ public DataSourceConfiguration withOAuth(ValueProvider<String> token) {
+ return builder().setOauthToken(token).build();
+ }
+
+ /**
+ * Sets key pair authentication.
+ *
+ * @param username - Snowflake username.
+ * @param privateKey - Private key.
+ */
+ public DataSourceConfiguration withKeyPairAuth(String username, PrivateKey
privateKey) {
+ return builder()
+ .setUsername(ValueProvider.StaticValueProvider.of(username))
+ .setPrivateKey(privateKey)
+ .build();
+ }
+
+ /**
+ * Sets key pair authentication.
+ *
+ * @param username - Snowflake username.
+ * @param privateKeyPath - Private key path.
+ * @param privateKeyPassphrase - Passphrase for provided private key.
+ */
+ public DataSourceConfiguration withKeyPairPathAuth(
+ ValueProvider<String> username,
+ String privateKeyPath,
+ ValueProvider<String> privateKeyPassphrase) {
+ String privateKey = KeyPairUtils.readPrivateKeyFile(privateKeyPath);
Review comment:
Exactly :)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
