accept either key file path or file itself

Author: vvkh

data_diff/databases/snowflake.py

@@ -162,7 +162,7 @@ class Snowflake(Database):
 
     _conn: Any
 
-    def __init__(self, *, schema: str, **kw):
+    def __init__(self, *, schema: str, key_file_path: str | None = None, key: bytes | None = None, **kw):
         super().__init__()
         snowflake, serialization, default_backend = import_snowflake()
         logging.getLogger("snowflake.connector").setLevel(logging.WARNING)
@@ -172,20 +172,26 @@ def __init__(self, *, schema: str, **kw):
         logging.getLogger("snowflake.connector.network").disabled = True
 
         assert '"' not in schema, "Schema name should not contain quotes!"
+        if key and key_file_path:
+            raise ConnectError("Only key value or key file path can be specified, not both")
+
+        if key_file_path:
+            with open(key_file_path, "rb") as f:
+                key = f.read()
+
         # If a private key is used, read it from the specified path and pass it as "private_key" to the connector.
-        if "key" in kw:
-            with open(kw.get("key"), "rb") as key:
-                if "password" in kw:
-                    raise ConnectError("Cannot use password and key at the same time")
-                if kw.get("private_key_passphrase"):
-                    encoded_passphrase = kw.get("private_key_passphrase").encode()
-                else:
-                    encoded_passphrase = None
-                p_key = serialization.load_pem_private_key(
-                    key.read(),
-                    password=encoded_passphrase,
-                    backend=default_backend(),
-                )
+        if key:
+            if "password" in kw:
+                raise ConnectError("Cannot use password and key at the same time")
+            if kw.get("private_key_passphrase"):
+                encoded_passphrase = kw.get("private_key_passphrase").encode()
+            else:
+                encoded_passphrase = None
+            p_key = serialization.load_pem_private_key(
+                key,
+                password=encoded_passphrase,
+                backend=default_backend(),
+            )
 
             kw["private_key"] = p_key.private_bytes(
                 encoding=serialization.Encoding.DER,

data_diff/dbt_parser.py

@@ -355,7 +355,7 @@ def set_connection(self):
             if credentials.get("private_key_path") is not None:
                 if credentials.get("password") is not None:
                     raise DataDiffDbtSnowflakeSetConnectionError("Cannot use password and key at the same time")
-                conn_info["key"] = credentials.get("private_key_path")
+                conn_info["key_file_path"] = credentials.get("private_key_path")
                 conn_info["private_key_passphrase"] = credentials.get("private_key_passphrase")
             elif credentials.get("authenticator") is not None:
                 conn_info["authenticator"] = credentials.get("authenticator")