Skip to content
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

Separate query rendering from statement creation in QueryBuilder #6688

Merged
merged 4 commits into from
Jan 22, 2021
Merged

Separate query rendering from statement creation in QueryBuilder #6688

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5d53a41
Simplify DruidJdbcClient.buildSql override
findepi Jan 22, 2021
7e4ba21
Separate query rendering from statement creation in QueryBuilder
findepi Jan 18, 2021
2105b31
Deprecate QueryBuilder.buildSql
findepi Jan 19, 2021
3fe3158
Remove redundant aliases in generated queries
findepi Jan 20, 2021
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
8 changes: 5 additions & 3 deletions plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -424,15 +424,17 @@ public Connection getConnection(ConnectorSession session, JdbcSplit split)
public PreparedStatement buildSql(ConnectorSession session, Connection connection, JdbcSplit split, JdbcTableHandle table, List<JdbcColumnHandle> columns)
throws SQLException
{
return new QueryBuilder(this).buildSql(
QueryBuilder queryBuilder = new QueryBuilder(this);
PreparedQuery preparedQuery = queryBuilder.prepareQuery(
session,
connection,
table.getRemoteTableName(),
table.getGroupingSets(),
columns,
table.getConstraint(),
split.getAdditionalPredicate(),
tryApplyLimit(table.getLimit()));
split.getAdditionalPredicate());
preparedQuery = preparedQuery.transformQuery(tryApplyLimit(table.getLimit()));
return queryBuilder.prepareStatement(session, connection, preparedQuery);
}

@Override
Expand Down
55 changes: 55 additions & 0 deletions plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/PreparedQuery.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
/*
* 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 io.trino.plugin.jdbc;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.collect.ImmutableList;

import java.util.List;
import java.util.function.Function;

import static java.util.Objects.requireNonNull;

public final class PreparedQuery
{
private final String query;
private final List<QueryParameter> parameters;

@JsonCreator
public PreparedQuery(String query, List<QueryParameter> parameters)
{
this.query = requireNonNull(query, "query is null");
this.parameters = ImmutableList.copyOf(requireNonNull(parameters, "parameters is null"));
}

@JsonProperty
public String getQuery()
{
return query;
}

@JsonProperty
public List<QueryParameter> getParameters()
{
return parameters;
}

public PreparedQuery transformQuery(Function<String, String> sqlFunction)
{
return new PreparedQuery(
sqlFunction.apply(query),
parameters);
}
}
127 changes: 74 additions & 53 deletions plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/QueryBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import io.trino.spi.predicate.Range;
import io.trino.spi.predicate.TupleDomain;
import io.trino.spi.predicate.ValueSet;
import io.trino.spi.type.Type;

import java.sql.Connection;
import java.sql.PreparedStatement;
Expand All @@ -32,6 +33,7 @@
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Consumer;
import java.util.function.Function;

import static com.google.common.base.Preconditions.checkArgument;
Expand All @@ -53,33 +55,15 @@ public class QueryBuilder

private final JdbcClient client;

private static class BoundValue
{
private final WriteFunction writeFunction;
private final Object value;

public BoundValue(WriteFunction writeFunction, Object value)
{
this.writeFunction = requireNonNull(writeFunction, "writeFunction is null");
this.value = requireNonNull(value, "value is null");
}

public WriteFunction getWriteFunction()
{
return writeFunction;
}

public Object getValue()
{
return value;
}
}

public QueryBuilder(JdbcClient client)
{
this.client = requireNonNull(client, "jdbcClient is null");
}

/**
* @deprecated Use #prepareSql and #prepareStatement instead.
*/
@Deprecated
public PreparedStatement buildSql(
ConnectorSession session,
Connection connection,
Expand All @@ -90,13 +74,34 @@ public PreparedStatement buildSql(
Optional<String> additionalPredicate,
Function<String, String> sqlFunction)
throws SQLException
{
PreparedQuery preparedQuery = prepareQuery(
session,
connection,
remoteTableName,
groupingSets,
columns,
tupleDomain,
additionalPredicate);
preparedQuery = preparedQuery.transformQuery(sqlFunction);
return prepareStatement(session, connection, preparedQuery);
}

public PreparedQuery prepareQuery(
ConnectorSession session,
Connection connection,
RemoteTableName remoteTableName,
Optional<List<List<JdbcColumnHandle>>> groupingSets,
List<JdbcColumnHandle> columns,
TupleDomain<ColumnHandle> tupleDomain,
Optional<String> additionalPredicate)
{
String sql = "SELECT " + getProjection(columns);
sql += " FROM " + getRelation(remoteTableName);

List<BoundValue> accumulator = new ArrayList<>();
ImmutableList.Builder<QueryParameter> accumulator = ImmutableList.builder();

List<String> clauses = toConjuncts(client, session, connection, tupleDomain, accumulator);
List<String> clauses = toConjuncts(client, session, connection, tupleDomain, accumulator::add);
if (additionalPredicate.isPresent()) {
clauses = ImmutableList.<String>builder()
.addAll(clauses)
Expand All @@ -109,16 +114,28 @@ public PreparedStatement buildSql(

sql += getGroupBy(groupingSets);

String query = sqlFunction.apply(sql);
log.debug("Preparing query: %s", query);
PreparedStatement statement = client.getPreparedStatement(connection, query);
return new PreparedQuery(sql, accumulator.build());
}

public PreparedStatement prepareStatement(
ConnectorSession session,
Connection connection,
PreparedQuery preparedQuery)
throws SQLException
{
log.debug("Preparing query: %s", preparedQuery.getQuery());
PreparedStatement statement = client.getPreparedStatement(connection, preparedQuery.getQuery());

for (int i = 0; i < accumulator.size(); i++) {
BoundValue boundValue = accumulator.get(i);
List<QueryParameter> parameters = preparedQuery.getParameters();
for (int i = 0; i < parameters.size(); i++) {
QueryParameter parameter = parameters.get(i);
int parameterIndex = i + 1;
WriteFunction writeFunction = boundValue.getWriteFunction();
WriteFunction writeFunction = getWriteFunction(session, connection, parameter.getJdbcType(), parameter.getType());
Class<?> javaType = writeFunction.getJavaType();
Object value = boundValue.getValue();
Object value = parameter.getValue()
// The value must be present, since QueryBuilder never creates null parameters. Values coming from Domain's ValueSet are non-null, and
// nullable domains are handled explicitly, with SQL syntax.
.orElseThrow(() -> new VerifyException("Value is missing"));
if (javaType == boolean.class) {
((BooleanWriteFunction) writeFunction).set(statement, parameterIndex, (boolean) value);
}
Expand Down Expand Up @@ -150,7 +167,14 @@ protected String getProjection(List<JdbcColumnHandle> columns)
return "1";
}
return columns.stream()
.map(jdbcColumnHandle -> format("%s AS %s", jdbcColumnHandle.toSqlExpression(client::quoted), client.quoted(jdbcColumnHandle.getColumnName())))
.map(jdbcColumnHandle -> {
String columnAlias = client.quoted(jdbcColumnHandle.getColumnName());
if (jdbcColumnHandle.getExpression().isEmpty()) {
return columnAlias;
}
String expression = jdbcColumnHandle.toSqlExpression(client::quoted);
return format("%s AS %s", expression, columnAlias);
})
.collect(joining(", "));
}

Expand All @@ -166,7 +190,7 @@ private List<String> toConjuncts(
ConnectorSession session,
Connection connection,
TupleDomain<ColumnHandle> tupleDomain,
List<BoundValue> accumulator)
Consumer<QueryParameter> accumulator)
{
if (tupleDomain.isNone()) {
return ImmutableList.of(ALWAYS_FALSE);
Expand All @@ -180,7 +204,7 @@ private List<String> toConjuncts(
return builder.build();
}

private String toPredicate(ConnectorSession session, Connection connection, JdbcColumnHandle column, Domain domain, List<BoundValue> accumulator)
private String toPredicate(ConnectorSession session, Connection connection, JdbcColumnHandle column, Domain domain, Consumer<QueryParameter> accumulator)
{
if (domain.getValues().isNone()) {
return domain.isNullAllowed() ? client.quoted(column.getColumnName()) + " IS NULL" : ALWAYS_FALSE;
Expand All @@ -197,7 +221,7 @@ private String toPredicate(ConnectorSession session, Connection connection, Jdbc
return format("(%s OR %s IS NULL)", predicate, client.quoted(column.getColumnName()));
}

private String toPredicate(ConnectorSession session, Connection connection, JdbcColumnHandle column, ValueSet valueSet, List<BoundValue> accumulator)
private String toPredicate(ConnectorSession session, Connection connection, JdbcColumnHandle column, ValueSet valueSet, Consumer<QueryParameter> accumulator)
{
checkArgument(!valueSet.isNone(), "none values should be handled earlier");

Expand All @@ -208,7 +232,9 @@ private String toPredicate(ConnectorSession session, Connection connection, Jdbc
}
}

WriteFunction writeFunction = getWriteFunction(session, connection, column);
JdbcTypeHandle jdbcType = column.getJdbcTypeHandle();
Type type = column.getColumnType();
WriteFunction writeFunction = getWriteFunction(session, connection, jdbcType, type);

List<String> disjuncts = new ArrayList<>();
List<Object> singleValues = new ArrayList<>();
Expand All @@ -222,10 +248,10 @@ private String toPredicate(ConnectorSession session, Connection connection, Jdbc
if (!range.getLow().isLowerUnbounded()) {
switch (range.getLow().getBound()) {
case ABOVE:
rangeConjuncts.add(toPredicate(column, writeFunction, ">", range.getLow().getValue(), accumulator));
rangeConjuncts.add(toPredicate(column, jdbcType, type, writeFunction, ">", range.getLow().getValue(), accumulator));
break;
case EXACTLY:
rangeConjuncts.add(toPredicate(column, writeFunction, ">=", range.getLow().getValue(), accumulator));
rangeConjuncts.add(toPredicate(column, jdbcType, type, writeFunction, ">=", range.getLow().getValue(), accumulator));
break;
case BELOW:
throw new IllegalArgumentException("Low marker should never use BELOW bound");
Expand All @@ -238,10 +264,10 @@ private String toPredicate(ConnectorSession session, Connection connection, Jdbc
case ABOVE:
throw new IllegalArgumentException("High marker should never use ABOVE bound");
case EXACTLY:
rangeConjuncts.add(toPredicate(column, writeFunction, "<=", range.getHigh().getValue(), accumulator));
rangeConjuncts.add(toPredicate(column, jdbcType, type, writeFunction, "<=", range.getHigh().getValue(), accumulator));
break;
case BELOW:
rangeConjuncts.add(toPredicate(column, writeFunction, "<", range.getHigh().getValue(), accumulator));
rangeConjuncts.add(toPredicate(column, jdbcType, type, writeFunction, "<", range.getHigh().getValue(), accumulator));
break;
default:
throw new AssertionError("Unhandled bound: " + range.getHigh().getBound());
Expand All @@ -260,11 +286,11 @@ private String toPredicate(ConnectorSession session, Connection connection, Jdbc

// Add back all of the possible single values either as an equality or an IN predicate
if (singleValues.size() == 1) {
disjuncts.add(toPredicate(column, writeFunction, "=", getOnlyElement(singleValues), accumulator));
disjuncts.add(toPredicate(column, jdbcType, type, writeFunction, "=", getOnlyElement(singleValues), accumulator));
}
else if (singleValues.size() > 1) {
for (Object value : singleValues) {
bindValue(writeFunction, value, accumulator);
accumulator.accept(new QueryParameter(jdbcType, type, Optional.of(value)));
}
String values = Joiner.on(",").join(nCopies(singleValues.size(), writeFunction.getBindExpression()));
disjuncts.add(client.quoted(column.getColumnName()) + " IN (" + values + ")");
Expand All @@ -277,18 +303,18 @@ else if (singleValues.size() > 1) {
return "(" + Joiner.on(" OR ").join(disjuncts) + ")";
}

private String toPredicate(JdbcColumnHandle column, WriteFunction writeFunction, String operator, Object value, List<BoundValue> accumulator)
private String toPredicate(JdbcColumnHandle column, JdbcTypeHandle jdbcType, Type type, WriteFunction writeFunction, String operator, Object value, Consumer<QueryParameter> accumulator)
{
bindValue(writeFunction, value, accumulator);
accumulator.accept(new QueryParameter(jdbcType, type, Optional.of(value)));
return format("%s %s %s", client.quoted(column.getColumnName()), operator, writeFunction.getBindExpression());
}

private WriteFunction getWriteFunction(ConnectorSession session, Connection connection, JdbcColumnHandle column)
private WriteFunction getWriteFunction(ConnectorSession session, Connection connection, JdbcTypeHandle jdbcType, Type type)
{
WriteFunction writeFunction = client.toColumnMapping(session, connection, column.getJdbcTypeHandle())
.orElseThrow(() -> new VerifyException(format("Unsupported type %s with handle %s for %s", column.getColumnType(), column.getJdbcTypeHandle(), column)))
WriteFunction writeFunction = client.toColumnMapping(session, connection, jdbcType)
.orElseThrow(() -> new VerifyException(format("Unsupported type %s with handle %s", type, jdbcType)))
.getWriteFunction();
verify(writeFunction.getJavaType() == column.getColumnType().getJavaType(), "Java type mismatch for %s: %s, %s", column, writeFunction, column.getColumnType());
verify(writeFunction.getJavaType() == type.getJavaType(), "Java type mismatch: %s, %s", writeFunction, type);
return writeFunction;
}

Expand Down Expand Up @@ -318,9 +344,4 @@ private String getGroupBy(Optional<List<List<JdbcColumnHandle>>> groupingSets)
.collect(joining(", ", "(", ")")))
.collect(joining(", ", "(", ")"));
}

private static void bindValue(WriteFunction writeFunction, Object value, List<BoundValue> accumulator)
{
accumulator.add(new BoundValue(writeFunction, value));
}
}
Loading