Support UUID type

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

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
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()