com.alibaba.wasp.plan.parser.druid.DruidDQLParser.java Source code

Java tutorial

Introduction

Here is the source code for com.alibaba.wasp.plan.parser.druid.DruidDQLParser.java

Source

/**
 * 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 com.alibaba.wasp.plan.parser.druid;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLAggregateExpr;
import com.alibaba.druid.sql.ast.expr.SQLAllColumnExpr;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOperator;
import com.alibaba.druid.sql.ast.statement.SQLSelect;
import com.alibaba.druid.sql.ast.statement.SQLSelectItem;
import com.alibaba.druid.sql.ast.statement.SQLSelectQuery;
import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
import com.alibaba.druid.sql.ast.statement.SQLUnionQuery;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
import com.alibaba.wasp.EntityGroupLocation;
import com.alibaba.wasp.FConstants;
import com.alibaba.wasp.ZooKeeperConnectionException;
import com.alibaba.wasp.client.FConnection;
import com.alibaba.wasp.client.FConnectionManager;
import com.alibaba.wasp.meta.FTable;
import com.alibaba.wasp.meta.Field;
import com.alibaba.wasp.meta.Index;
import com.alibaba.wasp.meta.RowBuilder;
import com.alibaba.wasp.meta.StorageTableNameBuilder;
import com.alibaba.wasp.plan.AggregateQueryPlan;
import com.alibaba.wasp.plan.DQLPlan;
import com.alibaba.wasp.plan.GlobalQueryPlan;
import com.alibaba.wasp.plan.LocalQueryPlan;
import com.alibaba.wasp.plan.action.Action;
import com.alibaba.wasp.plan.action.ColumnStruct;
import com.alibaba.wasp.plan.action.GetAction;
import com.alibaba.wasp.plan.action.ScanAction;
import com.alibaba.wasp.plan.parser.AggregateInfo;
import com.alibaba.wasp.plan.parser.Condition;
import com.alibaba.wasp.plan.parser.Condition.ConditionType;
import com.alibaba.wasp.plan.parser.ParseContext;
import com.alibaba.wasp.plan.parser.QueryInfo;
import com.alibaba.wasp.plan.parser.UnsupportedException;
import com.alibaba.wasp.util.ParserUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.hbase.util.Pair;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;

/**
 * Use Druid (https://github.com/AlibabaTech/druid) to parse the sql and
 * generate QueryPlan
 * 
 */
public class DruidDQLParser extends DruidParser {

    private static final Log LOG = LogFactory.getLog(DruidDQLParser.class);

    private final FConnection connection;

    /**
     * @param conf
     * @throws com.alibaba.wasp.ZooKeeperConnectionException
     */
    public DruidDQLParser(Configuration conf) throws ZooKeeperConnectionException {
        this(conf, FConnectionManager.getConnection(conf));
    }

    /**
     * @param conf
     */
    public DruidDQLParser(Configuration conf, FConnection connection) {
        super();
        this.setConf(conf);
        this.connection = connection;
    }

    /**
     * @see com.alibaba.wasp.plan.parser.Parser#generatePlan(com.alibaba.wasp.plan.parser.ParseContext)
     */
    @Override
    public void generatePlan(ParseContext context) throws IOException {
        SQLStatement stmt = context.getStmt();
        MetaEventOperation metaEventOperation = new FMetaEventOperation(context.getTsr());
        if (stmt instanceof SQLSelectStatement) {
            // This is a select SQL
            getSelectPlan(context, (SQLSelectStatement) stmt, metaEventOperation);
        } else {
            throw new UnsupportedException("Unsupported SQLStatement " + stmt);
        }
    }

