JdbcToEnumerableConverter.java
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to you 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 org.apache.calcite.adapter.jdbc;
import org.apache.calcite.DataContext;
import org.apache.calcite.adapter.enumerable.EnumerableRel;
import org.apache.calcite.adapter.enumerable.EnumerableRelImplementor;
import org.apache.calcite.adapter.enumerable.JavaRowFormat;
import org.apache.calcite.adapter.enumerable.PhysType;
import org.apache.calcite.adapter.enumerable.PhysTypeImpl;
import org.apache.calcite.adapter.enumerable.RexImpTable;
import org.apache.calcite.adapter.java.JavaTypeFactory;
import org.apache.calcite.config.CalciteSystemProperty;
import org.apache.calcite.linq4j.tree.BlockBuilder;
import org.apache.calcite.linq4j.tree.BlockStatement;
import org.apache.calcite.linq4j.tree.ConstantExpression;
import org.apache.calcite.linq4j.tree.Expression;
import org.apache.calcite.linq4j.tree.Expressions;
import org.apache.calcite.linq4j.tree.ParameterExpression;
import org.apache.calcite.linq4j.tree.Primitive;
import org.apache.calcite.linq4j.tree.UnaryExpression;
import org.apache.calcite.plan.ConventionTraitDef;
import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.plan.RelOptCost;
import org.apache.calcite.plan.RelOptPlanner;
import org.apache.calcite.plan.RelTraitSet;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.convert.ConverterImpl;
import org.apache.calcite.rel.metadata.RelMetadataQuery;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.runtime.Hook;
import org.apache.calcite.runtime.SqlFunctions;
import org.apache.calcite.schema.Schemas;
import org.apache.calcite.sql.SqlDialect;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.calcite.sql.util.SqlString;
import org.apache.calcite.util.BuiltInMethod;
import org.checkerframework.checker.nullness.qual.Nullable;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.List;
import java.util.TimeZone;
import java.util.stream.Collectors;
import javax.sql.DataSource;
import static org.apache.calcite.linq4j.Nullness.castNonNull;
import static java.util.Objects.requireNonNull;
/**
* Relational expression representing a scan of a table in a JDBC data source.
*/
public class JdbcToEnumerableConverter
extends ConverterImpl
implements EnumerableRel {
protected JdbcToEnumerableConverter(
RelOptCluster cluster,
RelTraitSet traits,
RelNode input) {
super(cluster, ConventionTraitDef.INSTANCE, traits, input);
}
@Override public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
return new JdbcToEnumerableConverter(
getCluster(), traitSet, sole(inputs));
}
@Override public @Nullable RelOptCost computeSelfCost(RelOptPlanner planner,
RelMetadataQuery mq) {
RelOptCost cost = super.computeSelfCost(planner, mq);
if (cost == null) {
return null;
}
return cost.multiplyBy(.1);
}
@Override public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
// Generate:
// ResultSetEnumerable.of(schema.getDataSource(), "select ...")
final BlockBuilder builder0 = new BlockBuilder(false);
final JdbcRel child = (JdbcRel) getInput();
final PhysType physType =
PhysTypeImpl.of(
implementor.getTypeFactory(), getRowType(),
pref.prefer(JavaRowFormat.CUSTOM));
final JdbcConvention jdbcConvention =
(JdbcConvention) requireNonNull(child.getConvention(),
() -> "child.getConvention() is null for " + child);
JdbcCorrelationDataContextBuilderImpl dataContextBuilder =
new JdbcCorrelationDataContextBuilderImpl(implementor, builder0, DataContext.ROOT);
SqlString sqlString = generateSql(jdbcConvention.dialect, dataContextBuilder);
String sql = sqlString.getSql();
if (CalciteSystemProperty.DEBUG.value()) {
System.out.println("[" + sql + "]");
}
Hook.QUERY_PLAN.run(sql);
final Expression sql_ =
builder0.append("sql", Expressions.constant(sql));
final int fieldCount = getRowType().getFieldCount();
BlockBuilder builder = new BlockBuilder();
final ParameterExpression resultSet_ =
Expressions.parameter(Modifier.FINAL, ResultSet.class,
builder.newName("resultSet"));
final SqlDialect.CalendarPolicy calendarPolicy =
jdbcConvention.dialect.getCalendarPolicy();
final Expression calendar_;
switch (calendarPolicy) {
case LOCAL:
calendar_ =
builder0.append("calendar",
Expressions.call(Calendar.class, "getInstance",
getTimeZoneExpression(implementor)));
break;
default:
calendar_ = null;
}
if (fieldCount == 0) {
// return (Object) null;
final ParameterExpression value_ =
Expressions.parameter(Object.class, builder.newName("value"));
builder.add(Expressions.declare(Modifier.FINAL, value_, null));
builder.add(Expressions.return_(null, value_));
} else if (fieldCount == 1) {
final ParameterExpression value_ =
Expressions.parameter(Object.class, builder.newName("value"));
builder.add(Expressions.declare(Modifier.FINAL, value_, null));
generateGet(implementor, physType, builder, resultSet_, 0, value_,
calendar_, calendarPolicy);
builder.add(Expressions.return_(null, value_));
} else {
final Expression values_ =
builder.append("values",
Expressions.newArrayBounds(Object.class, 1,
Expressions.constant(fieldCount)));
for (int i = 0; i < fieldCount; i++) {
generateGet(implementor, physType, builder, resultSet_, i,
Expressions.arrayIndex(values_, Expressions.constant(i)),
calendar_, calendarPolicy);
}
builder.add(
Expressions.return_(null, values_));
}
final ParameterExpression e_ =
Expressions.parameter(SQLException.class, builder.newName("e"));
BlockStatement valueBlock;
if (fieldCount == 0) {
// For some queries, we don't need to rely on a specific column value in the data source,
// we only need to know the number of rows in the table.
// For example: "select row_number() over () from dept"
// we can't push down the "row_number() over ()"
// So the generated code should be:
// public org.apache.calcite.linq4j.function.Function0 apply(
// final java.sql.ResultSet resultSet) {
// return new org.apache.calcite.linq4j.function.Function0() {
// public Object apply() {
// return (Object) null;
// }
// };
// }
valueBlock = builder.toBlock();
} else {
// public org.apache.calcite.linq4j.function.Function0 apply(
// final java.sql.ResultSet resultSet) {
// return new org.apache.calcite.linq4j.function.Function0() {
// public Object apply() {
// try {
// return resultSet.getString(1);
// } catch (java.sql.SQLException e) {
// throw new RuntimeException(
// e);
// }
// }
// };
// }
valueBlock =
Expressions.block(
Expressions.tryCatch(
builder.toBlock(),
Expressions.catch_(
e_,
Expressions.throw_(
Expressions.new_(
RuntimeException.class,
e_)))));
}
final Expression rowBuilderFactory_ =
builder0.append("rowBuilderFactory",
Expressions.lambda(
Expressions.block(
Expressions.return_(null,
Expressions.lambda(valueBlock))), resultSet_));
final Expression enumerable;
if (sqlString.getDynamicParameters() != null
&& !sqlString.getDynamicParameters().isEmpty()) {
final Expression preparedStatementConsumer_ =
builder0.append("preparedStatementConsumer",
Expressions.call(BuiltInMethod.CREATE_ENRICHER.method,
Expressions.newArrayInit(Integer.class, 1,
toIndexesTableExpression(sqlString)),
dataContextBuilder.build()));
enumerable =
builder0.append("enumerable",
Expressions.call(
BuiltInMethod.RESULT_SET_ENUMERABLE_OF_PREPARED.method,
Schemas.unwrap(jdbcConvention.expression, DataSource.class),
sql_,
rowBuilderFactory_,
preparedStatementConsumer_));
} else {
enumerable =
builder0.append("enumerable",
Expressions.call(
BuiltInMethod.RESULT_SET_ENUMERABLE_OF.method,
Schemas.unwrap(jdbcConvention.expression, DataSource.class),
sql_,
rowBuilderFactory_));
}
builder0.add(
Expressions.statement(
Expressions.call(enumerable,
BuiltInMethod.RESULT_SET_ENUMERABLE_SET_TIMEOUT.method,
DataContext.ROOT)));
builder0.add(
Expressions.return_(null, enumerable));
return implementor.result(physType, builder0.toBlock());
}
private static List<ConstantExpression> toIndexesTableExpression(SqlString sqlString) {
return requireNonNull(sqlString.getDynamicParameters(),
() -> "sqlString.getDynamicParameters() is null for " + sqlString).stream()
.map(Expressions::constant)
.collect(Collectors.toList());
}
private static UnaryExpression getTimeZoneExpression(
EnumerableRelImplementor implementor) {
return Expressions.convert_(
Expressions.call(
implementor.getRootExpression(),
"get",
Expressions.constant("timeZone")),
TimeZone.class);
}
private static void generateGet(EnumerableRelImplementor implementor,
PhysType physType, BlockBuilder builder, ParameterExpression resultSet_,
int i, Expression target, @Nullable Expression calendar_,
SqlDialect.CalendarPolicy calendarPolicy) {
final Primitive primitive = Primitive.ofBoxOr(physType.fieldClass(i));
final RelDataType fieldType =
physType.getRowType().getFieldList().get(i).getType();
final List<Expression> dateTimeArgs = new ArrayList<>();
dateTimeArgs.add(Expressions.constant(i + 1));
SqlTypeName sqlTypeName = fieldType.getSqlTypeName();
boolean offset = false;
switch (calendarPolicy) {
case LOCAL:
requireNonNull(calendar_, "calendar_");
dateTimeArgs.add(calendar_);
break;
case NULL:
// We don't specify a calendar at all, so we don't add an argument and
// instead use the version of the getXXX that doesn't take a Calendar
break;
case DIRECT:
sqlTypeName = SqlTypeName.ANY;
break;
case SHIFT:
switch (sqlTypeName) {
case TIMESTAMP:
case DATE:
offset = true;
break;
default:
break;
}
break;
default:
break;
}
final Expression source;
switch (sqlTypeName) {
case DATE:
case TIME:
case TIMESTAMP:
source =
Expressions.call(
getMethod(sqlTypeName, fieldType.isNullable(), offset),
Expressions.<Expression>list()
.append(
Expressions.call(resultSet_,
getMethod2(sqlTypeName), dateTimeArgs))
.appendIf(offset, getTimeZoneExpression(implementor)));
break;
case ARRAY:
final Expression x =
Expressions.convert_(
Expressions.call(resultSet_, jdbcGetMethod(primitive),
Expressions.constant(i + 1)),
java.sql.Array.class);
source = Expressions.call(BuiltInMethod.JDBC_ARRAY_TO_LIST.method, x);
break;
case NULL:
source = RexImpTable.NULL_EXPR;
break;
default:
source =
Expressions.call(resultSet_, jdbcGetMethod(primitive),
Expressions.constant(i + 1));
}
builder.add(
Expressions.statement(
Expressions.assign(
target, source)));
// [CALCITE-596] If primitive type columns contain null value, returns null
// object
if (primitive != null) {
builder.add(
Expressions.ifThen(
Expressions.call(resultSet_, "wasNull"),
Expressions.statement(
Expressions.assign(target,
Expressions.constant(null)))));
}
}
private static Method getMethod(SqlTypeName sqlTypeName, boolean nullable,
boolean offset) {
switch (sqlTypeName) {
case DATE:
return (nullable
? (offset
? BuiltInMethod.DATE_TO_INT_OPTIONAL_OFFSET
: BuiltInMethod.DATE_TO_INT_OPTIONAL)
: (offset
? BuiltInMethod.DATE_TO_INT_OFFSET
: BuiltInMethod.DATE_TO_INT)).method;
case TIME:
return (nullable
? BuiltInMethod.TIME_TO_INT_OPTIONAL
: BuiltInMethod.TIME_TO_INT).method;
case TIMESTAMP:
return (nullable
? (offset
? BuiltInMethod.TIMESTAMP_TO_LONG_OPTIONAL_OFFSET
: BuiltInMethod.TIMESTAMP_TO_LONG_OPTIONAL)
: (offset
? BuiltInMethod.TIMESTAMP_TO_LONG_OFFSET
: BuiltInMethod.TIMESTAMP_TO_LONG)).method;
default:
throw new AssertionError(sqlTypeName + ":" + nullable);
}
}
private static Method getMethod2(SqlTypeName sqlTypeName) {
switch (sqlTypeName) {
case DATE:
return BuiltInMethod.RESULT_SET_GET_DATE2.method;
case TIME:
return BuiltInMethod.RESULT_SET_GET_TIME2.method;
case TIMESTAMP:
return BuiltInMethod.RESULT_SET_GET_TIMESTAMP2.method;
default:
throw new AssertionError(sqlTypeName);
}
}
/** E,g, {@code jdbcGetMethod(int)} returns "getInt". */
private static String jdbcGetMethod(@Nullable Primitive primitive) {
return primitive == null
? "getObject"
: "get" + SqlFunctions.initcap(castNonNull(primitive.primitiveName));
}
private SqlString generateSql(SqlDialect dialect,
JdbcCorrelationDataContextBuilder dataContextBuilder) {
final JdbcImplementor jdbcImplementor =
new JdbcImplementor(dialect,
(JavaTypeFactory) getCluster().getTypeFactory(), dataContextBuilder);
final JdbcImplementor.Result result =
jdbcImplementor.visitRoot(this.getInput());
return result.asStatement().toSqlString(dialect);
}
}