Support decimal, date, time, timestamp with time zone and timestamp

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -10,7 +10,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-from datetime import datetime
+from datetime import datetime, time
+from decimal import Decimal
 
 import pytest
 import pytz
@@ -123,22 +124,114 @@ def test_string_query_param(trino_connection):
     assert rows[0][0] == "six'"
 
 
+def test_decimal_query_param(trino_connection):
+    cur = trino_connection.cursor()
+
+    cur.execute("SELECT ?", params=(Decimal('0.142857'),))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == Decimal('0.142857')
+
+
 def test_datetime_query_param(trino_connection):
     cur = trino_connection.cursor()
 
-    cur.execute("SELECT ?", params=(datetime(2020, 1, 1, 0, 0, 0),))
+    params = datetime(2020, 1, 1, 16, 43, 22, 320000)
+
+    cur.execute("SELECT ?", params=(params,))
     rows = cur.fetchall()
 
-    assert rows[0][0] == "2020-01-01 00:00:00.000"
+    assert rows[0][0] == params
+    assert cur.description[0][1] == "timestamp"
+
+
+def test_datetime_with_time_zone_query_param(trino_connection):
+    cur = trino_connection.cursor()
+
+    params = datetime(2020, 1, 1, 16, 43, 22, 320000, tzinfo=pytz.timezone('CET'))
 
     cur.execute("SELECT ?",
-                params=(datetime(2020, 1, 1, 0, 0, 0, tzinfo=pytz.utc),))
+                params=(params,))
     rows = cur.fetchall()
 
-    assert rows[0][0] == "2020-01-01 00:00:00.000 UTC"
+    assert rows[0][0] == params
     assert cur.description[0][1] == "timestamp with time zone"
 
 
