Support UUID type

trinodb · Apr 6, 2023 · 9529547 · 9529547
1 parent e10bc53
commit 9529547
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/tests/integration/test_types_integration.py b/tests/integration/test_types_integration.py
@@ -1,5 +1,6 @@
 import math
 import re
+import uuid
 from datetime import date, datetime, time, timedelta, timezone, tzinfo
 from decimal import Decimal
 
@@ -774,7 +775,8 @@ def test_ipaddress(trino_connection):
 def test_uuid(trino_connection):
     SqlTest(trino_connection) \
         .add_field(sql="CAST(null AS UUID)", python=None) \
-        .add_field(sql="UUID '12151fd2-7586-11e9-8f9e-2a86e4085a59'", python='12151fd2-7586-11e9-8f9e-2a86e4085a59') \
+        .add_field(sql="UUID '12151fd2-7586-11e9-8f9e-2a86e4085a59'",
+                   python=uuid.UUID('12151fd2-7586-11e9-8f9e-2a86e4085a59')) \
         .execute()
 
 

diff --git a/trino/client.py b/trino/client.py
@@ -43,6 +43,7 @@
 import re
 import threading
 import urllib.parse
+import uuid
 import warnings
 from dataclasses import dataclass
 from datetime import date, datetime, time, timedelta, timezone, tzinfo
@@ -1185,6 +1186,13 @@ def map(self, values: Any) -> Optional[Dict[Any, Optional[Any]]]:
         }
 
 
+class UuidValueMapper(ValueMapper[uuid.UUID]):
+    def map(self, value: Any) -> Optional[uuid.UUID]:
+        if value is None:
+            return None
+        return uuid.UUID(value)
+
+
 class NoOpRowMapper:
     """
     No-op RowMapper which does not perform any transformation
@@ -1245,6 +1253,8 @@ def _create_value_mapper(self, column) -> ValueMapper:
             return DateValueMapper()
         elif col_type == 'varbinary':
             return BinaryValueMapper()
+        elif col_type == 'uuid':
+            return UuidValueMapper()
         else:
             return NoOpValueMapper()
 

diff --git a/trino/sqlalchemy/datatype.py b/trino/sqlalchemy/datatype.py
@@ -25,6 +25,10 @@ class DOUBLE(sqltypes.Float):
     __visit_name__ = "DOUBLE"
 
 
+class UUID(TypeEngine):
+    __visit_name__ = "UUID"
+
+
 class MAP(TypeEngine):
     __visit_name__ = "MAP"
 
@@ -122,7 +126,7 @@ def get_col_spec(self, **kw):
     #
     # === Mixed ===
     # 'ipaddress': IPADDRESS
-    # 'uuid': UUID,
+    'uuid': UUID,
     # 'hyperloglog': HYPERLOGLOG,
     # 'p4hyperloglog': P4HYPERLOGLOG,
     # 'qdigest': QDIGEST,
@@ -219,6 +223,8 @@ def parse_sqltype(type_str: str) -> TypeEngine:
             attr_type = parse_sqltype(attr_type_str)
             attr_types.append((attr_name, attr_type))
         return ROW(attr_types)
+    elif type_name == "uuid":
+        return UUID()
 
     if type_name not in _type_map:
         util.warn(f"Did not recognize type '{type_name}'")