Fix SnowflakePrivateKeyPemProfileMapping private_key reference (#501)

Renamed key for Snowflake profile mapping from private_key_content to private_key as expected by DBT Tested locally with Airflow Snowflake Connection containing extra.private_key_content Closes #479
astronomer · Aug 31, 2023 · 8c9e768 · 8c9e768
1 parent 9387c45
commit 8c9e768
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git 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
@@ -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"),