Merge pull request #468 from dlawin/issue_442

add basic snowflake key/pair support

Author: dlawin

data_diff/dbt.py

@@ -480,12 +480,11 @@ def set_connection(self):
         credentials, conn_type = self._get_connection_creds()
 
         if conn_type == "snowflake":
-            if credentials.get("password") is None or credentials.get("private_key_path") is not None:
-                raise Exception("Only password authentication is currently supported for 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.")
             conn_info = {
                 "driver": conn_type,
                 "user": credentials.get("user"),
-                "password": credentials.get("password"),
                 "account": credentials.get("account"),
                 "database": credentials.get("database"),
                 "warehouse": credentials.get("warehouse"),
@@ -494,6 +493,15 @@ def set_connection(self):
             }
             self.threads = credentials.get("threads")
             self.requires_upper = True
+
+            if credentials.get("private_key_path") is not None:
+                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")
+            elif credentials.get("password") is not None:
+                conn_info["password"] = credentials.get("password")
+            else:
+                raise Exception("Password or key authentication not provided.")
         elif conn_type == "bigquery":
             method = credentials.get("method")
             # there are many connection types https://docs.getdbt.com/reference/warehouse-setups/bigquery-setup#oauth-via-gcloud

tests/test_dbt.py

@@ -136,7 +136,7 @@ def test_get_project_dict(self, mock_open):
         self.assertEqual(project_dict, expected_dict)
         mock_open.assert_called_once_with(Path(PROJECT_FILE))
 
-    def test_set_connection_snowflake_success(self):
+    def test_set_connection_snowflake_success_password(self):
         expected_driver = "snowflake"
         expected_credentials = {"user": "user", "password": "password"}
         mock_self = Mock()
@@ -148,9 +148,25 @@ def test_set_connection_snowflake_success(self):
         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"), expected_credentials["password"])
+        self.assertEqual(mock_self.connection.get("key"), None)
         self.assertEqual(mock_self.requires_upper, True)
 
-    def test_set_connection_snowflake_no_password(self):
+    def test_set_connection_snowflake_success_key(self):
+        expected_driver = "snowflake"
+        expected_credentials = {"user": "user", "private_key_path": "private_key_path"}
+        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.requires_upper, True)
+
+    def test_set_connection_snowflake_no_key_or_password(self):
         expected_driver = "snowflake"
         expected_credentials = {"user": "user"}
         mock_self = Mock()
@@ -160,6 +176,39 @@ def test_set_connection_snowflake_no_password(self):
             DbtParser.set_connection(mock_self)
 
         self.assertNotIsInstance(mock_self.connection, dict)
+    
+    def test_set_connection_snowflake_authenticator(self):
+        expected_driver = "snowflake"
+        expected_credentials = {"user": "user", "authenticator": "authenticator"}
+        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_snowflake_key_and_password(self):
+        expected_driver = "snowflake"
+        expected_credentials = {"user": "user", "private_key_path": "private_key_path", "password": "password"}
+        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_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"