    /**
     * Process select Statement and generate QueryPlan
     *
     */
    private void getSelectPlan(ParseContext context, SQLSelectStatement sqlSelectStatement,
            MetaEventOperation metaEventOperation) throws IOException {
        SQLSelect select = sqlSelectStatement.getSelect();
        SQLSelectQuery sqlSelectQuery = select.getQuery();
        if (sqlSelectQuery instanceof MySqlSelectQueryBlock) {
            MySqlSelectQueryBlock sqlSelectQueryBlock = (MySqlSelectQueryBlock) sqlSelectQuery;
            // SELECT
            // FROM
            // WHERE

            // Parse The FROM clause
            String fTableName = parseFromClause(sqlSelectQueryBlock.getFrom());
            LOG.debug("SELECT SQL TableSource " + sqlSelectQueryBlock.getFrom());
            // check if table exists and get Table info(WTable)
            FTable table = metaEventOperation.checkAndGetTable(fTableName, false);

            LinkedHashSet<String> selectItem = null;
            AggregateInfo aggregateInfo = parseAggregateClause(sqlSelectQueryBlock.getSelectList(), table);
            if (aggregateInfo == null) {
                // Parse The SELECT clause
                if (sqlSelectQueryBlock.getSelectList().size() == 1
                        && sqlSelectQueryBlock.getSelectList().get(0).getExpr() instanceof SQLAllColumnExpr) {
                    // This is SELECT * clause
                    selectItem = parseFTable(table);
                } else {
                    selectItem = parseSelectClause(sqlSelectQueryBlock.getSelectList());
                }
            } else {
                selectItem = new LinkedHashSet<String>();
                if (aggregateInfo.getField() == null) {
                    //TODO
                }
                if (!aggregateInfo.getField().getName().equals("*")) {
                    selectItem.add(aggregateInfo.getField().getName());
                }
            }
            LOG.debug("SELECT SQL:Select columns " + sqlSelectQueryBlock.getSelectList());
            // check if table has this columns
            metaEventOperation.checkAndGetFields(table, selectItem);

            // Parse The WHERE clause
            SQLExpr where = sqlSelectQueryBlock.getWhere();
            LOG.debug("SELECT SQL:where " + where);
            QueryInfo actionInfo = parseWhereClause(table, metaEventOperation, where,
                    sqlSelectQueryBlock.isForUpdate());
            LOG.debug("ActionInfo " + actionInfo.toString());

            // Parse The Limit clause
            SQLExpr rowCount = null;
            if (sqlSelectQueryBlock.getLimit() != null) {
                rowCount = sqlSelectQueryBlock.getLimit().getRowCount();
            }
            int limit = -1;
            if (rowCount != null) {
                limit = convertToInt(rowCount);
            }

            // Convert to QueryPlan
            if (aggregateInfo == null) {
                convertToQueryPlan(table, context, actionInfo, metaEventOperation, selectItem, limit);
            } else {
                actionInfo.setType(QueryInfo.QueryType.AGGREGATE);
                actionInfo.setAggregateInfo(aggregateInfo);
                convertToQueryPlan(table, context, actionInfo, metaEventOperation);
            }
        } else if (sqlSelectQuery instanceof SQLUnionQuery) {
            throw new UnsupportedException("Union clause Unsupported");
        }
    }

    private void convertToQueryPlan(FTable table, ParseContext context, QueryInfo queryInfo,
            MetaEventOperation metaEventOperation) throws IOException {
        // should check the fields in where clause
        List<String> conditionFields = queryInfo.getAllConditionFieldName();
        // check if table has this columns
        metaEventOperation.checkAndGetFields(table, conditionFields);

        List<ColumnStruct> conditionColumns = buildAllConditionColumns(table, queryInfo);
        ScanAction scanAction = new ScanAction(table.getTableName());
        scanAction.setNotIndexConditionColumns(conditionColumns);

        DQLPlan qp = new AggregateQueryPlan(scanAction, table);
        qp.setQueryInfo(queryInfo);
        LOG.debug("QueryPlan " + qp);
        context.setPlan(qp);

    }

