Support TIME WITH TIME ZONE in prepared statements

Author: hovaesco

tests/integration/test_dbapi_integration.py

@@ -527,23 +527,27 @@ def test_time_query_param(trino_connection):
 
 
 def test_time_with_named_time_zone_query_param(trino_connection):
-    with pytest.raises(trino.exceptions.NotSupportedError):
-        cur = trino_connection.cursor()
+    cur = trino_connection.cursor()
 
-        params = time(16, 43, 22, 320000, tzinfo=pytz.timezone('Asia/Shanghai'))
+    params = time(16, 43, 22, 320000, tzinfo=pytz.timezone('Asia/Shanghai'))
 
-        cur.execute("SELECT ?", params=(params,))
+    cur.execute("SELECT ?", params=(params,))
+    rows = cur.fetchall()
+
+    # Asia/Shanghai
+    assert rows[0][0].tzinfo == timezone(timedelta(seconds=28800))
 
 
 def test_time_with_numeric_offset_time_zone_query_param(trino_connection):
-    with pytest.raises(trino.exceptions.NotSupportedError):
-        cur = trino_connection.cursor()
+    cur = trino_connection.cursor()
 
-        tz = timezone(-timedelta(hours=8, minutes=0))
+    tz = timezone(-timedelta(hours=8, minutes=0))
+    params = time(16, 43, 22, 320000, tzinfo=tz)
 
-        params = time(16, 43, 22, 320000, tzinfo=tz)
+    cur.execute("SELECT ?", params=(params,))
+    rows = cur.fetchall()
 
-        cur.execute("SELECT ?", params=(params,))
+    assert rows[0][0] == params
 
 
 def test_time(trino_connection):

trino/dbapi.py

@@ -25,6 +25,8 @@
 from typing import Any, Dict, List, NamedTuple, Optional  # NOQA for mypy types
 from urllib.parse import urlparse
 
+import pytz
+
 import trino.client
 import trino.exceptions
 import trino.logging
@@ -433,6 +435,15 @@ def _format_prepared_param(self, param):
             time_str = param.strftime("%H:%M:%S.%f")
             return "TIME '%s'" % time_str
 
+        if isinstance(param, datetime.time) and param.tzinfo is not None:
+            time_str = param.strftime("%H:%M:%S.%f")
+            # named timezones
+            if hasattr(param.tzinfo, 'zone'):
+                utc_offset = datetime.datetime.now(pytz.timezone(param.tzinfo.zone)).strftime('%z')
+                return "TIME '%s %s:%s'" % (time_str, utc_offset[:3], utc_offset[3:])
+            # offset-based timezones
+            return "TIME '%s %s'" % (time_str, param.strftime('%Z')[3:])
+
         if isinstance(param, datetime.date):
             date_str = param.strftime("%Y-%m-%d")
             return "DATE '%s'" % date_str