Add support for TIMEZONE

Author: ulisesojeda

.pre-commit-config.yaml

@@ -13,3 +13,4 @@ repos:
         additional_dependencies:
           - "types-pytz"
           - "types-requests"
+          - "types-tzlocal"

README.md

@@ -357,6 +357,20 @@ conn = trino.dbapi.connect(
 )
 ```
 
+## Timezone
+
+The time zone for the session can be explicitly set using the official IANA time zone name. When not set the time zone defaults to the client side local timezone.
+
+```python
+import trino
+conn = trino.dbapi.connect(
+    host='localhost',
+    port=443,
+    user='username',
+    timezone="Europe/Brussels",
+)
+```
+
 ## SSL
 
 ### SSL verification

setup.py

@@ -76,7 +76,7 @@
         "Topic :: Database :: Front-Ends",
     ],
     python_requires='>=3.7',
-    install_requires=["pytz", "requests"],
+    install_requires=["pytz", "requests", "tzlocal"],
     extras_require={
         "all": all_require,
         "kerberos": kerberos_require,

tests/integration/test_dbapi_integration.py

@@ -16,6 +16,7 @@
 import pytest
 import pytz
 import requests
+from tzlocal import get_localzone_name
 
 import trino
 from tests.integration.conftest import trino_version
@@ -1120,3 +1121,31 @@ def test_prepared_statements(run_trino):
     cur.execute('DEALLOCATE PREPARE test_prepared_statements')
     cur.fetchall()
     assert cur._request._client_session.prepared_statements == {}
+
+
+def test_set_timezone_in_connection(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", catalog="tpch", timezone="Europe/Brussels"
+    )
+    cur = trino_connection.cursor()
+    cur.execute('SELECT current_timezone()')
+    res = cur.fetchall()
+    assert res[0][0] == "Europe/Brussels"
+
+
+def test_connection_without_timezone(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", catalog="tpch"
+    )
+    cur = trino_connection.cursor()
+    cur.execute('SELECT current_timezone()')
+    res = cur.fetchall()
+    session_tz = res[0][0]
+    localzone = get_localzone_name()
+    assert session_tz == localzone or \
+        (session_tz == "UTC" and localzone == "Etc/UTC") \
+        # Workaround for difference between Trino timezone and tzlocal for UTC

tests/unit/test_client.py

@@ -13,6 +13,7 @@
 import threading
 import time
 import uuid
+from tzlocal import get_localzone_name
 from typing import Optional, Dict
 from unittest import mock
 from urllib.parse import urlparse
@@ -32,6 +33,12 @@
 from trino.client import TrinoQuery, TrinoRequest, TrinoResult, ClientSession, _DelayExponential, _retry_with, \
     _RetryWithExponentialBackoff
 
+try:
+    from zoneinfo._common import ZoneInfoNotFoundError  # type: ignore
+
+except ModuleNotFoundError:
+    from backports.zoneinfo._common import ZoneInfoNotFoundError  # type: ignore
+
 
 @mock.patch("trino.client.TrinoRequest.http")
 def test_trino_initial_request(mock_requests, sample_post_response_data):
@@ -65,6 +72,7 @@ def test_request_headers(mock_get_and_post):
     schema = "test_schema"
     user = "test_user"
     source = "test_source"
+    timezone = "Europe/Brussels"
     accept_encoding_header = "accept-encoding"
     accept_encoding_value = "identity,deflate,gzip"
     client_info_header = constants.HEADER_CLIENT_INFO
@@ -78,6 +86,7 @@ def test_request_headers(mock_get_and_post):
             source=source,
             catalog=catalog,
             schema=schema,
+            timezone=timezone,
             headers={
                 accept_encoding_header: accept_encoding_value,
                 client_info_header: client_info_value,
@@ -93,9 +102,10 @@ def assert_headers(headers):
         assert headers[constants.HEADER_SOURCE] == source
         assert headers[constants.HEADER_USER] == user
         assert headers[constants.HEADER_SESSION] == ""
+        assert headers[constants.HEADER_TIMEZONE] == timezone
         assert headers[accept_encoding_header] == accept_encoding_value
         assert headers[client_info_header] == client_info_value
-        assert len(headers.keys()) == 8
+        assert len(headers.keys()) == 9
 
     req.post("URL")
     _, post_kwargs = post.call_args
@@ -1056,3 +1066,62 @@ def test_request_headers_role_empty(mock_get_and_post):
     req.get("URL")
     _, get_kwargs = get.call_args
     assert_headers_with_roles(post_kwargs["headers"], None)
+
+
+def assert_headers_timezone(headers: Dict[str, str], timezone: str):
+    assert headers[constants.HEADER_TIMEZONE] == timezone
+
+
+def test_request_headers_with_timezone(mock_get_and_post):
+    get, post = mock_get_and_post
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=8080,
+        client_session=ClientSession(
+            user="test_user",
+            timezone="Europe/Brussels"
+        ),
+    )
+
+    req.post("URL")
+    _, post_kwargs = post.call_args
+    assert_headers_timezone(post_kwargs["headers"], "Europe/Brussels")
+
+    req.get("URL")
+    _, get_kwargs = get.call_args
+    assert_headers_timezone(post_kwargs["headers"], "Europe/Brussels")
+
+
+def test_request_headers_without_timezone(mock_get_and_post):
+    get, post = mock_get_and_post
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=8080,
+        client_session=ClientSession(
+            user="test_user",
+        ),
+    )
+    localzone = get_localzone_name()
+
+    req.post("URL")
+    _, post_kwargs = post.call_args
+    assert_headers_timezone(post_kwargs["headers"], localzone)
+
+    req.get("URL")
+    _, get_kwargs = get.call_args
+    assert_headers_timezone(post_kwargs["headers"], localzone)
+
+
+def test_request_with_invalid_timezone(mock_get_and_post):
+    with pytest.raises(ZoneInfoNotFoundError) as zinfo_error:
+        TrinoRequest(
+            host="coordinator",
+            port=8080,
+            client_session=ClientSession(
+                user="test_user",
+                timezone="INVALID_TIMEZONE"
+            ),
+        )
+    assert str(zinfo_error.value).startswith("'No time zone found with key")

trino/client.py

@@ -43,6 +43,7 @@
 import urllib.parse
 from datetime import datetime, timedelta, timezone
 from decimal import Decimal
+from tzlocal import get_localzone_name
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import pytz
@@ -51,6 +52,13 @@
 import trino.logging
 from trino import constants, exceptions
 
+try:
+    from zoneinfo import ZoneInfo  # type: ignore
+
+except ModuleNotFoundError:
+    from backports.zoneinfo import ZoneInfo  # type: ignore
+
+
 __all__ = ["ClientSession", "TrinoQuery", "TrinoRequest", "PROXIES"]
 
 logger = trino.logging.get_logger(__name__)
@@ -100,6 +108,8 @@ class ClientSession(object):
     :param client_tags: Client tags as list of strings.
     :param roles: roles for the current session. Some connectors do not
                  support role management. See connector documentation for more details.
+    :param timezone: The timezone for query processing. Defaults to the timezone
+                      of the Trino cluster, and not the timezone of the client.
     """
 
     def __init__(
@@ -114,6 +124,7 @@ def __init__(
         extra_credential: List[Tuple[str, str]] = None,
         client_tags: List[str] = None,
         roles: Dict[str, str] = None,
+        timezone: str = None,
     ):
         self._user = user
         self._catalog = catalog
@@ -127,6 +138,9 @@ def __init__(
         self._roles = roles.copy() if roles is not None else {}
         self._prepared_statements: Dict[str, str] = {}
         self._object_lock = threading.Lock()
+        if timezone:  # Check timezone validity
+            ZoneInfo(timezone)
+        self._timezone = timezone or get_localzone_name()
 
     @property
     def user(self):
@@ -207,6 +221,11 @@ def prepared_statements(self, prepared_statements):
         with self._object_lock:
             self._prepared_statements = prepared_statements
 
+    @property
+    def timezone(self):
+        with self._object_lock:
+            return self._timezone
+
     def __getstate__(self):
         state = self.__dict__.copy()
         del state["_object_lock"]
@@ -408,6 +427,7 @@ def http_headers(self) -> Dict[str, str]:
         headers[constants.HEADER_SCHEMA] = self._client_session.schema
         headers[constants.HEADER_SOURCE] = self._client_session.source
         headers[constants.HEADER_USER] = self._client_session.user
+        headers[constants.HEADER_TIMEZONE] = self._client_session.timezone
         if len(self._client_session.roles.values()):
             headers[constants.HEADER_ROLE] = ",".join(
                 # ``name`` must not contain ``=``

trino/constants.py

@@ -34,6 +34,7 @@
 HEADER_CLIENT_INFO = "X-Trino-Client-Info"
 HEADER_CLIENT_TAGS = "X-Trino-Client-Tags"
 HEADER_EXTRA_CREDENTIAL = "X-Trino-Extra-Credential"
+HEADER_TIMEZONE = "X-Trino-Time-Zone"
 
 HEADER_SESSION = "X-Trino-Session"
 HEADER_SET_SESSION = "X-Trino-Set-Session"

trino/dbapi.py

@@ -111,6 +111,7 @@ def __init__(
         client_tags=None,
         experimental_python_types=False,
         roles=None,
+        timezone=None,
     ):
         self.host = host
         self.port = port
@@ -130,6 +131,7 @@ def __init__(
             extra_credential=extra_credential,
             client_tags=client_tags,
             roles=roles,
+            timezone=timezone,
         )
         # mypy cannot follow module import
         if http_session is None: