Java tutorial
/* * 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.verifier; import com.facebook.presto.sql.parser.SqlParser; import com.facebook.presto.sql.tree.CreateTable; import com.facebook.presto.sql.tree.CreateTableAsSelect; import com.facebook.presto.sql.tree.DropTable; import com.facebook.presto.sql.tree.Expression; import com.facebook.presto.sql.tree.FunctionCall; import com.facebook.presto.sql.tree.Identifier; import com.facebook.presto.sql.tree.Insert; import com.facebook.presto.sql.tree.LikeClause; import com.facebook.presto.sql.tree.LongLiteral; import com.facebook.presto.sql.tree.QualifiedName; import com.facebook.presto.sql.tree.QueryBody; import com.facebook.presto.sql.tree.QuerySpecification; import com.facebook.presto.sql.tree.Select; import com.facebook.presto.sql.tree.SelectItem; import com.facebook.presto.sql.tree.SingleColumn; import com.facebook.presto.sql.tree.Statement; import com.facebook.presto.sql.tree.Table; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.SimpleTimeLimiter; import com.google.common.util.concurrent.TimeLimiter; import io.airlift.units.Duration; import java.sql.Connection; import java.sql.DriverManager; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLClientInfoException; import java.sql.SQLException; import java.sql.Types; import java.util.ArrayList; import java.util.List; import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.concurrent.TimeUnit; import static com.facebook.presto.sql.SqlFormatter.formatSql; import static com.facebook.presto.sql.tree.LikeClause.PropertiesOption.INCLUDING; import static com.facebook.presto.verifier.PrestoVerifier.statementToQueryType; import static com.facebook.presto.verifier.QueryType.READ; import static java.lang.String.format; import static java.util.Objects.requireNonNull; public class QueryRewriter { private static final Set<Integer> APPROXIMATE_TYPES = ImmutableSet.of(Types.REAL, Types.FLOAT, Types.DOUBLE); private final SqlParser parser; private final String gatewayUrl; private final QualifiedName rewritePrefix; private final Optional<String> catalogOverride; private final Optional<String> schemaOverride; private final Optional<String> usernameOverride; private final Optional<String> passwordOverride; private final int doublePrecision; private final Duration timeout; public QueryRewriter(SqlParser parser, String gatewayUrl, QualifiedName rewritePrefix, Optional<String> catalogOverride, Optional<String> schemaOverride, Optional<String> usernameOverride, Optional<String> passwordOverride, int doublePrecision, Duration timeout) { this.parser = requireNonNull(parser, "parser is null"); this.gatewayUrl = requireNonNull(gatewayUrl, "gatewayUrl is null"); this.rewritePrefix = requireNonNull(rewritePrefix, "rewritePrefix is null"); this.catalogOverride = requireNonNull(catalogOverride, "catalogOverride is null"); this.schemaOverride = requireNonNull(schemaOverride, "schemaOverride is null"); this.usernameOverride = requireNonNull(usernameOverride, "usernameOverride is null"); this.passwordOverride = requireNonNull(passwordOverride, "passwordOverride is null"); this.doublePrecision = doublePrecision; this.timeout = requireNonNull(timeout, "timeout is null"); } public Query shadowQuery(Query query) throws QueryRewriteException, SQLException { if (statementToQueryType(parser, query.getQuery()) == READ) { return query; } if (!query.getPreQueries().isEmpty()) { throw new QueryRewriteException("Cannot rewrite queries that use pre-queries"); } if (!query.getPostQueries().isEmpty()) { throw new QueryRewriteException("Cannot rewrite queries that use post-queries"); } Statement statement = parser.createStatement(query.getQuery()); try (Connection connection = DriverManager.getConnection(gatewayUrl, usernameOverride.orElse(query.getUsername()), passwordOverride.orElse(query.getPassword()))) { trySetConnectionProperties(query, connection); if (statement instanceof CreateTableAsSelect) { return rewriteCreateTableAsSelect(connection, query, (CreateTableAsSelect) statement); } else if (statement instanceof Insert) { return rewriteInsertQuery(connection, query, (Insert) statement); } } throw new QueryRewriteException("Unsupported query type: " + statement.getClass()); } private Query rewriteCreateTableAsSelect(Connection connection, Query query, CreateTableAsSelect statement) throws SQLException, QueryRewriteException { QualifiedName temporaryTableName = generateTemporaryTableName(statement.getName()); Statement rewritten = new CreateTableAsSelect(temporaryTableName, statement.getQuery(), statement.isNotExists(), statement.getProperties(), statement.isWithData(), Optional.empty()); String createTableAsSql = formatSql(rewritten, Optional.empty()); String checksumSql = checksumSql(getColumns(connection, statement), temporaryTableName); String dropTableSql = dropTableSql(temporaryTableName); return new Query(query.getCatalog(), query.getSchema(), ImmutableList.of(createTableAsSql), checksumSql, ImmutableList.of(dropTableSql), query.getUsername(), query.getPassword(), query.getSessionProperties()); } private Query rewriteInsertQuery(Connection connection, Query query, Insert statement) throws SQLException, QueryRewriteException { QualifiedName temporaryTableName = generateTemporaryTableName(statement.getTarget()); Statement createTemporaryTable = new CreateTable(temporaryTableName, ImmutableList.of(new LikeClause(statement.getTarget(), Optional.of(INCLUDING))), true, ImmutableMap.of(), Optional.empty()); String createTemporaryTableSql = formatSql(createTemporaryTable, Optional.empty()); String insertSql = formatSql(new Insert(temporaryTableName, statement.getColumns(), statement.getQuery()), Optional.empty()); String checksumSql = checksumSql(getColumnsForTable(connection, query.getCatalog(), query.getSchema(), statement.getTarget().toString()), temporaryTableName); String dropTableSql = dropTableSql(temporaryTableName); return new Query(query.getCatalog(), query.getSchema(), ImmutableList.of(createTemporaryTableSql, insertSql), checksumSql, ImmutableList.of(dropTableSql), query.getUsername(), query.getPassword(), query.getSessionProperties()); } private QualifiedName generateTemporaryTableName(QualifiedName originalName) { List<String> parts = new ArrayList<>(); int originalSize = originalName.getOriginalParts().size(); int prefixSize = rewritePrefix.getOriginalParts().size(); if (originalSize > prefixSize) { parts.addAll(originalName.getOriginalParts().subList(0, originalSize - prefixSize)); } parts.addAll(rewritePrefix.getOriginalParts()); parts.set(parts.size() - 1, createTemporaryTableName()); return QualifiedName.of(parts); } private void trySetConnectionProperties(Query query, Connection connection) throws SQLException { // Required for jdbc drivers that do not implement all/some of these functions (eg. impala jdbc driver) // For these drivers, set the database default values in the query database try { connection.setClientInfo("ApplicationName", "verifier-rewrite"); connection.setCatalog(catalogOverride.orElse(query.getCatalog())); connection.setSchema(schemaOverride.orElse(query.getSchema())); } catch (SQLClientInfoException ignored) { // Do nothing } } private String createTemporaryTableName() { return rewritePrefix.getSuffix() + UUID.randomUUID().toString().replace("-", ""); } private List<Column> getColumnsForTable(Connection connection, String catalog, String schema, String table) throws SQLException { ResultSet columns = connection.getMetaData().getColumns(catalog, escapeLikeExpression(connection, schema), escapeLikeExpression(connection, table), null); ImmutableList.Builder<Column> columnBuilder = new ImmutableList.Builder<>(); while (columns.next()) { String name = columns.getString("COLUMN_NAME"); int type = columns.getInt("DATA_TYPE"); columnBuilder.add(new Column(name, APPROXIMATE_TYPES.contains(type))); } return columnBuilder.build(); } private List<Column> getColumns(Connection connection, CreateTableAsSelect createTableAsSelect) throws SQLException { com.facebook.presto.sql.tree.Query createSelectClause = createTableAsSelect.getQuery(); // Rewrite the query to select zero rows, so that we can get the column names and types QueryBody innerQuery = createSelectClause.getQueryBody(); com.facebook.presto.sql.tree.Query zeroRowsQuery; if (innerQuery instanceof QuerySpecification) { QuerySpecification querySpecification = (QuerySpecification) innerQuery; innerQuery = new QuerySpecification(querySpecification.getSelect(), querySpecification.getFrom(), querySpecification.getWhere(), querySpecification.getGroupBy(), querySpecification.getHaving(), querySpecification.getOrderBy(), Optional.of("0")); zeroRowsQuery = new com.facebook.presto.sql.tree.Query(createSelectClause.getWith(), innerQuery, Optional.empty(), Optional.empty()); } else { zeroRowsQuery = new com.facebook.presto.sql.tree.Query(createSelectClause.getWith(), innerQuery, Optional.empty(), Optional.of("0")); } ImmutableList.Builder<Column> columns = ImmutableList.builder(); try (java.sql.Statement jdbcStatement = connection.createStatement()) { TimeLimiter limiter = new SimpleTimeLimiter(); java.sql.Statement limitedStatement = limiter.newProxy(jdbcStatement, java.sql.Statement.class, timeout.toMillis(), TimeUnit.MILLISECONDS); try (ResultSet resultSet = limitedStatement.executeQuery(formatSql(zeroRowsQuery, Optional.empty()))) { ResultSetMetaData metaData = resultSet.getMetaData(); for (int i = 1; i <= metaData.getColumnCount(); i++) { String name = metaData.getColumnName(i); int type = metaData.getColumnType(i); columns.add(new Column(name, APPROXIMATE_TYPES.contains(type))); } } } return columns.build(); } private String checksumSql(List<Column> columns, QualifiedName table) throws SQLException, QueryRewriteException { if (columns.isEmpty()) { throw new QueryRewriteException("Table " + table + " has no columns"); } ImmutableList.Builder<SelectItem> selectItems = ImmutableList.builder(); for (Column column : columns) { Expression expression = new Identifier(column.getName()); if (column.isApproximateType()) { expression = new FunctionCall(QualifiedName.of("round"), ImmutableList.of(expression, new LongLiteral(Integer.toString(doublePrecision)))); } selectItems.add( new SingleColumn(new FunctionCall(QualifiedName.of("checksum"), ImmutableList.of(expression)))); } Select select = new Select(false, selectItems.build()); return formatSql(new QuerySpecification(select, Optional.of(new Table(table)), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()), Optional.empty()); } private static String dropTableSql(QualifiedName table) { return formatSql(new DropTable(table, true), Optional.empty()); } private static String escapeLikeExpression(Connection connection, String value) throws SQLException { String escapeString = connection.getMetaData().getSearchStringEscape(); return value.replace(escapeString, escapeString + escapeString).replace("_", escapeString + "_") .replace("%", escapeString + "%"); } public static class QueryRewriteException extends Exception { public QueryRewriteException(String messageFormat, Object... args) { super(format(messageFormat, args)); } } private static class Column { private final String name; private final boolean approximateType; private Column(String name, boolean approximateType) { this.name = name; this.approximateType = approximateType; } public String getName() { return name; } public boolean isApproximateType() { return approximateType; } } }