From 887ee8db0b1d618050f5095c2cdffb523aa7c5f6 Mon Sep 17 00:00:00 2001 From: Bartosz Jankiewicz Date: Tue, 20 Jun 2023 06:44:27 +0000 Subject: [PATCH 1/6] Validate private key file path and size in Snowflake hook --- airflow/providers/snowflake/hooks/snowflake.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py index fd5a2e968a56d..c38153be46b85 100644 --- a/airflow/providers/snowflake/hooks/snowflake.py +++ b/airflow/providers/snowflake/hooks/snowflake.py @@ -216,7 +216,16 @@ def _get_conn_params(self) -> dict[str, str | None]: "Please remove one." ) elif private_key_file: - private_key_pem = Path(private_key_file).read_bytes() + private_key_file_path = Path(private_key_file) + if not private_key_file_path.is_file() or private_key_file_path.stat().st_size == 0: + raise AirflowException( + "The private_key_file path points to an empty or invalid file." + ) + if private_key_file_path.stat().st_size > 4096: + raise AirflowException( + "File size of private_key_file is too big." + ) + private_key_pem = Path(private_key_file_path).read_bytes() elif private_key_content: private_key_pem = private_key_content.encode() From aa78f9d7f2fd2e9bb5821178447c8f627ced73fc Mon Sep 17 00:00:00 2001 From: Bartosz Jankiewicz Date: Tue, 20 Jun 2023 07:15:46 +0000 Subject: [PATCH 2/6] Test update for snowflake hook --- .../snowflake/hooks/test_snowflake.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/providers/snowflake/hooks/test_snowflake.py b/tests/providers/snowflake/hooks/test_snowflake.py index e731c283972ef..fb691b4ade6ef 100644 --- a/tests/providers/snowflake/hooks/test_snowflake.py +++ b/tests/providers/snowflake/hooks/test_snowflake.py @@ -31,6 +31,8 @@ from airflow.models import Connection from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook +from airflow import AirflowException + _PASSWORD = 'snowflake42' BASE_CONNECTION_KWARGS: dict = { @@ -342,6 +344,24 @@ def test_get_conn_params_should_support_private_auth_with_unencrypted_key( ), pytest.raises(TypeError, match="Password was given but private key is not encrypted."): SnowflakeHook(snowflake_conn_id='test_conn')._get_conn_params() + def test_get_conn_params_should_fail_on_invalid_key(self): + connection_kwargs = { + **BASE_CONNECTION_KWARGS, + 'password': None, + 'extra': { + 'database': 'db', + 'account': 'airflow', + 'warehouse': 'af_wh', + 'region': 'af_region', + 'role': 'af_role', + 'private_key_file': "/dev/urandom", + }, + } + with unittest.mock.patch.dict( + 'os.environ', AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri() + ), pytest.raises(AirflowException, match="The private_key_file path points to an empty or invalid file."): + SnowflakeHook(snowflake_conn_id='test_conn').get_conn() + def test_should_add_partner_info(self): with unittest.mock.patch.dict( 'os.environ', From 4504f4e44f65967143b3c74ff8a02e1d0a7052c0 Mon Sep 17 00:00:00 2001 From: Bartosz Jankiewicz Date: Tue, 20 Jun 2023 09:53:42 +0000 Subject: [PATCH 3/6] Code review fixes --- airflow/providers/snowflake/hooks/snowflake.py | 4 ++-- tests/providers/snowflake/hooks/test_snowflake.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py index c38153be46b85..e0251d5ffe8b1 100644 --- a/airflow/providers/snowflake/hooks/snowflake.py +++ b/airflow/providers/snowflake/hooks/snowflake.py @@ -218,11 +218,11 @@ def _get_conn_params(self) -> dict[str, str | None]: elif private_key_file: private_key_file_path = Path(private_key_file) if not private_key_file_path.is_file() or private_key_file_path.stat().st_size == 0: - raise AirflowException( + raise ValueError( "The private_key_file path points to an empty or invalid file." ) if private_key_file_path.stat().st_size > 4096: - raise AirflowException( + raise ValueError( "File size of private_key_file is too big." ) private_key_pem = Path(private_key_file_path).read_bytes() diff --git a/tests/providers/snowflake/hooks/test_snowflake.py b/tests/providers/snowflake/hooks/test_snowflake.py index fb691b4ade6ef..ecbcfda30afc8 100644 --- a/tests/providers/snowflake/hooks/test_snowflake.py +++ b/tests/providers/snowflake/hooks/test_snowflake.py @@ -359,7 +359,7 @@ def test_get_conn_params_should_fail_on_invalid_key(self): } with unittest.mock.patch.dict( 'os.environ', AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri() - ), pytest.raises(AirflowException, match="The private_key_file path points to an empty or invalid file."): + ), pytest.raises(ValueError, match="The private_key_file path points to an empty or invalid file."): SnowflakeHook(snowflake_conn_id='test_conn').get_conn() def test_should_add_partner_info(self): From 92bdf79d29635984d79a230d0e2f3fc0156b4b2c Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 6 Jul 2023 14:33:21 +0800 Subject: [PATCH 4/6] Fix import Co-authored-by: xrmr <17521278+xrmr@users.noreply.github.com> --- tests/providers/snowflake/hooks/test_snowflake.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/snowflake/hooks/test_snowflake.py b/tests/providers/snowflake/hooks/test_snowflake.py index afd56cfabf451..55976abaf5902 100644 --- a/tests/providers/snowflake/hooks/test_snowflake.py +++ b/tests/providers/snowflake/hooks/test_snowflake.py @@ -406,7 +406,7 @@ def test_get_conn_params_should_fail_on_invalid_key(self): 'private_key_file': "/dev/urandom", }, } - with unittest.mock.patch.dict( + with mock.patch.dict( 'os.environ', AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri() ), pytest.raises(ValueError, match="The private_key_file path points to an empty or invalid file."): SnowflakeHook(snowflake_conn_id='test_conn').get_conn() From faeba981942e32feb49950ff0803334fcae7870d Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Fri, 7 Jul 2023 13:50:45 +0200 Subject: [PATCH 5/6] Fixes --- .../providers/snowflake/hooks/snowflake.py | 8 ++------ .../snowflake/hooks/test_snowflake.py | 20 +++++++++---------- 2 files changed, 12 insertions(+), 16 deletions(-) diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py index 353975b5fe493..42eba8fe78305 100644 --- a/airflow/providers/snowflake/hooks/snowflake.py +++ b/airflow/providers/snowflake/hooks/snowflake.py @@ -250,13 +250,9 @@ def _get_conn_params(self) -> dict[str, str | None]: elif private_key_file: private_key_file_path = Path(private_key_file) if not private_key_file_path.is_file() or private_key_file_path.stat().st_size == 0: - raise ValueError( - "The private_key_file path points to an empty or invalid file." - ) + raise ValueError("The private_key_file path points to an empty or invalid file.") if private_key_file_path.stat().st_size > 4096: - raise ValueError( - "File size of private_key_file is too big." - ) + raise ValueError("File size of private_key_file is too big.") private_key_pem = Path(private_key_file_path).read_bytes() elif private_key_content: private_key_pem = private_key_content.encode() diff --git a/tests/providers/snowflake/hooks/test_snowflake.py b/tests/providers/snowflake/hooks/test_snowflake.py index 55976abaf5902..6a738952d90f8 100644 --- a/tests/providers/snowflake/hooks/test_snowflake.py +++ b/tests/providers/snowflake/hooks/test_snowflake.py @@ -396,20 +396,20 @@ def test_get_conn_params_should_support_private_auth_with_unencrypted_key( def test_get_conn_params_should_fail_on_invalid_key(self): connection_kwargs = { **BASE_CONNECTION_KWARGS, - 'password': None, - 'extra': { - 'database': 'db', - 'account': 'airflow', - 'warehouse': 'af_wh', - 'region': 'af_region', - 'role': 'af_role', - 'private_key_file': "/dev/urandom", + "password": None, + "extra": { + "database": "db", + "account": "airflow", + "warehouse": "af_wh", + "region": "af_region", + "role": "af_role", + "private_key_file": "/dev/urandom", }, } with mock.patch.dict( - 'os.environ', AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri() + "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri() ), pytest.raises(ValueError, match="The private_key_file path points to an empty or invalid file."): - SnowflakeHook(snowflake_conn_id='test_conn').get_conn() + SnowflakeHook(snowflake_conn_id="test_conn").get_conn() def test_should_add_partner_info(self): with mock.patch.dict( From 1c1cedeead7b07ba4497289020d9f925b38de5ef Mon Sep 17 00:00:00 2001 From: bjankiewicz Date: Fri, 7 Jul 2023 16:31:27 +0200 Subject: [PATCH 6/6] Fix error message. --- airflow/providers/snowflake/hooks/snowflake.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py index 42eba8fe78305..59199cf8cd2e2 100644 --- a/airflow/providers/snowflake/hooks/snowflake.py +++ b/airflow/providers/snowflake/hooks/snowflake.py @@ -252,7 +252,7 @@ def _get_conn_params(self) -> dict[str, str | None]: if not private_key_file_path.is_file() or private_key_file_path.stat().st_size == 0: raise ValueError("The private_key_file path points to an empty or invalid file.") if private_key_file_path.stat().st_size > 4096: - raise ValueError("File size of private_key_file is too big.") + raise ValueError("The private_key_file size is too big. Please keep it less than 4 KB.") private_key_pem = Path(private_key_file_path).read_bytes() elif private_key_content: private_key_pem = private_key_content.encode()