From 328efaf81e729f8fb9879c0509dfba5270df8099 Mon Sep 17 00:00:00 2001 From: Aseem Bansal Date: Fri, 31 Mar 2023 20:03:19 +0530 Subject: [PATCH 1/2] feat(snowflake): better error message on key pair authentication --- .../datahub/ingestion/source_config/sql/snowflake.py | 8 ++++++++ metadata-ingestion/tests/unit/test_snowflake_source.py | 10 ++++++++++ 2 files changed, 18 insertions(+) diff --git a/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py b/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py index ba3aa50b62e74..bd2c8eb6eb8e0 100644 --- a/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py +++ b/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py @@ -179,6 +179,14 @@ def authenticator_type_is_valid(cls, v, values, field): f"unsupported authenticator type '{v}' was provided," f" use one of {list(VALID_AUTH_TYPES.keys())}" ) + if ( + values.get("private_key") is not None + or values.get("private_key_path") is not None + ) and v != "KEY_PAIR_AUTHENTICATOR": + raise ValueError( + f"Either `private_key` and `private_key_path` is set but `authentication_type` is {v}. " + f"Should be set to 'KEY_PAIR_AUTHENTICATOR' when using key pair authentication" + ) if v == "KEY_PAIR_AUTHENTICATOR": # If we are using key pair auth, we need the private key path and password to be set if ( diff --git a/metadata-ingestion/tests/unit/test_snowflake_source.py b/metadata-ingestion/tests/unit/test_snowflake_source.py index 94720faa8d0b8..6db344692bef1 100644 --- a/metadata-ingestion/tests/unit/test_snowflake_source.py +++ b/metadata-ingestion/tests/unit/test_snowflake_source.py @@ -238,6 +238,16 @@ def test_snowflake_config_with_no_connect_args_returns_base_connect_args(): } +def test_private_key_set_but_auth_not_changed(): + with pytest.raises(Exception): + SnowflakeV2Config.parse_obj( + { + "account_id": "acctname", + "private_key_path": "/a/random/path", + } + ) + + def test_snowflake_config_with_connect_args_overrides_base_connect_args(): config: SnowflakeV2Config = SnowflakeV2Config.parse_obj( { From 8286e055c5ea27ef08d19301d74c18e90984c558 Mon Sep 17 00:00:00 2001 From: Harshal Sheth Date: Tue, 4 Apr 2023 14:49:06 -0700 Subject: [PATCH 2/2] narrow exception type --- metadata-ingestion/tests/unit/test_snowflake_source.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/metadata-ingestion/tests/unit/test_snowflake_source.py b/metadata-ingestion/tests/unit/test_snowflake_source.py index 6db344692bef1..630768c80dfea 100644 --- a/metadata-ingestion/tests/unit/test_snowflake_source.py +++ b/metadata-ingestion/tests/unit/test_snowflake_source.py @@ -239,7 +239,7 @@ def test_snowflake_config_with_no_connect_args_returns_base_connect_args(): def test_private_key_set_but_auth_not_changed(): - with pytest.raises(Exception): + with pytest.raises(ValidationError): SnowflakeV2Config.parse_obj( { "account_id": "acctname",