Merge remote-tracking branch 'origin/main' into feature/match_by_column_name

Author: chikamura

.github/workflows/main.yml

@@ -10,10 +10,20 @@ on:
     branches:
       - 'main'
     types: [opened, synchronize]
+  pull_request_target:
+    branches:
+      - 'main'
+    types: [labeled]
 
 jobs:
   main:
     runs-on: ubuntu-latest
+    if: >
+      ${{
+        github.event_name == 'pull_request' ||
+        (github.event_name == 'pull_request_target' && contains(github.event.pull_request.labels.*.name, 'safe to test')) ||
+        startsWith(github.ref, 'refs/tags/')
+      }}
     permissions:
       packages: write
       contents: read

README.md

@@ -13,11 +13,13 @@ Snowflake output plugin for Embulk loads records to Snowflake.
 - **host**: database host name (string, required)
 - **user**: database login user name (string, required)
 - **password**: database login password (string, default: "")
-- **privateKey**: database login using key-pair authentication(string, default: ""). This authentication method requires a 2048-bit (minimum) RSA key pair. 
+- **privateKey**: database login using key-pair authentication(string, default: ""). This authentication method requires a 2048-bit (minimum) RSA key pair.
+- **private_key_passphrase**: passphrase for private_key (string, default: "")
 - **warehouse**: destination warehouse name (string, required)
 - **database**: destination database name (string, required)
 - **schema**: destination schema name (string, default: "public")
 - **table**: destination table name (string, required)
+- **role**: role to execute queries (string, default: "")
 - **retry_limit**: max retry count for database operations (integer, default: 12). When intermediate table to create already created by another process, this plugin will retry with another table name to avoid collision.
 - **retry_wait**: initial retry wait time in milliseconds (integer, default: 1000 (1 second))
 - **max_retry_wait**: upper limit of retry wait, which will be doubled at every retry (integer, default: 1800000 (30 minutes))
@@ -60,10 +62,6 @@ Snowflake output plugin for Embulk loads records to Snowflake.
 
 ## Build
 
-## Not implement
-- Passphrase for `privateKey` in key-pair authentication.
-
-
 ```
 $ ./gradlew gem  # -t to watch change of files and rebuild continuously
 ```

src/main/java/org/embulk/output/SnowflakeOutputPlugin.java

@@ -7,6 +7,8 @@
 import java.sql.Types;
 import java.util.*;
 import java.util.function.BiFunction;
+import net.snowflake.client.jdbc.internal.org.bouncycastle.operator.OperatorCreationException;
+import net.snowflake.client.jdbc.internal.org.bouncycastle.pkcs.PKCSException;
 import org.embulk.config.ConfigDiff;
 import org.embulk.config.ConfigException;
 import org.embulk.config.TaskSource;
@@ -47,6 +49,10 @@ public interface SnowflakePluginTask extends PluginTask {
     @ConfigDefault("\"\"")
     String getPrivateKey();
 
+    @Config("private_key_passphrase")
+    @ConfigDefault("\"\"")
+    String getPrivateKeyPassphrase();
+
     @Config("database")
     public String getDatabase();
 
@@ -57,6 +63,10 @@ public interface SnowflakePluginTask extends PluginTask {
     @ConfigDefault("\"public\"")
     public String getSchema();
 
+    @Config("role")
+    @ConfigDefault("\"\"")
+    public String getRole();
+
     @Config("delete_stage")
     @ConfigDefault("false")
     public boolean getDeleteStage();
@@ -146,17 +156,21 @@ protected JdbcOutputConnector getConnector(PluginTask task, boolean retryableMet
       props.setProperty("password", t.getPassword());
     } else if (!t.getPrivateKey().isEmpty()) {
       try {
-        props.put("privateKey", PrivateKeyReader.get(t.getPrivateKey()));
-      } catch (IOException e) {
-        // Because the source of newConnection definition does not assume IOException, change it to
-        // ConfigException.
+        props.put(
+            "privateKey", PrivateKeyReader.get(t.getPrivateKey(), t.getPrivateKeyPassphrase()));
+      } catch (IOException | OperatorCreationException | PKCSException e) {
+        // Since this method is not allowed to throw any checked exception,
+        // wrap it with ConfigException, which is unchecked.
         throw new ConfigException(e);
       }
     }
 
     props.setProperty("warehouse", t.getWarehouse());
     props.setProperty("db", t.getDatabase());
     props.setProperty("schema", t.getSchema());
+    if (!t.getRole().isEmpty()) {
+      props.setProperty("role", t.getRole());
+    }
 
     // When CLIENT_METADATA_REQUEST_USE_CONNECTION_CTX is false (default),
     // getMetaData().getColumns() returns columns of the tables which table name is

src/main/java/org/embulk/output/snowflake/PrivateKeyReader.java

@@ -8,18 +8,31 @@
 import net.snowflake.client.jdbc.internal.org.bouncycastle.jce.provider.BouncyCastleProvider;
 import net.snowflake.client.jdbc.internal.org.bouncycastle.openssl.PEMParser;
 import net.snowflake.client.jdbc.internal.org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
+import net.snowflake.client.jdbc.internal.org.bouncycastle.openssl.jcajce.JceOpenSSLPKCS8DecryptorProviderBuilder;
+import net.snowflake.client.jdbc.internal.org.bouncycastle.operator.InputDecryptorProvider;
+import net.snowflake.client.jdbc.internal.org.bouncycastle.operator.OperatorCreationException;
+import net.snowflake.client.jdbc.internal.org.bouncycastle.pkcs.PKCS8EncryptedPrivateKeyInfo;
+import net.snowflake.client.jdbc.internal.org.bouncycastle.pkcs.PKCSException;
 
 // ref:
 // https://docs.snowflake.com/en/developer-guide/jdbc/jdbc-configure#privatekey-property-in-connection-properties
 public class PrivateKeyReader {
-  public static PrivateKey get(String pemString) throws IOException {
+  public static PrivateKey get(String pemString, String passphrase)
+      throws IOException, OperatorCreationException, PKCSException {
     Security.addProvider(new BouncyCastleProvider());
     PEMParser pemParser = new PEMParser(new StringReader(pemString));
     Object pemObject = pemParser.readObject();
     pemParser.close();
 
     PrivateKeyInfo privateKeyInfo;
-    if (pemObject instanceof PrivateKeyInfo) {
+    if (pemObject instanceof PKCS8EncryptedPrivateKeyInfo) {
+      // Handle the case where the private key is encrypted.
+      PKCS8EncryptedPrivateKeyInfo encryptedPrivateKeyInfo =
+          (PKCS8EncryptedPrivateKeyInfo) pemObject;
+      InputDecryptorProvider pkcs8Prov =
+          new JceOpenSSLPKCS8DecryptorProviderBuilder().build(passphrase.toCharArray());
+      privateKeyInfo = encryptedPrivateKeyInfo.decryptPrivateKeyInfo(pkcs8Prov);
+    } else if (pemObject instanceof PrivateKeyInfo) {
       privateKeyInfo = (PrivateKeyInfo) pemObject;
     } else {
       throw new IllegalArgumentException("Provided PEM does not contain a valid Private Key");

src/test/java/org/embulk/output/snowflake/TestSnowflakeOutputPlugin.java

@@ -208,6 +208,7 @@ public void testConfigDefault() throws Exception {
     assertEquals("", task.getUser());
     assertEquals("", task.getPassword());
     assertEquals("public", task.getSchema());
+    assertEquals("", task.getRole());
     assertEquals(false, task.getDeleteStage());
   }