Skip to content

Add support for UUID in PostgreSQL connector #16875

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tdcmeehan merged 1 commit into from
Dec 20, 2022
Merged

Add support for UUID in PostgreSQL connector #16875

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/BaseJdbcClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import com.facebook.presto.common.type.CharType;
import com.facebook.presto.common.type.DecimalType;
import com.facebook.presto.common.type.Type;
import com.facebook.presto.common.type.UuidType;
import com.facebook.presto.common.type.VarcharType;
import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.spi.ColumnMetadata;
Expand Down Expand Up @@ -107,6 +108,7 @@ public class BaseJdbcClient
.put(TIME_WITH_TIME_ZONE, "time with timezone")
.put(TIMESTAMP, "timestamp")
.put(TIMESTAMP_WITH_TIME_ZONE, "timestamp with timezone")
.put(UuidType.UUID, "uuid")
Copy link
Contributor

@vivek-bharathan vivek-bharathan Nov 10, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

static import UUID

Copy link
Member Author

@agrawalreetika agrawalreetika Nov 11, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is UUID import from util package exist in this class, which is why I didn't use static import for this

.build();

protected final String connectorId;
Expand Down
6 changes: 6 additions & 0 deletions presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcPageSink.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import com.facebook.presto.common.block.Block;
import com.facebook.presto.common.type.DecimalType;
import com.facebook.presto.common.type.Type;
import com.facebook.presto.common.type.UuidType;
import com.facebook.presto.spi.ConnectorPageSink;
import com.facebook.presto.spi.ConnectorSession;
import com.facebook.presto.spi.PrestoException;
Expand Down Expand Up @@ -46,6 +47,7 @@
import static com.facebook.presto.common.type.RealType.REAL;
import static com.facebook.presto.common.type.SmallintType.SMALLINT;
import static com.facebook.presto.common.type.TinyintType.TINYINT;
import static com.facebook.presto.common.type.UuidType.prestoUuidToJavaUuid;
import static com.facebook.presto.common.type.VarbinaryType.VARBINARY;
import static com.facebook.presto.common.type.Varchars.isVarcharType;
import static com.facebook.presto.plugin.jdbc.JdbcErrorCode.JDBC_ERROR;
Expand Down Expand Up @@ -163,6 +165,10 @@ else if (DATE.equals(type)) {
long localMillis = getInstanceUTC().getZone().getMillisKeepLocal(DateTimeZone.getDefault(), utcMillis);
statement.setDate(parameter, new Date(localMillis));
}
else if (UuidType.UUID.equals(type)) {
Slice slice = type.getSlice(block, position);
statement.setObject(parameter, prestoUuidToJavaUuid(slice));
}
else {
throw new PrestoException(NOT_SUPPORTED, "Unsupported column type: " + type.getDisplayName());
}
Expand Down
2 changes: 2 additions & 0 deletions presto-client/src/main/java/com/facebook/presto/client/FixJsonDataUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@
import static com.facebook.presto.common.type.StandardTypes.TIMESTAMP_WITH_TIME_ZONE;
import static com.facebook.presto.common.type.StandardTypes.TIME_WITH_TIME_ZONE;
import static com.facebook.presto.common.type.StandardTypes.TINYINT;
import static com.facebook.presto.common.type.StandardTypes.UUID;
import static com.facebook.presto.common.type.StandardTypes.VARCHAR;
import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature;
import static com.google.common.base.Preconditions.checkArgument;
Expand Down Expand Up @@ -185,6 +186,7 @@ private static Object fixValue(TypeSignature signature, Object value)
case DECIMAL:
case CHAR:
case GEOMETRY:
case UUID:
return String.class.cast(value);
case BING_TILE:
// Bing tiles are serialized as strings when used as map keys,
Expand Down
5 changes: 1 addition & 4 deletions ...va/com/facebook/presto/type/UuidType.java → ...facebook/presto/common/type/UuidType.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,17 +11,14 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.facebook.presto.type;
package com.facebook.presto.common.type;

import com.facebook.presto.common.block.Block;
import com.facebook.presto.common.block.BlockBuilder;
import com.facebook.presto.common.block.BlockBuilderStatus;
import com.facebook.presto.common.block.Int128ArrayBlockBuilder;
import com.facebook.presto.common.block.PageBuilderStatus;
import com.facebook.presto.common.function.SqlFunctionProperties;
import com.facebook.presto.common.type.AbstractPrimitiveType;
import com.facebook.presto.common.type.FixedWidthType;
import com.facebook.presto.common.type.StandardTypes;
import io.airlift.slice.Slice;
import io.airlift.slice.XxHash64;

