/*
* 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.drill.exec.planner.sql.handlers;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.rex.RexBuilder;
import org.apache.calcite.rex.RexInputRef;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.schema.Table;
import org.apache.calcite.sql.SqlNodeList;
import org.apache.calcite.sql.SqlWriter;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.tools.RelConversionException;
import org.apache.drill.common.exceptions.DrillRuntimeException;
import org.apache.drill.common.exceptions.UserException;
import org.apache.drill.exec.planner.StarColumnHelper;
import org.apache.drill.exec.planner.common.DrillRelOptUtil;
import org.apache.drill.exec.store.AbstractSchema;
import org.apache.calcite.tools.ValidationException;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.plan.RelOptUtil;
import org.apache.calcite.rel.type.RelDataType;
import java.io.IOException;
import java.util.AbstractList;
import java.util.HashSet;
import java.util.List;
public class SqlHandlerUtil {
private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(SqlHandlerUtil.class);
/**
* Resolve final RelNode of the new table (or view) for given table field list and new table definition.
*
* @param isNewTableView Is the new table created a view? This doesn't affect the functionality, but it helps format
* better error messages.
* @param tableFieldNames List of fields specified in new table/view field list. These are the fields given just after
* new table name.
* Ex. CREATE TABLE newTblName(col1, medianOfCol2, avgOfCol3) AS
* SELECT col1, median(col2), avg(col3) FROM sourcetbl GROUP BY col1;
* @throws ValidationException If table's fields list and field list specified in table definition are not valid.
* @throws RelConversionException If failed to convert the table definition into a RelNode.
*/
public static RelNode resolveNewTableRel(boolean isNewTableView, List<String> tableFieldNames,
RelDataType validatedRowtype, RelNode queryRelNode) throws ValidationException, RelConversionException {
// TypedSqlNode validatedSqlNodeWithType = planner.validateAndGetType(newTableQueryDef);
// Get the row type of view definition query.
// Reason for getting the row type from validated SqlNode than RelNode is because SqlNode -> RelNode involves
// renaming duplicate fields which is not desired when creating a view or table.
// For ex: SELECT region_id, region_id FROM cp.`region.json` LIMIT 1 returns
// +------------+------------+
// | region_id | region_id0 |
// +------------+------------+
// | 0 | 0 |
// +------------+------------+
// which is not desired when creating new views or tables.
// final RelDataType queryRowType = validatedRowtype;
if (tableFieldNames.size() > 0) {
// Field count should match.
if (tableFieldNames.size() != validatedRowtype.getFieldCount()) {
final String tblType = isNewTableView ? "view" : "table";
throw UserException.validationError()
.message("%s's field list and the %s's query field list have different counts.", tblType, tblType)
.build(logger);
}
// CTAS's query field list shouldn't have "*" when table's field list is specified.
for (String field : validatedRowtype.getFieldNames()) {
if (field.equals("*")) {
final String tblType = isNewTableView ? "view" : "table";
throw UserException.validationError()
.message("%s's query field list has a '*', which is invalid when %s's field list is specified.",
tblType, tblType)
.build(logger);
}
}
// validate the given field names to make sure there are no duplicates
ensureNoDuplicateColumnNames(tableFieldNames);
// CTAS statement has table field list (ex. below), add a project rel to rename the query fields.
// Ex. CREATE TABLE tblname(col1, medianOfCol2, avgOfCol3) AS
// SELECT col1, median(col2), avg(col3) FROM sourcetbl GROUP BY col1 ;
// Similary for CREATE VIEW.
return DrillRelOptUtil.createRename(queryRelNode, tableFieldNames);
}
// As the column names of the view are derived from SELECT query, make sure the query has no duplicate column names
ensureNoDuplicateColumnNames(validatedRowtype.getFieldNames());
return queryRelNode;
}
private static void ensureNoDuplicateColumnNames(List<String> fieldNames) throws ValidationException {
final HashSet<String> fieldHashSet = Sets.newHashSetWithExpectedSize(fieldNames.size());
for(String field : fieldNames) {
if (fieldHashSet.contains(field.toLowerCase())) {
throw new ValidationException(String.format("Duplicate column name [%s]", field));
}
fieldHashSet.add(field.toLowerCase());
}
}
/**
* Resolve the partition columns specified in "PARTITION BY" clause of CTAS statement.
*
* A partition column is resolved, either (1) the same column appear in the select list of CTAS
* or (2) CTAS has a * in select list.
*
* In the second case, a PROJECT with ITEM expression would be created and returned.
* Throw validation error if a partition column is not resolved correctly.
*
* @param input : the RelNode represents the select statement in CTAS.
* @param partitionColumns : the list of partition columns.
* @return : 1) the original RelNode input, if all partition columns are in select list of CTAS
* 2) a New Project, if a partition column is resolved to * column in select list
* 3) validation error, if partition column is not resolved.
*/
public static RelNode qualifyPartitionCol(RelNode input, List<String> partitionColumns) {
final RelDataType inputRowType = input.getRowType();
final List<RexNode> colRefStarExprs = Lists.newArrayList();
final List<String> colRefStarNames = Lists.newArrayList();
final RexBuilder builder = input.getCluster().getRexBuilder();
final int originalFieldSize = inputRowType.getFieldCount();
for (final String col : partitionColumns) {
final RelDataTypeField field = inputRowType.getField(col, false, false);
if (field == null) {
throw UserException.validationError()
.message("Partition column %s is not in the SELECT list of CTAS!", col)
.build(logger);
} else {
if (field.getName().startsWith(StarColumnHelper.STAR_COLUMN)) {
colRefStarNames.add(col);
final List<RexNode> operands = Lists.newArrayList();
operands.add(new RexInputRef(field.getIndex(), field.getType()));
operands.add(builder.makeLiteral(col));
final RexNode item = builder.makeCall(SqlStdOperatorTable.ITEM, operands);
colRefStarExprs.add(item);
}
}
}
if (colRefStarExprs.isEmpty()) {
return input;
} else {
final List<String> names =
new AbstractList<String>() {
@Override
public String get(int index) {
if (index < originalFieldSize) {
return inputRowType.getFieldNames().get(index);
} else {
return colRefStarNames.get(index - originalFieldSize);
}
}
@Override
public int size() {
return originalFieldSize + colRefStarExprs.size();
}
};
final List<RexNode> refs =
new AbstractList<RexNode>() {
public int size() {
return originalFieldSize + colRefStarExprs.size();
}
public RexNode get(int index) {
if (index < originalFieldSize) {
return RexInputRef.of(index, inputRowType.getFieldList());
} else {
return colRefStarExprs.get(index - originalFieldSize);
}
}
};
return RelOptUtil.createProject(input, refs, names, false);
}
}
public static Table getTableFromSchema(AbstractSchema drillSchema, String tblName) {
try {
return drillSchema.getTable(tblName);
} catch (Exception e) {
// TODO: Move to better exception types.
throw new DrillRuntimeException(
String.format("Failure while trying to check if a table or view with given name [%s] already exists " +
"in schema [%s]: %s", tblName, drillSchema.getFullSchemaName(), e.getMessage()), e);
}
}
public static void unparseSqlNodeList(SqlWriter writer, int leftPrec, int rightPrec, SqlNodeList fieldList) {
writer.keyword("(");
fieldList.get(0).unparse(writer, leftPrec, rightPrec);
for (int i = 1; i<fieldList.size(); i++) {
writer.keyword(",");
fieldList.get(i).unparse(writer, leftPrec, rightPrec);
}
writer.keyword(")");
}
/**
* Drops table from schema.
* If drop has failed makes concurrency check: checks if table still exists.
* If table exists, throws {@link @UserException} since drop was unsuccessful,
* otherwise assumes that other user had dropped the view and exists without error.
*
* @param drillSchema drill schema
* @param tableName table name
*/
public static void dropTableFromSchema(AbstractSchema drillSchema, String tableName) {
try {
drillSchema.dropTable(tableName);
} catch (Exception e) {
if (SqlHandlerUtil.getTableFromSchema(drillSchema, tableName) != null) {
throw e;
}
}
}
/**
* Drops view from schema.
* If drop has failed makes concurrency check: checks if view still exists.
* If view exists, throws {@link @UserException} since drop was unsuccessful,
* otherwise assumes that other user had dropped the view and exists without error.
*
* @param drillSchema drill schema
* @param viewName view name
*/
public static void dropViewFromSchema(AbstractSchema drillSchema, String viewName) throws IOException {
try {
drillSchema.dropView(viewName);
} catch (Exception e) {
if (SqlHandlerUtil.getTableFromSchema(drillSchema, viewName) != null) {
throw e;
}
}
}
}