    /**
     * Convert to QueryPlan
     *
     * @param table
     * @param context
     * @param queryInfo
     * @param metaEventOperation
     * @param selectItem
     * @throws java.io.IOException
     */
    private void convertToQueryPlan(FTable table, ParseContext context, QueryInfo queryInfo,
            MetaEventOperation metaEventOperation, LinkedHashSet<String> selectItem, int limit) throws IOException {
        // should check the fields in where clause
        List<String> conditionColumns = queryInfo.getAllConditionFieldName();
        // check if table has this columns
        metaEventOperation.checkAndGetFields(table, conditionColumns);

        List<Pair<String, byte[]>> primaryKeyPairs = metaEventOperation.getPrimaryKeyPairList(table,
                queryInfo.getEqConditions(), queryInfo.getRangeConditions());

        if ((queryInfo.getEqConditions().size() + queryInfo.getRangeConditions().size()) == 0) {
            throw new UnsupportedException("Unsupported null condition.");
        }

        if (primaryKeyPairs == null || primaryKeyPairs.isEmpty()) {
            // conditions do not contain PK.
            queryInfo.setType(QueryInfo.QueryType.SCAN);
        } else {
            queryInfo.setType(QueryInfo.QueryType.GET);
            if (primaryKeyPairs.size() != conditionColumns.size()) {
                throw new UnsupportedException(
                        "When you have specified the pk, you'd better not to specify additional filter conditions.");
            }
        }

        Action action = null;
        DQLPlan qp = null;

        Condition entityGroupKeyCondition = queryInfo.getField(table.getEntityGroupKey().getName());
        if (queryInfo.getType() == QueryInfo.QueryType.GET) {
            byte[] primaryKey = RowBuilder.build().genRowkey(primaryKeyPairs);
            // check if the column is table's primary key.
            action = new GetAction(context.getReadModel(), table.getTableName(), primaryKey,
                    this.buildEntityColumnsForGet(table, metaEventOperation, selectItem));
            ((GetAction) action).setForUpdate(queryInfo.isForUpdate());
            if (context.isGenWholePlan()) {
                // get entityGroupLocation according to entity group key.
                EntityGroupLocation entityGroupLocation = this.connection.locateEntityGroup(
                        Bytes.toBytes(table.getTableName()),
                        DruidParser.convert(table.getColumn(entityGroupKeyCondition.getFieldName()),
                                entityGroupKeyCondition.getValue()));
                action.setEntityGroupLocation(entityGroupLocation);
            }
            qp = new LocalQueryPlan((GetAction) action);
            LOG.debug(QueryInfo.QueryType.GET + "  " + Bytes.toStringBinary(primaryKey) + " from "
                    + table.getTableName());
        } else if (queryInfo.getType() == QueryInfo.QueryType.SCAN) {
            Index index = metaEventOperation.checkAndGetIndex(table, queryInfo.getAllConditionFieldName());

            if (index == null) {
                throw new UnsupportedException("Don't get a Index!");
            }

            boolean isJustUseIndex = index.getIndexKeys().size() >= queryInfo.getAllConditionFieldName().size();
            Pair<byte[], byte[]> startKeyAndEndKey = metaEventOperation.getStartkeyAndEndkey(index, queryInfo);

            Pair<List<ColumnStruct>, List<ColumnStruct>> columnActionPair = this.buildEntityColumnsForScan(table,
                    index, metaEventOperation, selectItem);
            List<ColumnStruct> selectEntityColumns = columnActionPair.getFirst();
            List<ColumnStruct> selectStoringColumns = columnActionPair.getSecond();
            List<ColumnStruct> conditionNotInIndex = isJustUseIndex ? Collections.<ColumnStruct>emptyList()
                    : buildColumnsNotInIndex(table, index, queryInfo);

            // instance scan action.
            action = new ScanAction(context.getReadModel(), StorageTableNameBuilder.buildIndexTableName(index),
                    table.getTableName(), startKeyAndEndKey.getFirst(), startKeyAndEndKey.getSecond(),
                    selectEntityColumns);
            ((ScanAction) action).setStoringColumns(selectStoringColumns);
            ((ScanAction) action).setLimit(limit);
            ((ScanAction) action).setNotIndexConditionColumns(conditionNotInIndex);

            if (entityGroupKeyCondition != null && entityGroupKeyCondition.getType() == ConditionType.EQUAL) {
                if (context.isGenWholePlan()) {
                    EntityGroupLocation entityGroupLocation = this.connection.locateEntityGroup(
                            Bytes.toBytes(table.getTableName()),
                            DruidParser.convert(table.getColumn(entityGroupKeyCondition.getFieldName()),
                                    entityGroupKeyCondition.getValue()));
                    action.setEntityGroupLocation(entityGroupLocation);
                }
                qp = new LocalQueryPlan((ScanAction) action);
                LOG.debug(
                        QueryInfo.QueryType.SCAN + " startKey " + Bytes.toStringBinary(startKeyAndEndKey.getFirst())
                                + " endKey " + Bytes.toStringBinary(startKeyAndEndKey.getSecond()));
            } else {
                qp = new GlobalQueryPlan((ScanAction) action, table);
            }
        }
        qp.setQueryInfo(queryInfo);
        LOG.debug("QueryPlan " + qp);
        context.setPlan(qp);
    }

