Support UUID type

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

Hide file tree

Showing 3 changed files with 15 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,7 @@ 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,14 @@ def map(self, values: Any) -> Optional[Dict[Any, Optional[Any]]]:
         }
 
 
+class UuidValueMapper(ValueMapper[uuid.UUID]):
+    def map(self, values: Any) -> Optional[uuid.UUID]:
+        try:
+            return uuid.UUID(values)
+        except:
+            return None
+
+
 class NoOpRowMapper:
     """
     No-op RowMapper which does not perform any transformation
@@ -1245,6 +1254,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
@@ -11,6 +11,7 @@
 # limitations under the License.
 import json
 import re
+import uuid
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, Union
 
 from sqlalchemy import util
@@ -122,7 +123,7 @@ def get_col_spec(self, **kw):
     #
     # === Mixed ===
     # 'ipaddress': IPADDRESS
-    # 'uuid': UUID,
+    'uuid': uuid.UUID,
     # 'hyperloglog': HYPERLOGLOG,
     # 'p4hyperloglog': P4HYPERLOGLOG,
     # 'qdigest': QDIGEST,