Expand Down
2 changes: 1 addition & 1 deletion ...in/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,7 @@
import static com.facebook.presto.common.type.TinyintType.TINYINT;
import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature;
import static com.facebook.presto.common.type.UnknownType.UNKNOWN;
import static com.facebook.presto.common.type.UuidType.UUID;
import static com.facebook.presto.common.type.VarbinaryType.VARBINARY;
import static com.facebook.presto.common.type.VarcharEnumParametricType.VARCHAR_ENUM;
import static com.facebook.presto.geospatial.SphericalGeographyType.SPHERICAL_GEOGRAPHY;
Expand Down Expand Up @@ -471,7 +472,6 @@
import static com.facebook.presto.type.Re2JRegexpType.RE2J_REGEXP;
import static com.facebook.presto.type.RowParametricType.ROW;
import static com.facebook.presto.type.TypeUtils.resolveTypes;
import static com.facebook.presto.type.UuidType.UUID;
import static com.facebook.presto.type.khyperloglog.KHyperLogLogType.K_HYPER_LOG_LOG;
import static com.facebook.presto.type.setdigest.SetDigestType.SET_DIGEST;
import static com.google.common.base.Preconditions.checkArgument;
Expand Down
6 changes: 3 additions & 3 deletions presto-main/src/main/java/com/facebook/presto/type/UuidOperators.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,10 @@
import static com.facebook.presto.common.function.OperatorType.LESS_THAN_OR_EQUAL;
import static com.facebook.presto.common.function.OperatorType.NOT_EQUAL;
import static com.facebook.presto.common.function.OperatorType.XX_HASH_64;
import static com.facebook.presto.common.type.UuidType.UUID;
import static com.facebook.presto.common.type.UuidType.javaUuidToPrestoUuid;
import static com.facebook.presto.common.type.UuidType.prestoUuidToJavaUuid;
import static com.facebook.presto.spi.StandardErrorCode.INVALID_CAST_ARGUMENT;
import static com.facebook.presto.type.UuidType.UUID;
import static com.facebook.presto.type.UuidType.javaUuidToPrestoUuid;
import static com.facebook.presto.type.UuidType.prestoUuidToJavaUuid;
import static io.airlift.slice.Slices.utf8Slice;
import static io.airlift.slice.Slices.wrappedBuffer;
import static java.util.UUID.randomUUID;
Expand Down
2 changes: 1 addition & 1 deletion presto-main/src/test/java/com/facebook/presto/type/TestUuidOperators.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,10 @@
import static com.facebook.presto.common.function.OperatorType.INDETERMINATE;
import static com.facebook.presto.common.type.BigintType.BIGINT;
import static com.facebook.presto.common.type.BooleanType.BOOLEAN;
import static com.facebook.presto.common.type.UuidType.UUID;
import static com.facebook.presto.common.type.VarbinaryType.VARBINARY;
import static com.facebook.presto.common.type.VarcharType.VARCHAR;
import static com.facebook.presto.type.UuidOperators.castFromVarcharToUuid;
import static com.facebook.presto.type.UuidType.UUID;
import static com.google.common.io.BaseEncoding.base16;
import static io.airlift.slice.Slices.utf8Slice;

Expand Down
2 changes: 1 addition & 1 deletion presto-main/src/test/java/com/facebook/presto/type/TestUuidType.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@
import io.airlift.slice.Slices;
import org.testng.annotations.Test;

import static com.facebook.presto.common.type.UuidType.UUID;
import static com.facebook.presto.type.UuidOperators.castFromVarcharToUuid;
import static com.facebook.presto.type.UuidType.UUID;
import static io.airlift.slice.Slices.utf8Slice;
import static org.testng.Assert.assertEquals;

Expand Down
20 changes: 20 additions & 0 deletions presto-postgresql/src/main/java/com/facebook/presto/plugin/postgresql/PostgreSqlClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Optional;
import java.util.UUID;

import static com.facebook.presto.common.type.VarbinaryType.VARBINARY;
import static com.facebook.presto.plugin.jdbc.JdbcErrorCode.JDBC_ERROR;
Expand All @@ -57,6 +58,7 @@
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.Slices.utf8Slice;
import static io.airlift.slice.Slices.wrappedLongArray;
import static java.lang.String.format;
import static java.nio.charset.StandardCharsets.UTF_8;

