Fix #21096: Support boolean in extra__snowflake__insecure_mode (#21155)

apache · Jan 28, 2022 · 534e9ae · 534e9ae
1 parent 14cfefa
commit 534e9ae
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 1 deletion.
diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py
@@ -32,6 +32,12 @@
 from airflow.utils.strings import to_boolean
 
 
+def _try_to_boolean(value: Any):
+    if isinstance(value, (str, type(None))):
+        return to_boolean(value)
+    return value
+
+
 class SnowflakeHook(DbApiHook):
     """
     A client to interact with Snowflake.
@@ -157,7 +163,7 @@ def _get_conn_params(self) -> Dict[str, Optional[str]]:
         schema = conn.schema or ''
         authenticator = conn.extra_dejson.get('authenticator', 'snowflake')
         session_parameters = conn.extra_dejson.get('session_parameters')
-        insecure_mode = to_boolean(
+        insecure_mode = _try_to_boolean(
             conn.extra_dejson.get(
                 'extra__snowflake__insecure_mode', conn.extra_dejson.get('insecure_mode', None)
             )

diff --git a/tests/providers/snowflake/hooks/test_snowflake.py b/tests/providers/snowflake/hooks/test_snowflake.py
@@ -233,6 +233,32 @@ class TestPytestSnowflakeHook:
                     'warehouse': 'af_wh',
                 },
             ),
+            (
+                {
+                    **BASE_CONNECTION_KWARGS,
+                    'extra': {
+                        **BASE_CONNECTION_KWARGS['extra'],
+                        'extra__snowflake__insecure_mode': False,
+                    },
+                },
+                (
+                    'snowflake://user:pw@airflow.af_region/db/public?'
+                    'application=AIRFLOW&authenticator=snowflake&role=af_role&warehouse=af_wh'
+                ),
+                {
+                    'account': 'airflow',
+                    'application': 'AIRFLOW',
+                    'authenticator': 'snowflake',
+                    'database': 'db',
+                    'password': 'pw',
+                    'region': 'af_region',
+                    'role': 'af_role',
+                    'schema': 'public',
+                    'session_parameters': None,
+                    'user': 'user',
+                    'warehouse': 'af_wh',
+                },
+            ),
         ],
     )
     def test_hook_should_support_prepare_basic_conn_params_and_uri(