Feature/trino provider timezone (#35963)

* feat(trino): support config timezone * add ut * Update docs/apache-airflow-providers-trino/connections.rst Co-authored-by: Philippe Gagnon <[email protected]> --------- Co-authored-by: Duyet Le <[email protected]> Co-authored-by: Elad Kalif <[email protected]> Co-authored-by: Philippe Gagnon <[email protected]>
apache · Dec 1, 2023 · 1c6bbe2 · 1c6bbe2
1 parent 86b1bd2
commit 1c6bbe2
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 1 deletion.
diff --git a/airflow/providers/trino/hooks/trino.py b/airflow/providers/trino/hooks/trino.py
@@ -156,6 +156,7 @@ def get_conn(self) -> Connection:
             verify=_boolify(extra.get("verify", True)),
             session_properties=extra.get("session_properties") or None,
             client_tags=extra.get("client_tags") or None,
+            timezone=extra.get("timezone") or None,
         )
 
         return trino_conn

diff --git a/docs/apache-airflow-providers-trino/connections.rst b/docs/apache-airflow-providers-trino/connections.rst
@@ -54,5 +54,6 @@ Extra (optional, connection parameters)
     * ``kerberos__service_name``, ``kerberos__config``, ``kerberos__mutual_authentication``, ``kerberos__force_preemptive``, ``kerberos__hostname_override``, ``kerberos__sanitize_mutual_error_response``, ``kerberos__principal``,``kerberos__delegate``, ``kerberos__ca_bundle`` - These parameters can be set when enabling ``kerberos`` authentication.
     * ``session_properties`` - JSON dictionary which allows to set session_properties. Example: ``{'session_properties':{'scale_writers':true,'task_writer_count:1'}}``
     * ``client_tags`` - List of comma separated tags. Example ``{'client_tags':['sales','cluster1']}```
+    * ``timezone`` - The time zone for the session can be explicitly set using the IANA time zone name. Example: ``{'timezone':'Asia/Jerusalem'}``.
 
     Note: If ``jwt__file`` and ``jwt__token`` are both given, ``jwt__file`` will take precedent.
diff --git a/tests/providers/trino/hooks/test_trino.py b/tests/providers/trino/hooks/test_trino.py
@@ -239,6 +239,14 @@ def test_get_conn_verify(self, mock_connect, mock_get_connection, current_verify
         TrinoHook().get_conn()
         self.assert_connection_called_with(mock_connect, verify=expected_verify)
 
+    @patch(HOOK_GET_CONNECTION)
+    @patch(TRINO_DBAPI_CONNECT)
+    def test_get_conn_timezone(self, mock_connect, mock_get_connection):
+        extras = {"timezone": "Asia/Jerusalem"}
+        self.set_get_connection_return_value(mock_get_connection, extra=json.dumps(extras))
+        TrinoHook().get_conn()
+        self.assert_connection_called_with(mock_connect, timezone="Asia/Jerusalem")
+
     @staticmethod
     def set_get_connection_return_value(mock_get_connection, extra=None, password=None):
         mocked_connection = Connection(
@@ -248,7 +256,13 @@ def set_get_connection_return_value(mock_get_connection, extra=None, password=No
 
     @staticmethod
     def assert_connection_called_with(
-        mock_connect, http_headers=mock.ANY, auth=None, verify=True, session_properties=None, client_tags=None
+        mock_connect,
+        http_headers=mock.ANY,
+        auth=None,
+        verify=True,
+        session_properties=None,
+        client_tags=None,
+        timezone=None,
     ):
         mock_connect.assert_called_once_with(
             catalog="hive",
@@ -264,6 +278,7 @@ def assert_connection_called_with(
             verify=verify,
             session_properties=session_properties,
             client_tags=client_tags,
+            timezone=timezone,
         )