(airflow) branch main updated: enable client_store_temporary_credential for snowflake provider (#44431)

[sh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

shahar 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 e1f0f1f2cc5 enable client_store_temporary_credential for snowflake provider (#44431)
e1f0f1f2cc5 is described below

commit e1f0f1f2cc5a14599180575656a4a2bcd11d1a5b
Author: yzliao <yz...@users.noreply.github.com>
AuthorDate: Fri Nov 29 22:27:26 2024 -0800

    enable client_store_temporary_credential for snowflake provider (#44431)
    
    Co-authored-by: Yizheng Liao <yi...@ziphq.com>
---
 .../src/airflow/providers/snowflake/hooks/snowflake.py      | 13 +++++++++++--
 providers/tests/snowflake/hooks/test_snowflake.py           |  4 ++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/providers/src/airflow/providers/snowflake/hooks/snowflake.py b/providers/src/airflow/providers/snowflake/hooks/snowflake.py
index fdf75939bda..f9e40942bb3 100644
--- a/providers/src/airflow/providers/snowflake/hooks/snowflake.py
+++ b/providers/src/airflow/providers/snowflake/hooks/snowflake.py
@@ -134,6 +134,7 @@ class SnowflakeHook(DbApiHook):
                         "private_key_file": "private key",
                         "session_parameters": "session parameters",
                         "client_request_mfa_token": "client request mfa token",
+                        "client_store_temporary_credential": "client store temporary credential (externalbrowser mode)",
                     },
                     indent=1,
                 ),
@@ -162,6 +163,7 @@ class SnowflakeHook(DbApiHook):
         self.authenticator = kwargs.pop("authenticator", None)
         self.session_parameters = kwargs.pop("session_parameters", None)
         self.client_request_mfa_token = kwargs.pop("client_request_mfa_token", None)
+        self.client_store_temporary_credential = kwargs.pop("client_store_temporary_credential", None)
         self.query_ids: list[str] = []
 
     def _get_field(self, extra_dict, field_name):
@@ -206,6 +208,9 @@ class SnowflakeHook(DbApiHook):
         )
         schema = conn.schema or ""
         client_request_mfa_token = _try_to_boolean(self._get_field(extra_dict, "client_request_mfa_token"))
+        client_store_temporary_credential = _try_to_boolean(
+            self._get_field(extra_dict, "client_store_temporary_credential")
+        )
 
         # authenticator and session_parameters never supported long name so we don't use _get_field
         authenticator = extra_dict.get("authenticator", "snowflake")
@@ -234,6 +239,9 @@ class SnowflakeHook(DbApiHook):
         if client_request_mfa_token:
             conn_config["client_request_mfa_token"] = client_request_mfa_token
 
+        if client_store_temporary_credential:
+            conn_config["client_store_temporary_credential"] = client_store_temporary_credential
+
         # If private_key_file is specified in the extra json, load the contents of the file as a private key.
         # If private_key_content is specified in the extra json, use it as a private key.
         # As a next step, specify this private key in the connection configuration.
@@ -308,13 +316,14 @@ class SnowflakeHook(DbApiHook):
                 for k, v in conn_params.items()
                 if v
                 and k
-                not in {
+                not in [
                     "session_parameters",
                     "insecure_mode",
                     "private_key",
                     "client_request_mfa_token",
+                    "client_store_temporary_credential",
                     "json_result_force_utf8_decoding",
-                }
+                ]
             }
         )
 
diff --git a/providers/tests/snowflake/hooks/test_snowflake.py b/providers/tests/snowflake/hooks/test_snowflake.py
index f8696f81f24..775e9382729 100644
--- a/providers/tests/snowflake/hooks/test_snowflake.py
+++ b/providers/tests/snowflake/hooks/test_snowflake.py
@@ -140,6 +140,7 @@ class TestPytestSnowflakeHook:
                         "extra__snowflake__insecure_mode": "True",
                         "extra__snowflake__json_result_force_utf8_decoding": "True",
                         "extra__snowflake__client_request_mfa_token": "True",
+                        "extra__snowflake__client_store_temporary_credential": "True",
                     },
                 },
                 (
@@ -161,6 +162,7 @@ class TestPytestSnowflakeHook:
                     "insecure_mode": True,
                     "json_result_force_utf8_decoding": True,
                     "client_request_mfa_token": True,
+                    "client_store_temporary_credential": True,
                 },
             ),
             (
@@ -175,6 +177,7 @@ class TestPytestSnowflakeHook:
                         "extra__snowflake__insecure_mode": "False",
                         "extra__snowflake__json_result_force_utf8_decoding": "False",
                         "extra__snowflake__client_request_mfa_token": "False",
+                        "extra__snowflake__client_store_temporary_credential": "False",
                     },
                 },
                 (
@@ -252,6 +255,7 @@ class TestPytestSnowflakeHook:
                         "extra__snowflake__insecure_mode": False,
                         "extra__snowflake__json_result_force_utf8_decoding": True,
                         "extra__snowflake__client_request_mfa_token": False,
+                        "extra__snowflake__client_store_temporary_credential": False,
                     },
                 },
                 (