Add session_parameters option to snowflake_hook (#12071)

airflow/providers/snowflake/hooks/snowflake.py

@@ -46,6 +46,7 @@ def __init__(self, *args, **kwargs) -> None:
         self.role = kwargs.pop("role", None)
         self.schema = kwargs.pop("schema", None)
         self.authenticator = kwargs.pop("authenticator", None)
+        self.session_parameters = kwargs.pop("session_parameters", None)
 
     def _get_conn_params(self) -> Dict[str, Optional[str]]:
         """
@@ -62,6 +63,7 @@ def _get_conn_params(self) -> Dict[str, Optional[str]]:
         role = conn.extra_dejson.get('role', '')
         schema = conn.schema or ''
         authenticator = conn.extra_dejson.get('authenticator', 'snowflake')
+        session_parameters = conn.extra_dejson.get('session_parameters')
 
         conn_config = {
             "user": conn.login,
@@ -73,6 +75,7 @@ def _get_conn_params(self) -> Dict[str, Optional[str]]:
             "region": self.region or region,
             "role": self.role or role,
             "authenticator": self.authenticator or authenticator,
+            "session_parameters": self.session_parameters or session_parameters,
         }
 
         # If private_key_file is specified in the extra json, load the contents of the file as a private

airflow/providers/snowflake/operators/snowflake.py

@@ -57,6 +57,9 @@ class SnowflakeOperator(BaseOperator):
         'https://<your_okta_account_name>.okta.com' to authenticate
         through native Okta.
     :type authenticator: str
+    :param session_parameters: You can set session-level parameters at
+        the time you connect to Snowflake
+    :type session_parameters: dict
     """
 
     template_fields = ('sql',)
@@ -76,6 +79,7 @@ def __init__(
         role: Optional[str] = None,
         schema: Optional[str] = None,
         authenticator: Optional[str] = None,
+        session_parameters: Optional[dict] = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -88,6 +92,7 @@ def __init__(
         self.role = role
         self.schema = schema
         self.authenticator = authenticator
+        self.session_parameters = session_parameters
 
     def get_hook(self) -> SnowflakeHook:
         """
@@ -102,6 +107,7 @@ def get_hook(self) -> SnowflakeHook:
             role=self.role,
             schema=self.schema,
             authenticator=self.authenticator,
+            session_parameters=self.session_parameters,
         )
 
     def execute(self, context: Any) -> None:

tests/providers/snowflake/hooks/test_snowflake.py

@@ -55,7 +55,7 @@ def get_conn(self):
             def get_connection(self, _):
                 return conn
 
-        self.db_hook = UnitTestSnowflakeHook()
+        self.db_hook = UnitTestSnowflakeHook(session_parameters={"QUERY_TAG": "This is a test hook"})
 
         self.non_encrypted_private_key = "/tmp/test_key.pem"
         self.encrypted_private_key = "/tmp/test_key.p8"
@@ -100,6 +100,7 @@ def test_get_conn_params(self):
             'region': 'af_region',
             'role': 'af_role',
             'authenticator': 'snowflake',
+            'session_parameters': {"QUERY_TAG": "This is a test hook"},
         }
         self.assertEqual(self.db_hook.snowflake_conn_id, 'snowflake_default')  # pylint: disable=no-member
         self.assertEqual(conn_params_shouldbe, self.db_hook._get_conn_params())