+def test_datetime_with_time_zone_numeric_offset(trino_connection):
+    cur = trino_connection.cursor()
+
+    cur.execute("SELECT TIMESTAMP '2001-08-22 03:04:05.321 -08:00'")
+    rows = cur.fetchall()
+
+    assert rows[0][0] == datetime.strptime("2001-08-22 03:04:05.321 -08:00", "%Y-%m-%d %H:%M:%S.%f %z")
+
+
+def test_special_datetimes_query_param(trino_connection):
+    cur = trino_connection.cursor()
+
+    for special_date in (
+            datetime.fromtimestamp(1603589478, pytz.timezone('Europe/Warsaw')),
+    ):
+        params = special_date
+
+        cur.execute("SELECT ?", params=(params,))
+        rows = cur.fetchall()
+
+        assert rows[0][0] == params
+
+
+def test_date_query_param(trino_connection):
+    cur = trino_connection.cursor()
+
+    params = datetime(2020, 1, 1, 0, 0, 0).date()
+
+    cur.execute("SELECT ?", params=(params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == params
+
+
+def test_special_dates_query_param(trino_connection):
+    cur = trino_connection.cursor()
+
+    for params in (
+        # datetime(-1, 1, 1, 0, 0, 0).date(),
+        # datetime(0, 1, 1, 0, 0, 0).date(),
+        datetime(1752, 9, 4, 0, 0, 0).date(),
+        datetime(1970, 1, 1, 0, 0, 0).date(),
+    ):
+        cur.execute("SELECT ?", params=(params,))
+        rows = cur.fetchall()
+
+        assert rows[0][0] == params
+
+
+def test_time_query_param(trino_connection):
+    cur = trino_connection.cursor()
+
+    params = time(12, 3, 44, 333000)
+
+    cur.execute("SELECT ?", params=(params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == params
+
+
+@pytest.mark.skip(reason="time with time zone currently not supported")
+def test_time_with_time_zone_query_param(trino_connection):
+    cur = trino_connection.cursor()
+
+    params = time(16, 43, 22, 320000, tzinfo=pytz.timezone('CET'))
+
+    cur.execute("SELECT ?",
+                params=(params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == params
+    assert cur.description[0][1] == "time with time zone"
+
+
 def test_array_query_param(trino_connection):
     cur = trino_connection.cursor()
 
@@ -158,6 +251,38 @@ def test_array_query_param(trino_connection):
     assert rows[0][0] == "array(integer)"
 
 
+def test_array_timestamp_query_param(trino_connection):
+    cur = trino_connection.cursor()
+
+    params = [datetime(2020, 1, 1, 0, 0, 0), datetime(2020, 1, 2, 0, 0, 0)]
+
+    cur.execute("SELECT ?", params=(params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == params
+
+    cur.execute("SELECT TYPEOF(?)", params=(params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == "array(timestamp(6))"
+
+
+def test_array_timestamp_with_timezone_query_param(trino_connection):
+    cur = trino_connection.cursor()
+
+    params = [datetime(2020, 1, 1, 0, 0, 0, tzinfo=pytz.utc), datetime(2020, 1, 2, 0, 0, 0, tzinfo=pytz.utc)]
+
+    cur.execute("SELECT ?", params=(params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == params
+
+    cur.execute("SELECT TYPEOF(?)", params=(params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == "array(timestamp(6) with time zone)"
+
+
 def test_dict_query_param(trino_connection):
     cur = trino_connection.cursor()
 

trino/client.py

@@ -36,6 +36,9 @@
 import copy
 import os
 import re
+from decimal import Decimal
+from datetime import datetime
+import pytz
 from typing import Any, Dict, List, Optional, Tuple, Union
 import urllib.parse
 
@@ -494,12 +497,41 @@ def __iter__(self):
             for row in rows:
                 self._rownumber += 1
                 logger.debug("row %s", row)
-                yield row
+                yield self._map_to_python_types(row, self._query.columns)
 
     @property
     def response_headers(self):
         return self._query.response_headers
 
+    @classmethod
+    def _map_to_python_type(cls, item: Tuple[Any, Dict]) -> Any:
+        (value, data_type) = item
+
+        raw_type = data_type["typeSignature"]["rawType"]
+        if isinstance(value, list):
+            raw_type = {
+                "typeSignature": data_type["typeSignature"]["arguments"][0]["value"]
+            }
+            return [cls._map_to_python_type((array_item, raw_type)) for array_item in value]
+        elif "decimal" in raw_type:
+            return Decimal(value)
+        elif raw_type == "date":
+            return datetime.strptime(value, "%Y-%m-%d").date()
+        elif raw_type == "timestamp with time zone":
+            dt, tz = value.rsplit(' ', 1)
+            if tz.startswith('+') or tz.startswith('-'):
+                return datetime.strptime(value, "%Y-%m-%d %H:%M:%S.%f %z")
+            return datetime.strptime(dt, "%Y-%m-%d %H:%M:%S.%f").replace(tzinfo=pytz.timezone(tz))
+        elif "timestamp" in raw_type:
+            return datetime.strptime(value, "%Y-%m-%d %H:%M:%S.%f")
+        elif "time" in raw_type:
+            return datetime.strptime(value, "%H:%M:%S.%f").time()
+        else:
+            return value
+
+    def _map_to_python_types(self, row: List[Any], columns: List[Dict[str, Any]]) -> List[Any]:
+        return list(map(self._map_to_python_type, zip(row, columns)))
+
 
 class TrinoQuery(object):
     """Represent the execution of a SQL statement by Trino."""

trino/dbapi.py

@@ -17,7 +17,7 @@
 Fetch methods returns rows as a list of lists on purpose to let the caller
 decide to convert then to a list of tuples.
 """
-
+from decimal import Decimal
 from typing import Any, List, Optional  # NOQA for mypy types
 
 import copy
@@ -365,6 +365,14 @@ def _format_prepared_param(self, param):
             datetime_str = datetime_str.rstrip(" ")
             return "TIMESTAMP '%s'" % datetime_str
 
+        if isinstance(param, datetime.time):
+            time_str = param.strftime("%H:%M:%S.%f")
+            return "TIME '%s'" % time_str
+
+        if isinstance(param, datetime.date):
+            date_str = param.strftime("%Y-%m-%d")
+            return "DATE '%s'" % date_str
+
         if isinstance(param, list):
             return "ARRAY[%s]" % ','.join(map(self._format_prepared_param, param))
 
@@ -379,6 +387,9 @@ def _format_prepared_param(self, param):
         if isinstance(param, uuid.UUID):
             return "UUID '%s'" % param
 
+        if isinstance(param, Decimal):
+            return "DECIMAL '%s'" % param
+
         raise trino.exceptions.NotSupportedError("Query parameter of type '%s' is not supported." % type(param))
 
     def _deallocate_prepare_statement(self, added_prepare_header, statement_name):