diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/ConnectorArguments.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/ConnectorArguments.java index a58baa6de..55fdd4bb7 100644 --- a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/ConnectorArguments.java +++ b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/ConnectorArguments.java @@ -112,6 +112,10 @@ public class ConnectorArguments extends DefaultArguments { public static final String OPT_QUERY_LOG_EARLIEST_TIMESTAMP = "query-log-earliest-timestamp"; public static final String OPT_QUERY_LOG_ALTERNATES = "query-log-alternates"; + // Snowflake + public static final String OPT_PRIVATE_KEY_FILE = "private-key-file"; + public static final String OPT_PRIVATE_KEY_PASSWORD = "private-key-password"; + // Cloudera public static final String OPT_YARN_APPLICATION_TYPES = "yarn-application-types"; public static final String OPT_PAGINATION_PAGE_SIZE = "pagination-page-size"; @@ -159,7 +163,8 @@ public class ConnectorArguments extends DefaultArguments { public static final String OPT_RANGER_SCHEME_DEFAULT = "http"; public static final String OPT_RANGER_DISABLE_TLS_VALIDATION = "ranger-disable-tls-validation"; - // These are blocking threads on the client side, so it doesn't really matter much. + // These are blocking threads on the client side, so it doesn't really matter + // much. public static final Integer OPT_THREAD_POOL_SIZE_DEFAULT = 32; private final OptionSpec connectorNameOption = @@ -230,6 +235,18 @@ public class ConnectorArguments extends DefaultArguments { .withOptionalArg() .describedAs("sekr1t"); + private final OptionSpec optionPrivateKeyFile = + parser + .accepts(OPT_PRIVATE_KEY_FILE, "Path to the Private Key file used for authentication.") + .withRequiredArg() + .describedAs("/path/to/ras_key.p8"); + private final OptionSpec optionPrivateKeyPassword = + parser + .accepts( + OPT_PRIVATE_KEY_PASSWORD, "Private Key file password. Required if file is encrypted.") + .withRequiredArg() + .describedAs("sekr1t"); + private final OptionSpec optionStartDate = parser .accepts( @@ -342,9 +359,11 @@ public class ConnectorArguments extends DefaultArguments { .withValuesConvertedBy(ZonedParser.withDefaultPattern(DayOffset.END_OF_DAY)) .describedAs("2001-01-15[ 00:00:00.[000]]"); - // This is intentionally NOT provided as a default value to the optionQueryLogEnd OptionSpec, + // This is intentionally NOT provided as a default value to the + // optionQueryLogEnd OptionSpec, // because some callers - // such as ZonedIntervalIterable want to be able to distinguish a user-specified value from this + // such as ZonedIntervalIterable want to be able to distinguish a user-specified + // value from this // dumper-specified default. private final ZonedDateTime OPT_QUERY_LOG_END_DEFAULT = ZonedDateTime.now(ZoneOffset.UTC); @@ -800,7 +819,8 @@ public Predicate getDatabasePredicate() { } /** Returns the name of the single database specified, if exactly one database was specified. */ - // This can be used to generate an output filename, but it makes 1 be a special case + // This can be used to generate an output filename, but it makes 1 be a special + // case // that I find a little uncomfortable from the Unix philosophy: // "Sometimes the output filename is different" is hard to automate around. @CheckForNull @@ -873,6 +893,20 @@ public boolean isPasswordFlagProvided() { return has(optionPass); } + public boolean isPrivateKeyFileProvided() { + return has(optionPrivateKeyFile); + } + + @CheckForNull + public String getPrivateKeyFile() { + return getOptions().valueOf(optionPrivateKeyFile); + } + + @CheckForNull + public String getPrivateKeyPassword() { + return getOptions().valueOf(optionPrivateKeyPassword); + } + @CheckForNull public String getCluster() { return getOptions().valueOf(optionCluster); diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/annotations/RespectsArgumentPrivateKeyFile.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/annotations/RespectsArgumentPrivateKeyFile.java new file mode 100644 index 000000000..8334f4aec --- /dev/null +++ b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/annotations/RespectsArgumentPrivateKeyFile.java @@ -0,0 +1,36 @@ +/* + * Copyright 2022-2025 Google LLC + * Copyright 2013-2021 CompilerWorks + * + * 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. + */ +package com.google.edwmigration.dumper.application.dumper.annotations; + +import com.google.edwmigration.dumper.application.dumper.ConnectorArguments; +import java.lang.annotation.Documented; +import java.lang.annotation.ElementType; +import java.lang.annotation.Inherited; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; + +@Target(ElementType.TYPE) +@Retention(RetentionPolicy.RUNTIME) +@Documented +@Inherited +@RespectsInput( + order = 400, + arg = ConnectorArguments.OPT_PRIVATE_KEY_FILE, + description = "Path to the Private Key file used for authentication.", + required = "If the database user uses keypair authentication.") +public @interface RespectsArgumentPrivateKeyFile {} diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/annotations/RespectsArgumentPrivateKeyPassword.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/annotations/RespectsArgumentPrivateKeyPassword.java new file mode 100644 index 000000000..ff81cb075 --- /dev/null +++ b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/annotations/RespectsArgumentPrivateKeyPassword.java @@ -0,0 +1,36 @@ +/* + * Copyright 2022-2025 Google LLC + * Copyright 2013-2021 CompilerWorks + * + * 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. + */ +package com.google.edwmigration.dumper.application.dumper.annotations; + +import com.google.edwmigration.dumper.application.dumper.ConnectorArguments; +import java.lang.annotation.Documented; +import java.lang.annotation.ElementType; +import java.lang.annotation.Inherited; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; + +@Target(ElementType.TYPE) +@Retention(RetentionPolicy.RUNTIME) +@Documented +@Inherited +@RespectsInput( + order = 400, + arg = ConnectorArguments.OPT_PRIVATE_KEY_PASSWORD, + description = "Private Key file password.", + required = "If the private key file is encrypted.") +public @interface RespectsArgumentPrivateKeyPassword {} diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/AbstractSnowflakeConnector.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/AbstractSnowflakeConnector.java index ecb75cccd..a5e1906fb 100644 --- a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/AbstractSnowflakeConnector.java +++ b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/AbstractSnowflakeConnector.java @@ -26,6 +26,8 @@ import com.google.edwmigration.dumper.application.dumper.annotations.RespectsArgumentHostUnlessUrl; import com.google.edwmigration.dumper.application.dumper.annotations.RespectsArgumentJDBCUri; import com.google.edwmigration.dumper.application.dumper.annotations.RespectsArgumentPassword; +import com.google.edwmigration.dumper.application.dumper.annotations.RespectsArgumentPrivateKeyFile; +import com.google.edwmigration.dumper.application.dumper.annotations.RespectsArgumentPrivateKeyPassword; import com.google.edwmigration.dumper.application.dumper.annotations.RespectsArgumentUser; import com.google.edwmigration.dumper.application.dumper.annotations.RespectsInput; import com.google.edwmigration.dumper.application.dumper.annotations.RespectsInputs; @@ -39,6 +41,7 @@ import java.sql.SQLException; import java.util.ArrayList; import java.util.List; +import java.util.Properties; import javax.annotation.Nonnull; import javax.sql.DataSource; import org.apache.commons.lang3.StringUtils; @@ -49,6 +52,8 @@ @RespectsArgumentHostUnlessUrl @RespectsArgumentUser @RespectsArgumentPassword +@RespectsArgumentPrivateKeyFile +@RespectsArgumentPrivateKeyPassword @RespectsInputs({ // Although RespectsInput is @Repeatable, errorprone fails on it. @RespectsInput( @@ -80,39 +85,77 @@ public AbstractSnowflakeConnector(@Nonnull String name) { @Override public Handle open(@Nonnull ConnectorArguments arguments) throws MetadataDumperUsageException, SQLException { - String url = arguments.getUri(); - if (url == null) { - StringBuilder buf = new StringBuilder("jdbc:snowflake://"); - String host = arguments.getHost("host.snowflakecomputing.com"); - buf.append(host).append("/"); - // FWIW we can/should totally use a Properties object here and pass it to - // SimpleDriverDataSource rather than messing with the URL. - List optionalArguments = new ArrayList<>(); - if (arguments.getWarehouse() != null) { - optionalArguments.add("warehouse=" + arguments.getWarehouse()); - } - if (arguments.getRole() != null) { - optionalArguments.add("role=" + arguments.getRole()); - } - if (!optionalArguments.isEmpty()) { - buf.append("?").append(Joiner.on("&").join(optionalArguments)); - } - url = buf.toString(); - } + validateConnectionArguments(arguments); + String url = arguments.getUri() != null ? arguments.getUri() : getUrlFromArguments(arguments); String databaseName = arguments.getDatabases().isEmpty() ? DEFAULT_DATABASE : sanitizeDatabaseName(arguments.getDatabases().get(0)); - Driver driver = - newDriver(arguments.getDriverPaths(), "net.snowflake.client.jdbc.SnowflakeDriver"); - String password = arguments.getPasswordIfFlagProvided().orElse(null); - DataSource dataSource = new SimpleDriverDataSource(driver, url, arguments.getUser(), password); + DataSource dataSource = + arguments.isPrivateKeyFileProvided() + ? createPrivateKeyDataSource(arguments, url) + : createUserPasswordDataSource(arguments, url); JdbcHandle jdbcHandle = new JdbcHandle(dataSource); + setCurrentDatabase(databaseName, jdbcHandle.getJdbcTemplate()); return jdbcHandle; } + private void validateConnectionArguments(@Nonnull ConnectorArguments arguments) + throws MetadataDumperUsageException { + if (arguments.isPasswordFlagProvided() && arguments.isPrivateKeyFileProvided()) { + throw new MetadataDumperUsageException( + "Private key authentication method can't be used together with user password. " + + "If the private key file is encrypted, please use --" + + ConnectorArguments.OPT_PRIVATE_KEY_PASSWORD + + " to specify the key password."); + } + } + + private DataSource createUserPasswordDataSource(@Nonnull ConnectorArguments arguments, String url) + throws SQLException { + Driver driver = + newDriver(arguments.getDriverPaths(), "net.snowflake.client.jdbc.SnowflakeDriver"); + String password = arguments.getPasswordOrPrompt(); + return new SimpleDriverDataSource(driver, url, arguments.getUser(), password); + } + + private DataSource createPrivateKeyDataSource(@Nonnull ConnectorArguments arguments, String url) + throws SQLException { + Driver driver = + newDriver(arguments.getDriverPaths(), "net.snowflake.client.jdbc.SnowflakeDriver"); + Properties prop = new Properties(); + + prop.put("private_key_file", arguments.getPrivateKeyFile()); + prop.put("user", arguments.getUser()); + if (arguments.getPrivateKeyPassword() != null) { + prop.put("private_key_pwd", arguments.getPrivateKeyPassword()); + } + + return new SimpleDriverDataSource(driver, url, prop); + } + + @Nonnull + private String getUrlFromArguments(@Nonnull ConnectorArguments arguments) { + StringBuilder buf = new StringBuilder("jdbc:snowflake://"); + String host = arguments.getHost("host.snowflakecomputing.com"); + buf.append(host).append("/"); + // FWIW we can/should totally use a Properties object here and pass it to + // SimpleDriverDataSource rather than messing with the URL. + List optionalArguments = new ArrayList<>(); + if (arguments.getWarehouse() != null) { + optionalArguments.add("warehouse=" + arguments.getWarehouse()); + } + if (arguments.getRole() != null) { + optionalArguments.add("role=" + arguments.getRole()); + } + if (!optionalArguments.isEmpty()) { + buf.append("?").append(Joiner.on("&").join(optionalArguments)); + } + return buf.toString(); + } + final ImmutableList> getSqlTasks( @Nonnull SnowflakeInput inputSource, @Nonnull Class> header, diff --git a/dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/AbstractSnowflakeConnectorTest.java b/dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/AbstractSnowflakeConnectorTest.java index a9531adc9..1c7253a36 100644 --- a/dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/AbstractSnowflakeConnectorTest.java +++ b/dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/AbstractSnowflakeConnectorTest.java @@ -82,4 +82,27 @@ public void openConnection_failsForMalformedInput() throws IOException { Assert.assertTrue( e.getMessage().contains("Database name has incorrectly placed double quote(s).")); } + + @Test + public void openConnection_failsForMixedPrivateKeyAndPassword() throws IOException { + List args = new ArrayList<>(ARGS); + args.add("--connector"); + args.add(metadataConnector.getName()); + + args.add("--private-key-file"); + args.add("/path/to/file.r8"); + + ConnectorArguments arguments = + new ConnectorArguments(args.toArray(ArrayUtils.EMPTY_STRING_ARRAY)); + MetadataDumperUsageException e = + Assert.assertThrows( + MetadataDumperUsageException.class, + () -> { + metadataConnector.open(arguments); + }); + Assert.assertTrue( + e.getMessage() + .contains( + "Private key authentication method can't be used together with user password")); + } }