diff --git a/tests/integration/test_types_integration.py b/tests/integration/test_types_integration.py index 201f627e..144f276f 100644 --- 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 index f66c207e..713cee44 100644 --- 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()