diff --git a/src/preset_cli/cli/superset/sync/dbt/lib.py b/src/preset_cli/cli/superset/sync/dbt/lib.py index 24e1d243..a6b8f3f4 100644 --- a/src/preset_cli/cli/superset/sync/dbt/lib.py +++ b/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: """ diff --git a/tests/cli/superset/sync/dbt/lib_test.py b/tests/cli/superset/sync/dbt/lib_test.py index 127cf188..147d0a79 100644 --- a/tests/cli/superset/sync/dbt/lib_test.py +++ b/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.