snowflake add support for private_key_passphrase

Author: dlawin

data_diff/dbt.py

@@ -480,8 +480,8 @@ def set_connection(self):
         credentials, conn_type = self._get_connection_creds()
 
         if conn_type == "snowflake":
-            if credentials.get("authenticator") is not None or credentials.get("private_key_passphrase") is not None:
-                raise Exception("Federated authentication and key/pair with passphrase is not yet supported for Snowflake.")
+            if credentials.get("authenticator") is not None:
+                raise Exception("Federated authentication is not yet supported for Snowflake.")
             conn_info = {
                 "driver": conn_type,
                 "user": credentials.get("user"),
@@ -498,6 +498,7 @@ def set_connection(self):
                 if credentials.get("password") is not None:
                     raise Exception("Cannot use password and key at the same time")
                 conn_info["key"] = credentials.get("private_key_path")
+                conn_info["private_key_passphrase"] = credentials.get("private_key_passphrase")
             elif credentials.get("password") is not None:
                 conn_info["password"] = credentials.get("password")
             else:

tests/test_dbt.py

@@ -166,6 +166,22 @@ def test_set_connection_snowflake_success_key(self):
         self.assertEqual(mock_self.connection.get("key"), expected_credentials["private_key_path"])
         self.assertEqual(mock_self.requires_upper, True)
 
+    def test_set_connection_snowflake_success_key_and_passphrase(self):
+        expected_driver = "snowflake"
+        expected_credentials = {"user": "user", "private_key_path": "private_key_path", "private_key_passphrase": "private_key_passphrase"}
+        mock_self = Mock()
+        mock_self._get_connection_creds.return_value = (expected_credentials, expected_driver)
+
+        DbtParser.set_connection(mock_self)
+
+        self.assertIsInstance(mock_self.connection, dict)
+        self.assertEqual(mock_self.connection.get("driver"), expected_driver)
+        self.assertEqual(mock_self.connection.get("user"), expected_credentials["user"])
+        self.assertEqual(mock_self.connection.get("password"), None)
+        self.assertEqual(mock_self.connection.get("key"), expected_credentials["private_key_path"])
+        self.assertEqual(mock_self.connection.get("private_key_passphrase"), expected_credentials["private_key_passphrase"])
+        self.assertEqual(mock_self.requires_upper, True)
+
     def test_set_connection_snowflake_no_key_or_password(self):
         expected_driver = "snowflake"
         expected_credentials = {"user": "user"}
@@ -199,17 +215,6 @@ def test_set_connection_snowflake_key_and_password(self):
 
         self.assertNotIsInstance(mock_self.connection, dict)
 
-    def test_set_connection_snowflake_private_key_passphrase(self):
-        expected_driver = "snowflake"
-        expected_credentials = {"user": "user", "private_key_passphrase": "private_key_passphrase"}
-        mock_self = Mock()
-        mock_self._get_connection_creds.return_value = (expected_credentials, expected_driver)
-
-        with self.assertRaises(Exception):
-            DbtParser.set_connection(mock_self)
-
-        self.assertNotIsInstance(mock_self.connection, dict)
-
     def test_set_connection_bigquery_success(self):
         expected_driver = "bigquery"
         expected_credentials = {