This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new c6f151fa8f Fix an issue that crashes Airflow Webserver when passed
invalid private key path to Snowflake (#32016)
c6f151fa8f is described below
commit c6f151fa8f569687f3d889bce04bc270f114d208
Author: Bartosz Jankiewicz <[email protected]>
AuthorDate: Sun Jul 9 00:32:13 2023 +0200
Fix an issue that crashes Airflow Webserver when passed invalid private key
path to Snowflake (#32016)
* Validate private key file path and size in Snowflake hook
* Test update for snowflake hook
* Code review fixes
* Fix import
Co-authored-by: xrmr <[email protected]>
* Fixes
* Fix error message.
---------
Co-authored-by: Tzu-ping Chung <[email protected]>
Co-authored-by: xrmr <[email protected]>
Co-authored-by: bjankiewicz <[email protected]>
---
airflow/providers/snowflake/hooks/snowflake.py | 7 ++++++-
tests/providers/snowflake/hooks/test_snowflake.py | 18 ++++++++++++++++++
2 files changed, 24 insertions(+), 1 deletion(-)
diff --git a/airflow/providers/snowflake/hooks/snowflake.py
b/airflow/providers/snowflake/hooks/snowflake.py
index 2a202441a7..59199cf8cd 100644
--- a/airflow/providers/snowflake/hooks/snowflake.py
+++ b/airflow/providers/snowflake/hooks/snowflake.py
@@ -248,7 +248,12 @@ class SnowflakeHook(DbApiHook):
"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 ValueError("The private_key_file path points to an empty
or invalid file.")
+ if private_key_file_path.stat().st_size > 4096:
+ 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()
diff --git a/tests/providers/snowflake/hooks/test_snowflake.py
b/tests/providers/snowflake/hooks/test_snowflake.py
index e1105fbb59..6a738952d9 100644
--- a/tests/providers/snowflake/hooks/test_snowflake.py
+++ b/tests/providers/snowflake/hooks/test_snowflake.py
@@ -393,6 +393,24 @@ class TestPytestSnowflakeHook:
), 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 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()
+
def test_should_add_partner_info(self):
with mock.patch.dict(
"os.environ",