    private List<ColumnStruct> buildColumnsNotInIndex(FTable table, Index index, QueryInfo queryInfo)
            throws UnsupportedException {
        List<ColumnStruct> columns = new ArrayList<ColumnStruct>();
        LinkedHashMap<String, Field> fields = table.getColumns();
        for (String fieldname : queryInfo.getAllConditionFieldName()) {
            if (!index.getIndexKeys().containsKey(fieldname)) {
                Field field = fields.get(fieldname);
                if (field != null) {
                    ColumnStruct column = buildColumnStruct(table, queryInfo, field);
                    columns.add(column);
                }
            }
        }
        return columns;
    }

    private List<ColumnStruct> buildAllConditionColumns(FTable table, QueryInfo queryInfo)
            throws UnsupportedException {
        List<ColumnStruct> columns = new ArrayList<ColumnStruct>();
        LinkedHashMap<String, Field> fields = table.getColumns();
        for (String fieldname : queryInfo.getAllConditionFieldName()) {
            Field field = fields.get(fieldname);
            if (field != null) {
                ColumnStruct column = buildColumnStruct(table, queryInfo, field);
                columns.add(column);
            }
        }
        return columns;
    }

    private ColumnStruct buildColumnStruct(FTable table, QueryInfo queryInfo, Field field)
            throws UnsupportedException {
        Condition condition = queryInfo.getField(field.getName());
        ColumnStruct column = null;
        if (condition.getType() == ConditionType.EQUAL) {
            column = new ColumnStruct(table.getTableName(), field.getFamily(), field.getName(), field.getType(),
                    DruidParser.convert(field, condition.getValue()),
                    ParserUtils.parseSQLBinaryOperatorToIntValue(SQLBinaryOperator.Equality));
        } else {
            SQLBinaryOperator leftOperator = condition.getLeftOperator();
            SQLBinaryOperator rightOperator = condition.getRightOperator();
            column = new ColumnStruct(table.getTableName(), field.getFamily(), field.getName(), field.getType(),
                    DruidParser.convert(field, leftOperator != null ? condition.getLeft() : condition.getRight()),
                    ParserUtils
                            .parseSQLBinaryOperatorToIntValue(leftOperator != null ? leftOperator : rightOperator));
        }
        return column;
    }

    /**
     * Build the query field list that get operation requires.
     * 
     */
    private List<ColumnStruct> buildEntityColumnsForGet(FTable table, MetaEventOperation metaEventOperation,
            LinkedHashSet<String> selectItem) throws IOException {
        List<ColumnStruct> selectEntityColumns = new ArrayList<ColumnStruct>();
        for (String item : selectItem) {
            // Get the column's familyName
            String familyName = metaEventOperation.getColumnFamily(table.getTableName(), item);
            Field field = table.getColumn(item);
            ColumnStruct column = new ColumnStruct(table.getTableName(), familyName, item, field.getType());
            selectEntityColumns.add(column);
        }
        return selectEntityColumns;
    }

