Support Trino's map data type

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -591,6 +591,16 @@ def test_dict_query_param(trino_connection):
     assert rows[0][0] == "map(varchar(3), varchar(3))"
 
 
+def test_dict_timestamp_query_param_types(trino_connection):
+    cur = trino_connection.cursor(experimental_python_types=True)
+
+    params = {"foo": datetime(2020, 1, 1, 16, 43, 22, 320000)}
+    cur.execute("SELECT ?", params=(params,))
+    rows = cur.fetchall()
+
+    assert rows[0][0] == params
+
+
 def test_boolean_query_param(trino_connection):
     cur = trino_connection.cursor()
 

trino/client.py

@@ -528,6 +528,21 @@ def _map_to_python_type(cls, item: Tuple[Any, Dict]) -> Any:
                     "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 = {
+                        "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 value
             elif "decimal" in raw_type:
                 return Decimal(value)
             elif raw_type == "double":

trino/dbapi.py

@@ -392,6 +392,7 @@ def _format_prepared_param(self, param):
             return "ARRAY[%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(