PostgreSqlClient.java
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.facebook.presto.plugin.postgresql;
import com.facebook.airlift.json.JsonObjectMapperProvider;
import com.facebook.presto.common.type.StandardTypes;
import com.facebook.presto.common.type.Type;
import com.facebook.presto.common.type.TypeManager;
import com.facebook.presto.common.type.TypeSignature;
import com.facebook.presto.plugin.jdbc.BaseJdbcClient;
import com.facebook.presto.plugin.jdbc.BaseJdbcConfig;
import com.facebook.presto.plugin.jdbc.DriverConnectionFactory;
import com.facebook.presto.plugin.jdbc.JdbcConnectorId;
import com.facebook.presto.plugin.jdbc.JdbcIdentity;
import com.facebook.presto.plugin.jdbc.JdbcTypeHandle;
import com.facebook.presto.plugin.jdbc.mapping.ReadMapping;
import com.facebook.presto.spi.ConnectorSession;
import com.facebook.presto.spi.ConnectorTableMetadata;
import com.facebook.presto.spi.PrestoException;
import com.facebook.presto.spi.SchemaTableName;
import com.fasterxml.jackson.core.JsonFactory;
import com.fasterxml.jackson.core.JsonFactoryBuilder;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.airlift.slice.DynamicSliceOutput;
import io.airlift.slice.Slice;
import io.airlift.slice.SliceOutput;
import org.postgresql.Driver;
import javax.inject.Inject;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
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;
import static com.facebook.presto.plugin.jdbc.mapping.ReadMapping.createSliceReadMapping;
import static com.facebook.presto.spi.StandardErrorCode.ALREADY_EXISTS;
import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
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.Long.reverseBytes;
import static java.lang.String.format;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.Locale.ENGLISH;
public class PostgreSqlClient
extends BaseJdbcClient
{
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);
@Inject
public PostgreSqlClient(JdbcConnectorId connectorId, BaseJdbcConfig config, TypeManager typeManager)
{
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
public PreparedStatement getPreparedStatement(ConnectorSession session, Connection connection, String sql)
throws SQLException
{
connection.setAutoCommit(false);
PreparedStatement statement = connection.prepareStatement(sql);
statement.setFetchSize(1000);
return statement;
}
@Override
protected ResultSet getTables(Connection connection, Optional<String> schemaName, Optional<String> tableName)
throws SQLException
{
DatabaseMetaData metadata = connection.getMetaData();
Optional<String> escape = Optional.ofNullable(metadata.getSearchStringEscape());
return metadata.getTables(
connection.getCatalog(),
escapeNamePattern(schemaName, escape).orElse(null),
escapeNamePattern(tableName, escape).orElse(null),
new String[] {"TABLE", "VIEW", "MATERIALIZED VIEW", "FOREIGN TABLE", "PARTITIONED TABLE"});
}
@Override
protected String toSqlType(Type type)
{
if (VARBINARY.equals(type)) {
return "bytea";
}
return super.toSqlType(type);
}
@Override
public Optional<ReadMapping> toPrestoType(ConnectorSession session, JdbcTypeHandle typeHandle)
{
if (typeHandle.getJdbcTypeName().equals("jsonb") || typeHandle.getJdbcTypeName().equals("json")) {
return Optional.of(jsonReadMapping());
}
else if (typeHandle.getJdbcTypeName().equals("uuid")) {
return Optional.of(uuidReadMapping());
}
return super.toPrestoType(session, typeHandle);
}
@Override
public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
{
try {
createTable(tableMetadata, session, tableMetadata.getTable().getTableName());
}
catch (SQLException e) {
if (DUPLICATE_TABLE_SQLSTATE.equals(e.getSQLState())) {
throw new PrestoException(ALREADY_EXISTS, e);
}
throw new PrestoException(JDBC_ERROR, e);
}
}
@Override
protected void renameTable(JdbcIdentity identity, String catalogName, SchemaTableName oldTable, SchemaTableName newTable)
{
// PostgreSQL does not allow qualifying the target of a rename
try (Connection connection = connectionFactory.openConnection(identity)) {
String sql = format(
"ALTER TABLE %s RENAME TO %s",
quoted(catalogName, oldTable.getSchemaName(), oldTable.getTableName()),
quoted(newTable.getTableName()));
execute(connection, sql);
}
catch (SQLException e) {
throw new PrestoException(JDBC_ERROR, e);
}
}
private ReadMapping jsonReadMapping()
{
return createSliceReadMapping(
jsonType,
(resultSet, columnIndex) -> jsonParse(utf8Slice(resultSet.getString(columnIndex))));
}
public static Slice jsonParse(Slice slice)
{
try (JsonParser parser = createJsonParser(JSON_FACTORY, slice)) {
byte[] in = slice.getBytes();
SliceOutput dynamicSliceOutput = new DynamicSliceOutput(in.length);
SORTED_MAPPER.writeValue((OutputStream) dynamicSliceOutput, SORTED_MAPPER.readValue(parser, Object.class));
// nextToken() returns null if the input is parsed correctly,
// but will throw an exception if there are trailing characters.
parser.nextToken();
return dynamicSliceOutput.slice();
}
catch (Exception e) {
throw new PrestoException(INVALID_FUNCTION_ARGUMENT, format("Cannot convert '%s' to JSON", slice.toStringUtf8()));
}
}
public static JsonParser createJsonParser(JsonFactory factory, Slice json)
throws IOException
{
// Jackson tries to detect the character encoding automatically when using InputStream
// so we pass an InputStreamReader instead.
return factory.createParser(new InputStreamReader(json.getInput(), UTF_8));
}
private ReadMapping uuidReadMapping()
{
return createSliceReadMapping(
uuidType,
(resultSet, columnIndex) -> uuidSlice((UUID) resultSet.getObject(columnIndex)));
}
private static Slice uuidSlice(UUID uuid)
{
return wrappedLongArray(
reverseBytes(uuid.getMostSignificantBits()),
reverseBytes(uuid.getLeastSignificantBits()));
}
@Override
public String normalizeIdentifier(ConnectorSession session, String identifier)
{
return caseSensitiveNameMatchingEnabled ? identifier : identifier.toLowerCase(ENGLISH);
}
}