Expand All @@ -65,6 +67,7 @@ public class PostgreSqlClient
{
protected final Type jsonType;
private static final String DUPLICATE_TABLE_SQLSTATE = "42P07";
private final Type uuidType;

private static final JsonFactory JSON_FACTORY = new JsonFactoryBuilder().configure(CANONICALIZE_FIELD_NAMES, false).build();
private static final ObjectMapper SORTED_MAPPER = new JsonObjectMapperProvider().get().configure(ORDER_MAP_ENTRIES_BY_KEYS, true);
Expand All @@ -74,6 +77,7 @@ public PostgreSqlClient(JdbcConnectorId connectorId, BaseJdbcConfig config, Type
{
super(connectorId, config, "\"", new DriverConnectionFactory(new Driver(), config));
this.jsonType = typeManager.getType(new TypeSignature(StandardTypes.JSON));
this.uuidType = typeManager.getType(new TypeSignature(StandardTypes.UUID));
}

@Override
Expand Down Expand Up @@ -115,6 +119,10 @@ public Optional<ReadMapping> toPrestoType(ConnectorSession session, JdbcTypeHand
if (typeHandle.getJdbcTypeName().equals("jsonb") || typeHandle.getJdbcTypeName().equals("json")) {
return Optional.of(jsonColumnMapping());
}

else if (typeHandle.getJdbcTypeName().equals("uuid")) {
return Optional.of(uuidColumnMapping());
}
return super.toPrestoType(session, typeHandle);
}

Expand Down Expand Up @@ -178,4 +186,16 @@ public static JsonParser createJsonParser(JsonFactory factory, Slice json)
// so we pass an InputStreamReader instead.
return factory.createParser(new InputStreamReader(json.getInput(), UTF_8));
}

private ReadMapping uuidColumnMapping()
{
return ReadMapping.sliceReadMapping(
uuidType,
(resultSet, columnIndex) -> uuidSlice((UUID) resultSet.getObject(columnIndex)));
}

private static Slice uuidSlice(UUID uuid)
{
return wrappedLongArray(uuid.getMostSignificantBits(), uuid.getLeastSignificantBits());
}
}
26 changes: 26 additions & 0 deletions ...gresql/src/test/java/com/facebook/presto/plugin/postgresql/TestPostgreSqlTypeMapping.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@

import static com.facebook.presto.common.type.JsonType.JSON;
import static com.facebook.presto.common.type.TimeZoneKey.UTC_KEY;
import static com.facebook.presto.common.type.UuidType.UUID;
import static com.facebook.presto.common.type.VarbinaryType.VARBINARY;
import static com.facebook.presto.plugin.postgresql.PostgreSqlQueryRunner.createPostgreSqlQueryRunner;
import static com.facebook.presto.tests.datatype.DataType.bigintDataType;
Expand Down Expand Up @@ -329,6 +330,31 @@ public static String formatStringLiteral(String value)
return "'" + value.replace("'", "''") + "'";
}

@Test
public void testUuid()
{
uuidTestCases(uuidDataType())
.execute(getQueryRunner(), postgresCreateAndInsert("tpch.postgresql_test_uuid"));
uuidTestCases(uuidDataType())
.execute(getQueryRunner(), prestoCreateAsSelect("tpch.presto_test_uuid"));
}

private DataTypeTest uuidTestCases(DataType<String> uuidDataType)
{
return DataTypeTest.create()
.addRoundTrip(uuidDataType, "00000000-0000-0000-0000-000000000000")
.addRoundTrip(uuidDataType, "71f9206a-75aa-4005-9ab6-4525c6fdec99");
}

private static DataType<String> uuidDataType()
{
return dataType(
"uuid",
UUID,
value -> "UUID " + formatStringLiteral(value),
Function.identity());
}

private void testUnsupportedDataType(String databaseDataType)
{
JdbcSqlExecutor jdbcSqlExecutor = new JdbcSqlExecutor(postgreSqlServer.getJdbcUrl());
Expand Down
4 changes: 4 additions & 0 deletions presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import com.facebook.presto.common.type.SqlTimestampWithTimeZone;
import com.facebook.presto.common.type.Type;
import com.facebook.presto.common.type.TypeWithName;
import com.facebook.presto.common.type.UuidType;
import com.facebook.presto.common.type.VarcharEnumType;
import com.facebook.presto.common.type.VarcharType;
import com.facebook.presto.server.testing.TestingPrestoServer;
Expand Down Expand Up @@ -263,6 +264,9 @@ else if (type instanceof DecimalType) {
else if (type instanceof JsonType) {
return value;
}
else if (type instanceof UuidType) {
return value;
}
else if (type instanceof VarcharEnumType) {
return value;
}
Expand Down