    /**
     * Separation storing query and entity query.
     * 
     */
    private Pair<List<ColumnStruct>, List<ColumnStruct>> buildEntityColumnsForScan(FTable table, Index index,
            MetaEventOperation metaEventOperation, LinkedHashSet<String> selectItem) throws IOException {
        List<ColumnStruct> selectStoringColumns = new ArrayList<ColumnStruct>();
        List<ColumnStruct> selectEntityColumns = new ArrayList<ColumnStruct>();
        String familyName = "";
        for (String item : selectItem) {
            // Get the column's familyName
            familyName = metaEventOperation.getColumnFamily(table.getTableName(), item);
            Field field = table.getColumn(item);
            if (index.getStoring().containsKey(item)) {
                ColumnStruct column = new ColumnStruct(table.getTableName(), FConstants.INDEX_STORING_FAMILY_STR,
                        item, field.getType());
                selectStoringColumns.add(column);
            } else {
                ColumnStruct column = new ColumnStruct(table.getTableName(), familyName, item, field.getType());
                selectEntityColumns.add(column);
            }
        }
        // if not only query storing columns, will be query the column with others.
        if (selectEntityColumns.size() > 0) {
            for (ColumnStruct selectStoringColumn : selectStoringColumns) {
                ColumnStruct column = new ColumnStruct(table.getTableName(), familyName,
                        selectStoringColumn.getColumnName(), selectStoringColumn.getDataType());
                selectEntityColumns.add(column);
            }
        }
        return new Pair<List<ColumnStruct>, List<ColumnStruct>>(selectEntityColumns, selectStoringColumns);
    }

    /**
     * Parse The WHERE Clause.
     * 
     */
    QueryInfo parseWhereClause(FTable table, MetaEventOperation metaEventOperation, SQLExpr where,
            boolean forUpdate) throws IOException {
        LinkedHashMap<String, Condition> conditions = new LinkedHashMap<String, Condition>();
        LinkedHashMap<String, Condition> ranges = new LinkedHashMap<String, Condition>();
        //List<Condition> ranges = new ArrayList<Condition>(5);
        ParserUtils.parse(where, conditions, ranges);
        return new QueryInfo(null, conditions, ranges, forUpdate);
    }

    /**
     * Parse The SQL SELECT Statement. Get which columns should be return
     * 
     */
    private LinkedHashSet<String> parseSelectClause(List<SQLSelectItem> select) throws UnsupportedException {
        LinkedHashSet<String> selectItem = new LinkedHashSet<String>(select.size());
        for (SQLSelectItem item : select) {
            SQLExpr expr = item.getExpr();
            if (expr instanceof SQLAggregateExpr) {

            }
            String columnName = parseColumn(expr);
            selectItem.add(columnName);
            LOG.debug(" SQLSelectItem " + columnName);
        }
        return selectItem;
    }

    private AggregateInfo parseAggregateClause(List<SQLSelectItem> select, FTable table)
            throws UnsupportedException {
        for (SQLSelectItem item : select) {
            SQLExpr expr = item.getExpr();
            if (expr instanceof SQLAggregateExpr) {
                SQLAggregateExpr aggregateExpr = (SQLAggregateExpr) expr;
                String method = aggregateExpr.getMethodName();
                String columnName = parseColumn(aggregateExpr.getArguments().get(0));
                LOG.debug(" SQLAggregatetItem column: " + columnName + " method: " + method);
                Field field = table.getColumns().get(columnName);
                if (method.equalsIgnoreCase("count")
                //&& (columnName.equals("*") || columnName.equals("1"))
                ) {
                    field = table.getEntityGroupKey();
                }
                return new AggregateInfo(AggregateInfo.getAggregateTypeByMethod(method), field);
            }
        }
        return null;
    }

    private LinkedHashSet<String> parseFTable(FTable ftable) throws UnsupportedException {
        LinkedHashSet<String> selectItem = new LinkedHashSet<String>(ftable.getColumns().size());
        for (Field item : ftable.getColumns().values()) {
            selectItem.add(item.getName());
            LOG.debug(" SQLSelectItem " + item.getName());
        }
        return selectItem;
    }
}