Support Trino's ROW datatype

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -616,6 +616,32 @@ def test_boolean_query_param(trino_connection):
     assert rows[0][0] is False
 
 
+def test_row(trino_connection):
+    cur = trino_connection.cursor(experimental_python_types=True)
+    params = (1, Decimal("2.0"), datetime(2020, 1, 1, 0, 0, 0))
+    cur.execute("SELECT ?", (params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == params
+
+
+def test_nested_row(trino_connection):
+    cur = trino_connection.cursor(experimental_python_types=True)
+    params = ((1, "test", Decimal("3.1")), Decimal("2.0"), datetime(2020, 1, 1, 0, 0, 0))
+    cur.execute("SELECT ?", (params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == params
+
+
+def test_named_row(trino_connection):
+    cur = trino_connection.cursor(experimental_python_types=True)
+    cur.execute("SELECT CAST(ROW(1, 2e0) AS ROW(x BIGINT, y DOUBLE))")
+    rows = cur.fetchall()
+
+    assert rows[0][0] == (1, 2.0)
+
+
 def test_float_query_param(trino_connection):
     cur = trino_connection.cursor()
     cur.execute("SELECT ?", params=(1.1,))

trino/client.py

@@ -530,25 +530,30 @@ def _map_to_python_type(cls, item: Tuple[Any, Dict]) -> Any:
 
         try:
             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]
-            if isinstance(value, dict):
-                if raw_type == "map":
-                    raw_key_type = {
+                if raw_type == "array":
+                    raw_type = {
                         "typeSignature": data_type["typeSignature"]["arguments"][0]["value"]
                     }
-                    raw_value_type = {
-                        "typeSignature": data_type["typeSignature"]["arguments"][1]["value"]
-                    }
-                    return {
-                        cls._map_to_python_type((key, raw_key_type)):
-                            cls._map_to_python_type((value[key], raw_value_type))
-                        for key in value
-                    }
-                # TODO: support row type
+                    return [cls._map_to_python_type((array_item, raw_type)) for array_item in value]
+                if raw_type == "row":
+                    raw_types = map(lambda arg: arg["value"], data_type["typeSignature"]["arguments"])
+                    return tuple(
+                        cls._map_to_python_type((array_item, raw_type))
+                        for (array_item, raw_type) in zip(value, raw_types)
+                    )
                 return value
+            if isinstance(value, dict):
+                raw_key_type = {
+                    "typeSignature": data_type["typeSignature"]["arguments"][0]["value"]
+                }
+                raw_value_type = {
+                    "typeSignature": data_type["typeSignature"]["arguments"][1]["value"]
+                }
+                return {
+                    cls._map_to_python_type((key, raw_key_type)):
+                        cls._map_to_python_type((value[key], raw_value_type))
+                    for key in value
+                }
             elif "decimal" in raw_type:
                 return Decimal(value)
             elif raw_type == "double":

trino/dbapi.py

@@ -394,8 +394,10 @@ def _format_prepared_param(self, param):
         if isinstance(param, list):
             return "ARRAY[%s]" % ','.join(map(self._format_prepared_param, param))
 
+        if isinstance(param, tuple):
+            return "ROW(%s)" % ','.join(map(self._format_prepared_param, param))
+
         if isinstance(param, dict):
-            # TODO: support mixed types in dicts and convert to ROW
             keys = list(param.keys())
             values = [param[key] for key in keys]
             return "MAP({}, {})".format(