Snowflake destination: add support for ecrypted private key

airbytehq · Jul 11, 2022 · 97d13d8 · 97d13d8
1 parent 167dc28
commit 97d13d8
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 22 deletions.
diff --git a/airbyte-integrations/bases/base-normalization/normalization/transform_config/transform.py b/airbyte-integrations/bases/base-normalization/normalization/transform_config/transform.py
@@ -224,6 +224,8 @@ def transform_snowflake(config: Dict[str, Any]):
             f.write(credentials["private_key"])
             f.close()
             dbt_config["private_key_path"] = "private_key_path.txt"
+            if credentials.get("passphrase"):
+                dbt_config["private_key_passphrase"] = credentials["passphrase"]
         elif credentials.get("password"):
             dbt_config["password"] = credentials["password"]
         else:

diff --git a/...wflake/src/main/java/io/airbyte/integrations/destination/snowflake/SnowflakeDatabase.java b/...wflake/src/main/java/io/airbyte/integrations/destination/snowflake/SnowflakeDatabase.java
@@ -51,6 +51,7 @@ public class SnowflakeDatabase {
       .build();
   public static final String PRIVATE_KEY_FILE_NAME = "rsa_key.p8";
   public static final String PRIVATE_KEY_FIELD_NAME = "private_key";
+  public static final String PASSPHRASE = "passphrase";
 
   public static HikariDataSource createDataSource(final JsonNode config) {
     final HikariDataSource dataSource = new HikariDataSource();
@@ -101,6 +102,9 @@ public static HikariDataSource createDataSource(final JsonNode config) {
       String privateKeyValue = credentials.get(PRIVATE_KEY_FIELD_NAME).asText();
       createPrivateKeyFile(PRIVATE_KEY_FILE_NAME, privateKeyValue);
       properties.put("private_key_file", PRIVATE_KEY_FILE_NAME);
+      if(credentials.has(PASSPHRASE)){
+        properties.put("private_key_file_pwd", credentials.get(PASSPHRASE).asText());
+      }
     } else {
       LOGGER.warn(
           "Obsolete User/password login mode is used. Please re-create a connection to use the latest connector's version");

diff --git a/.../airbyte/integrations/destination/snowflake/SnowflakeInsertDestinationAcceptanceTest.java b/.../airbyte/integrations/destination/snowflake/SnowflakeInsertDestinationAcceptanceTest.java
@@ -199,27 +199,6 @@ void testCheckWithKeyPairAuth() throws Exception {
     assertEquals(AirbyteConnectionStatus.Status.SUCCEEDED, check.getStatus());
   }
 
-  @ParameterizedTest
-  @ArgumentsSource(DataArgumentsProvider.class)
-  public void testSyncWithNormalization(final String messagesFilename, final String catalogFilename) throws Exception {
-    if (!normalizationFromSpec()) {
-      return;
-    }
-
-    final AirbyteCatalog catalog = Jsons.deserialize(MoreResources.readResource(catalogFilename), AirbyteCatalog.class);
-    final ConfiguredAirbyteCatalog configuredCatalog = CatalogHelpers.toDefaultConfiguredCatalog(catalog);
-    final List<AirbyteMessage> messages = MoreResources.readResource(messagesFilename).lines()
-        .map(record -> Jsons.deserialize(record, AirbyteMessage.class)).collect(Collectors.toList());
-
-    final JsonNode config = Jsons.deserialize(IOs.readFile(Path.of("secrets/config_key_pair.json")));
-    runSyncAndVerifyStateOutput(config, messages, configuredCatalog, true);
-
-    final String defaultSchema = getDefaultSchema(config);
-    final List<AirbyteRecordMessage> actualMessages = retrieveNormalizedRecords(catalog, defaultSchema);
-    assertSameMessages(messages, actualMessages, true);
-  }
-
-
   //normalization test
 
   /**

diff --git a/...integrations/destination/snowflake/SnowflakeInternalStagingDestinationAcceptanceTest.java b/...integrations/destination/snowflake/SnowflakeInternalStagingDestinationAcceptanceTest.java
@@ -8,7 +8,18 @@
 import com.google.common.base.Preconditions;
 import io.airbyte.commons.io.IOs;
 import io.airbyte.commons.json.Jsons;
+import io.airbyte.commons.resources.MoreResources;
+import io.airbyte.integrations.standardtest.destination.DataArgumentsProvider;
+import io.airbyte.protocol.models.AirbyteCatalog;
+import io.airbyte.protocol.models.AirbyteMessage;
+import io.airbyte.protocol.models.AirbyteRecordMessage;
+import io.airbyte.protocol.models.CatalogHelpers;
+import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
 import java.nio.file.Path;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ArgumentsSource;
 
 public class SnowflakeInternalStagingDestinationAcceptanceTest extends SnowflakeInsertDestinationAcceptanceTest {
 
@@ -18,5 +29,32 @@ public JsonNode getStaticConfig() {
     Preconditions.checkArgument(!SnowflakeDestinationResolver.isGcsCopy(internalStagingConfig));
     return internalStagingConfig;
   }
+  @ParameterizedTest
+  @ArgumentsSource(DataArgumentsProvider.class)
+  public void testSyncWithNormalization(final String messagesFilename, final String catalogFilename) throws Exception {
+    testSyncWithNormalizationWithKeyPairAuth(messagesFilename, catalogFilename, "secrets/config_key_pair.json");
+  }
+  @ParameterizedTest
+  @ArgumentsSource(DataArgumentsProvider.class)
+  public void testSyncWithNormalizationWithKeyPairEncrypt(final String messagesFilename, final String catalogFilename) throws Exception {
+    testSyncWithNormalizationWithKeyPairAuth(messagesFilename, catalogFilename, "secrets/config_key_pair_encrypted.json");
+  }
+  private void testSyncWithNormalizationWithKeyPairAuth(String messagesFilename, String catalogFilename, String configName) throws Exception {
+    if (!normalizationFromSpec()) {
+      return;
+    }
+
+    final AirbyteCatalog catalog = Jsons.deserialize(MoreResources.readResource(catalogFilename), AirbyteCatalog.class);
+    final ConfiguredAirbyteCatalog configuredCatalog = CatalogHelpers.toDefaultConfiguredCatalog(catalog);
+    final List<AirbyteMessage> messages = MoreResources.readResource(messagesFilename).lines()
+        .map(record -> Jsons.deserialize(record, AirbyteMessage.class)).collect(Collectors.toList());
+
+    final JsonNode config = Jsons.deserialize(IOs.readFile(Path.of(configName)));
+    runSyncAndVerifyStateOutput(config, messages, configuredCatalog, true);
+
+    final String defaultSchema = getDefaultSchema(config);
+    final List<AirbyteRecordMessage> actualMessages = retrieveNormalizedRecords(catalog, defaultSchema);
+    assertSameMessages(messages, actualMessages, true);
+  }
 
 }
diff --git a/docs/integrations/destinations/snowflake.md b/docs/integrations/destinations/snowflake.md
@@ -166,7 +166,9 @@ Field | Description |
 
 
 ### Key pair authentication
-    You need to generate private key - openssl genrsa 2048 | openssl pkcs8 -topk8 -inform PEM -out rsa_key.p8 -nocrypt
+    You need to generate private key without encryption - openssl genrsa 2048 | openssl pkcs8 -topk8 -inform PEM -out rsa_key.p8 -nocrypt
+
+    You can generate private key with encryption - openssl genrsa 2048 | openssl pkcs8 -topk8 -inform PEM -v1 PBE-SHA1-RC4-128 -out rsa_key.p8
 
     Based on private key we need to generate public key with - openssl rsa -in rsa_key.p8 -pubout -out rsa_key.pub
     Add generated public key to snowflake user - alter user <user_name> set rsa_public_key=<public_key_value>;