diff --git a/cosmos/profiles/snowflake/user_privatekey.py b/cosmos/profiles/snowflake/user_privatekey.py index c3827fee4..3aa0454ec 100644 --- a/cosmos/profiles/snowflake/user_privatekey.py +++ b/cosmos/profiles/snowflake/user_privatekey.py @@ -27,10 +27,10 @@ class SnowflakePrivateKeyPemProfileMapping(BaseProfileMapping): "database", "warehouse", "schema", - "private_key_content", + "private_key", ] secret_fields = [ - "private_key_content", + "private_key", ] airflow_param_mapping = { "account": "extra.account", @@ -39,7 +39,7 @@ class SnowflakePrivateKeyPemProfileMapping(BaseProfileMapping): "warehouse": "extra.warehouse", "schema": "schema", "role": "extra.role", - "private_key_content": "extra.private_key_content", + "private_key": "extra.private_key_content", } @property @@ -68,7 +68,7 @@ def profile(self) -> dict[str, Any | None]: **self.mapped_params, **self.profile_args, # private_key should always get set as env var - "private_key_content": self.get_env_var_format("private_key_content"), + "private_key": self.get_env_var_format("private_key"), } # remove any null values diff --git a/tests/profiles/snowflake/test_snowflake_user_privatekey.py b/tests/profiles/snowflake/test_snowflake_user_privatekey.py index 1c399b4e9..2692792cf 100644 --- a/tests/profiles/snowflake/test_snowflake_user_privatekey.py +++ b/tests/profiles/snowflake/test_snowflake_user_privatekey.py @@ -133,7 +133,7 @@ def test_profile_args( assert profile_mapping.profile == { "type": mock_snowflake_conn.conn_type, "user": mock_snowflake_conn.login, - "private_key_content": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY_CONTENT') }}", + "private_key": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY') }}", "schema": mock_snowflake_conn.schema, "account": mock_snowflake_conn.extra_dejson.get("account"), "database": mock_snowflake_conn.extra_dejson.get("database"), @@ -158,7 +158,7 @@ def test_profile_args_overrides( assert profile_mapping.profile == { "type": mock_snowflake_conn.conn_type, "user": mock_snowflake_conn.login, - "private_key_content": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY_CONTENT') }}", + "private_key": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY') }}", "schema": mock_snowflake_conn.schema, "account": mock_snowflake_conn.extra_dejson.get("account"), "database": "my_db_override", @@ -176,7 +176,7 @@ def test_profile_env_vars( mock_snowflake_conn.conn_id, ) assert profile_mapping.env_vars == { - "COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY_CONTENT": mock_snowflake_conn.extra_dejson.get("private_key_content"), + "COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY": mock_snowflake_conn.extra_dejson.get("private_key_content"), } @@ -204,7 +204,7 @@ def test_old_snowflake_format() -> None: assert profile_mapping.profile == { "type": conn.conn_type, "user": conn.login, - "private_key_content": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY_CONTENT') }}", + "private_key": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY') }}", "schema": conn.schema, "account": conn.extra_dejson.get("account"), "database": conn.extra_dejson.get("database"), @@ -237,7 +237,7 @@ def test_appends_region() -> None: assert profile_mapping.profile == { "type": conn.conn_type, "user": conn.login, - "private_key_content": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY_CONTENT') }}", + "private_key": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY') }}", "schema": conn.schema, "account": f"{conn.extra_dejson.get('account')}.{conn.extra_dejson.get('region')}", "database": conn.extra_dejson.get("database"),