From 2fce1f6dd541610369fbe640d014c20d5c98c9a8 Mon Sep 17 00:00:00 2001 From: Duyet Le Date: Tue, 24 Oct 2023 14:47:36 +0700 Subject: [PATCH 1/3] feat(trino): support config timezone --- airflow/providers/trino/hooks/trino.py | 1 + docs/apache-airflow-providers-trino/connections.rst | 1 + 2 files changed, 2 insertions(+) diff --git a/airflow/providers/trino/hooks/trino.py b/airflow/providers/trino/hooks/trino.py index 798109dc3f84a..3369af13f0839 100644 --- 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 index ee25fbbce7183..7cb5ce928439d 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/Ho_Chi_Minh'}``. Note: If ``jwt__file`` and ``jwt__token`` are both given, ``jwt__file`` will take precedent. From fa33eb34acf010d0bc32ac509acae4d72be34d95 Mon Sep 17 00:00:00 2001 From: Duyet Le Date: Tue, 24 Oct 2023 14:47:36 +0700 Subject: [PATCH 2/3] add ut --- .../connections.rst | 2 +- tests/providers/trino/hooks/test_trino.py | 17 ++++++++++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/docs/apache-airflow-providers-trino/connections.rst b/docs/apache-airflow-providers-trino/connections.rst index 7cb5ce928439d..da0e62e3a56a2 100644 --- a/docs/apache-airflow-providers-trino/connections.rst +++ b/docs/apache-airflow-providers-trino/connections.rst @@ -54,6 +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/Ho_Chi_Minh'}``. + * ``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 5d61c056bc3b2..c5b82bc9a43a6 100644 --- 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, ) From bce59a1872fa86e3e911fdbace10a97fa177f67f Mon Sep 17 00:00:00 2001 From: Elad Kalif <45845474+eladkal@users.noreply.github.com> Date: Fri, 1 Dec 2023 18:54:55 +0200 Subject: [PATCH 3/3] Update docs/apache-airflow-providers-trino/connections.rst Co-authored-by: Philippe Gagnon <12717218+pgagnon@users.noreply.github.com> --- docs/apache-airflow-providers-trino/connections.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/apache-airflow-providers-trino/connections.rst b/docs/apache-airflow-providers-trino/connections.rst index da0e62e3a56a2..5b2a9b36478f9 100644 --- a/docs/apache-airflow-providers-trino/connections.rst +++ b/docs/apache-airflow-providers-trino/connections.rst @@ -54,6 +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'}``. + * ``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.