This is an automated email from the ASF dual-hosted git repository.
eladkal 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 1c6bbe2841 Feature/trino provider timezone (#35963)
1c6bbe2841 is described below
commit 1c6bbe2841fe846957f7a1ce68eb978c30669896
Author: rom sharon <[email protected]>
AuthorDate: Fri Dec 1 19:21:58 2023 +0200
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]>
---
airflow/providers/trino/hooks/trino.py | 1 +
docs/apache-airflow-providers-trino/connections.rst | 1 +
tests/providers/trino/hooks/test_trino.py | 17 ++++++++++++++++-
3 files changed, 18 insertions(+), 1 deletion(-)
diff --git a/airflow/providers/trino/hooks/trino.py
b/airflow/providers/trino/hooks/trino.py
index 798109dc3f..3369af13f0 100644
--- a/airflow/providers/trino/hooks/trino.py
+++ b/airflow/providers/trino/hooks/trino.py
@@ -156,6 +156,7 @@ class TrinoHook(DbApiHook):
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
index ee25fbbce7..5b2a9b3647 100644
--- 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
index 5d61c056bc..c5b82bc9a4 100644
--- a/tests/providers/trino/hooks/test_trino.py
+++ b/tests/providers/trino/hooks/test_trino.py
@@ -239,6 +239,14 @@ class TestTrinoHookConn:
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 @@ class TestTrinoHookConn:
@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 @@ class TestTrinoHookConn:
verify=verify,
session_properties=session_properties,
client_tags=client_tags,
+ timezone=timezone,
)
