Add support for UUID type in PostgreSQL connector

presto-postgresql/src/main/java/io/prestosql/plugin/postgresql/PostgreSqlClient.java

@@ -69,11 +69,14 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.function.BiFunction;
 
 import static com.fasterxml.jackson.core.JsonFactory.Feature.CANONICALIZE_FIELD_NAMES;
 import static com.fasterxml.jackson.databind.SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS;
+import static io.airlift.slice.SizeOf.SIZE_OF_LONG;
 import static io.airlift.slice.Slices.utf8Slice;
+import static io.airlift.slice.Slices.wrappedLongArray;
 import static io.prestosql.plugin.jdbc.ColumnMapping.DISABLE_PUSHDOWN;
 import static io.prestosql.plugin.jdbc.JdbcErrorCode.JDBC_ERROR;
 import static io.prestosql.plugin.jdbc.StandardColumnMappings.timestampColumnMapping;
@@ -103,13 +106,15 @@ public class PostgreSqlClient
     private static final String DUPLICATE_TABLE_SQLSTATE = "42P07";
 
     private final Type jsonType;
+    private final Type uuidType;
     private final boolean supportArrays;
 
     @Inject
     public PostgreSqlClient(BaseJdbcConfig config, PostgreSqlConfig postgreSqlConfig, TypeManager typeManager)
     {
         super(config, "\"", new DriverConnectionFactory(new Driver(), config));
         this.jsonType = typeManager.getType(new TypeSignature(StandardTypes.JSON));
+        this.uuidType = typeManager.getType(new TypeSignature(StandardTypes.UUID));
 
         switch (postgreSqlConfig.getArrayMapping()) {
             case DISABLED:
@@ -248,6 +253,8 @@ public Optional<ColumnMapping> toPrestoType(ConnectorSession session, Connection
                 .orElseThrow(() -> new PrestoException(JDBC_ERROR, "Type name is missing: " + typeHandle));
 
         switch (jdbcTypeName) {
+            case "uuid":
+                return Optional.of(uuidColumnMapping());
             case "jsonb":
             case "json":
                 return Optional.of(jsonColumnMapping());
@@ -306,6 +313,9 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)
         if (type.getTypeSignature().getBase().equals(StandardTypes.JSON)) {
             return WriteMapping.sliceMapping("jsonb", typedVarcharWriteFunction("json"));
         }
+        if (type.getTypeSignature().getBase().equals(StandardTypes.UUID)) {
+            return WriteMapping.sliceMapping("uuid", uuidWriteFunction());
+        }
         if (type instanceof ArrayType && supportArrays) {
             Type elementType = ((ArrayType) type).getElementType();
             String elementDataType = toWriteMapping(session, elementType).getDataType();
@@ -417,6 +427,27 @@ private static SliceWriteFunction typedVarcharWriteFunction(String jdbcTypeName)
         };
     }
 
+    private static SliceWriteFunction uuidWriteFunction()
+    {
+        return (statement, index, value) -> {
+            UUID uuid = new UUID(value.getLong(0), value.getLong(SIZE_OF_LONG));
+            statement.setObject(index, uuid, Types.OTHER);
+        };
+    }
+
+    private static Slice uuidSlice(UUID uuid)
+    {
+        return wrappedLongArray(uuid.getMostSignificantBits(), uuid.getLeastSignificantBits());
+    }
+
+    private ColumnMapping uuidColumnMapping()
+    {
+        return ColumnMapping.sliceMapping(
+                uuidType,
+                (resultSet, columnIndex) -> uuidSlice((UUID) resultSet.getObject(columnIndex)),
+                uuidWriteFunction());
+    }
+
     private static final JsonFactory JSON_FACTORY = new JsonFactory()
             .disable(CANONICALIZE_FIELD_NAMES);
 

presto-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java

@@ -42,6 +42,7 @@
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
 import java.util.List;
+import java.util.UUID;
 import java.util.function.Function;
 
 import static com.google.common.base.Preconditions.checkState;
@@ -67,6 +68,7 @@
 import static io.prestosql.tests.datatype.DataType.varbinaryDataType;
 import static io.prestosql.tests.datatype.DataType.varcharDataType;
 import static io.prestosql.type.JsonType.JSON;
+import static io.prestosql.type.UuidType.UUID;
 import static java.lang.String.format;
 import static java.nio.charset.StandardCharsets.UTF_16LE;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -654,6 +656,20 @@ private DataTypeTest jsonTestCases(DataType<String> jsonDataType)
                 .addRoundTrip(jsonDataType, "[]");
     }
 
+    @Test
+    public void testUuid()
+    {
+        uuidTestCases(uuidDataType())
+                .execute(getQueryRunner(), prestoCreateAsSelect("presto_test_uuid"));
+    }
+
+    private DataTypeTest uuidTestCases(DataType<java.util.UUID> uuidDataType)
+    {
+        return DataTypeTest.create()
+                .addRoundTrip(uuidDataType, java.util.UUID.fromString("00000000-0000-0000-0000-000000000000"))
+                .addRoundTrip(uuidDataType, java.util.UUID.fromString("123e4567-e89b-12d3-a456-426655440000"));
+    }
+
     private void testUnsupportedDataType(String databaseDataType)
     {
         JdbcSqlExecutor jdbcSqlExecutor = new JdbcSqlExecutor(postgreSqlServer.getJdbcUrl());
@@ -698,6 +714,15 @@ public static DataType<String> jsonbDataType()
                 identity());
     }
 
+    public static DataType<java.util.UUID> uuidDataType()
+    {
+        return dataType(
+                "uuid",
+                UUID,
+                value -> "UUID " + formatStringLiteral(value.toString()),
+                identity());
+    }
+
     private static DataType<byte[]> byteaDataType()
     {
         return dataType(

presto-tests/src/main/java/io/prestosql/tests/TestingPrestoClient.java

@@ -73,6 +73,7 @@
 import static io.prestosql.type.IntervalDayTimeType.INTERVAL_DAY_TIME;
 import static io.prestosql.type.IntervalYearMonthType.INTERVAL_YEAR_MONTH;
 import static io.prestosql.type.JsonType.JSON;
+import static io.prestosql.type.UuidType.UUID;
 import static java.util.stream.Collectors.toList;
 
 public class TestingPrestoClient
@@ -192,6 +193,9 @@ else if (DOUBLE.equals(type)) {
         else if (REAL.equals(type)) {
             return ((Number) value).floatValue();
         }
+        else if (UUID.equals(type)) {
+            return java.util.UUID.fromString((String) value);
+        }
         else if (type instanceof VarcharType) {
             return value;
         }