Merge pull request #108 from preset-io/snowflake_pk

feat: snowflake pk

src/preset_cli/cli/superset/sync/dbt/lib.py

@@ -135,7 +135,7 @@ def build_bigquery_sqlalchemy_params(target: Dict[str, Any]) -> Dict[str, Any]:
 
     with open(target["keyfile"], encoding="utf-8") as input_:
         credentials_info = json.load(input_)
-        parameters["masked_encrypted_extra"] = json.dumps(
+        parameters["encrypted_extra"] = json.dumps(
             {"credentials_info": credentials_info},
         )
 
@@ -147,12 +147,12 @@ def build_snowflake_sqlalchemy_params(target: Dict[str, Any]) -> Dict[str, Any]:
     Build the SQLAlchemy URI for a Snowflake target.
     """
     username = target["user"]
-    password = target["password"] or None
+    password = target.get("password", "") or None
     database = target["database"]
     host = target["account"]
     query = {"role": target["role"], "warehouse": target["warehouse"]}
 
-    return {
+    parameters = {
         "sqlalchemy_uri": str(
             URL(
                 drivername="snowflake",
@@ -165,6 +165,22 @@ def build_snowflake_sqlalchemy_params(target: Dict[str, Any]) -> Dict[str, Any]:
         ),
     }
 
+    if "private_key_path" in target:
+        with open(target["private_key_path"], encoding="utf-8") as input_:
+            pk_body = input_.read()
+
+        parameters["encrypted_extra"] = json.dumps(
+            {
+                "auth_method": "keypair",
+                "auth_params": {
+                    "privatekey_body": pk_body,
+                    "privatekey_pass": target.get("private_key_passphrase", ""),
+                },
+            },
+        )
+
+    return parameters
+
 
 def env_var(var: str, default: Optional[str] = None) -> str:
     """

tests/cli/superset/sync/dbt/lib_test.py

@@ -86,7 +86,7 @@ def test_build_sqlalchemy_params_bigquery(fs: FakeFilesystem) -> None:
     }
     assert build_sqlalchemy_params(config) == {
         "sqlalchemy_uri": "bigquery://my_project/",
-        "masked_encrypted_extra": json.dumps({"credentials_info": {"Hello": "World!"}}),
+        "encrypted_extra": json.dumps({"credentials_info": {"Hello": "World!"}}),
     }
 
 
@@ -108,7 +108,7 @@ def test_build_sqlalchemy_params_bigquery_with_priority(fs: FakeFilesystem) -> N
     }
     assert build_sqlalchemy_params(config) == {
         "sqlalchemy_uri": "bigquery://my_project/?priority=INTERACTIVE",
-        "masked_encrypted_extra": json.dumps({"credentials_info": {"Hello": "World!"}}),
+        "encrypted_extra": json.dumps({"credentials_info": {"Hello": "World!"}}),
     }
 
 
@@ -151,6 +151,40 @@ def test_build_snowflake_sqlalchemy_params() -> None:
     }
 
 
+def test_build_snowflake_sqlalchemy_params_pk(fs: FakeFilesystem) -> None:
+    """
+    Test ``build_snowflake_sqlalchemy_params`` for Snowflake with private keys.
+    """
+    fs.create_file("/path/to/key", contents="-----BEGIN ENCRYPTED PRIVATE KEY")
+
+    config = {
+        "type": "snowflake",
+        "account": "abc123.eu-west-1.aws",
+        "user": "jdoe",
+        "password": "secret",
+        "role": "admin",
+        "database": "default",
+        "warehouse": "dunder-mifflin",
+        "private_key_path": "/path/to/key",
+        "private_key_passphrase": "XXX",
+    }
+    assert build_sqlalchemy_params(config) == {
+        "sqlalchemy_uri": (
+            "snowflake://jdoe:secret@abc123.eu-west-1.aws/default?"
+            "role=admin&warehouse=dunder-mifflin"
+        ),
+        "encrypted_extra": json.dumps(
+            {
+                "auth_method": "keypair",
+                "auth_params": {
+                    "privatekey_body": "-----BEGIN ENCRYPTED PRIVATE KEY",
+                    "privatekey_pass": "XXX",
+                },
+            }
+        ),
+    }
+
+
 def test_build_sqlalchemy_params_unsupported() -> None:
     """
     Test ``build_sqlalchemy_params`` for databases currently unsupported.