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

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -11,6 +11,7 @@
 # limitations under the License.
 import math
 from datetime import datetime
+from decimal import Decimal
 
 import pytest
 import pytz
@@ -123,19 +124,28 @@ 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),))
     rows = cur.fetchall()
 
-    assert rows[0][0] == "2020-01-01 00:00:00.000"
+    assert rows[0][0] == datetime(2020, 1, 1, 0, 0, 0)
 
     cur.execute("SELECT ?",
                 params=(datetime(2020, 1, 1, 0, 0, 0, tzinfo=pytz.utc),))
     rows = cur.fetchall()
 
-    assert rows[0][0] == "2020-01-01 00:00:00.000 UTC"
+    assert rows[0][0] == datetime(2020, 1, 1, 0, 0, 0, tzinfo=pytz.utc)
     assert cur.description[0][1] == "timestamp with time zone"
 
 
@@ -158,6 +168,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,10 @@
 import copy
 import os
 import re
+from decimal import Decimal
+from datetime import datetime
+import pytz
+from time import strptime
 from typing import Any, Dict, List, Optional, Tuple, Union
 import urllib.parse
 
@@ -494,12 +498,39 @@ 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 strptime(value, "%Y-%m-%d")
+        elif raw_type == "time":
+            return strptime(value, "%H:%M:%S")
+        elif "timestamp with time zone" in raw_type:
+            time_str = strptime(value, "%Y-%m-%d %H:%M:%S.%f %Z")
+            return datetime(*time_str[:3], tzinfo=pytz.timezone(time_str.tm_zone))
+        elif "timestamp" in raw_type:
+            return datetime(*strptime(value, "%Y-%m-%d %H:%M:%S.%f")[:3])
+        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
@@ -379,6 +379,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):