[b/404818949] Add private key auth method to snowflake connectors

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<String> connectorNameOption =
@@ -230,6 +235,18 @@ public class ConnectorArguments extends DefaultArguments {
           .withOptionalArg()
           .describedAs("sekr1t");
 
+  private final OptionSpec<String> 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<String> optionPrivateKeyPassword =
+      parser
+          .accepts(
+              OPT_PRIVATE_KEY_PASSWORD, "Private Key file password. Required if file is encrypted.")
+          .withRequiredArg()
+          .describedAs("sekr1t");
+
   private final OptionSpec<ZonedDateTime> 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<String> 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);

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 {}

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 {}

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<String> 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<String> 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<Task<?>> getSqlTasks(
       @Nonnull SnowflakeInput inputSource,
       @Nonnull Class<? extends Enum<?>> header,

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<String> 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"));
+  }
 }