diff --git a/presto-postgresql/src/main/java/io/prestosql/plugin/postgresql/PostgreSqlClient.java b/presto-postgresql/src/main/java/io/prestosql/plugin/postgresql/PostgreSqlClient.java index decbdd68d5e8a..f2ebed899797a 100644 --- a/presto-postgresql/src/main/java/io/prestosql/plugin/postgresql/PostgreSqlClient.java +++ b/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,6 +106,7 @@ public class PostgreSqlClient private static final String DUPLICATE_TABLE_SQLSTATE = "42P07"; private final Type jsonType; + private final Type uuidType; private final boolean supportArrays; @Inject @@ -110,6 +114,7 @@ public PostgreSqlClient(BaseJdbcConfig config, PostgreSqlConfig postgreSqlConfig { 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 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); diff --git a/presto-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java b/presto-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java index 2d2dda1e5eb43..b8eade5912fb1 100644 --- a/presto-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java +++ b/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 jsonDataType) .addRoundTrip(jsonDataType, "[]"); } + @Test + public void testUuid() + { + uuidTestCases(uuidDataType()) + .execute(getQueryRunner(), prestoCreateAsSelect("presto_test_uuid")); + } + + private DataTypeTest uuidTestCases(DataType 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 jsonbDataType() identity()); } + public static DataType uuidDataType() + { + return dataType( + "uuid", + UUID, + value -> "UUID " + formatStringLiteral(value.toString()), + identity()); + } + private static DataType byteaDataType() { return dataType( diff --git a/presto-tests/src/main/java/io/prestosql/tests/TestingPrestoClient.java b/presto-tests/src/main/java/io/prestosql/tests/TestingPrestoClient.java index 5a2f12cfbd95a..89fadb29cfcb6 100644 --- a/presto-tests/src/main/java/io/prestosql/tests/TestingPrestoClient.java +++ b/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; }