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

Java tutorial

Introduction

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

Source

/**
 * Copyright 2010 The Apache Software Foundation
 *
 * 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.statement.SQLDeleteStatement;
import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
import com.alibaba.druid.sql.ast.statement.SQLInsertStatement.ValuesClause;
import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem;
import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
import com.alibaba.wasp.EntityGroupLocation;
import com.alibaba.wasp.NotMatchPrimaryKeyException;
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.RowBuilder;
import com.alibaba.wasp.plan.DeletePlan;
import com.alibaba.wasp.plan.InsertPlan;
import com.alibaba.wasp.plan.UpdatePlan;
import com.alibaba.wasp.plan.action.ColumnStruct;
import com.alibaba.wasp.plan.action.DeleteAction;
import com.alibaba.wasp.plan.action.InsertAction;
import com.alibaba.wasp.plan.action.UpdateAction;
import com.alibaba.wasp.plan.parser.Condition;
import com.alibaba.wasp.plan.parser.ParseContext;
import com.alibaba.wasp.plan.parser.UnsupportedException;
import com.alibaba.wasp.plan.parser.druid.dialect.WaspSqlInsertStatement;
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.Arrays;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map.Entry;
import java.util.Set;

/**
 * Use Druid (https://github.com/AlibabaTech/druid) to parse the sql and
 * generate QueryPlan
 * 
 */
public class DruidDMLParser extends DruidParser {
    private static final Log LOG = LogFactory.getLog(DruidDMLParser.class);

    private final FConnection connection;

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

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

    /**
     * Parse sql and generate QueryPlan SELECT,UPDATE,INSERT,DELETE
     */
    @Override
    public void generatePlan(ParseContext context) throws IOException {
        SQLStatement stmt = context.getStmt();
        MetaEventOperation metaEventOperation = new FMetaEventOperation(context.getTsr());
        if (stmt instanceof SQLUpdateStatement) {
            // This is a Update SQL
            getUpdatePlan(context, (SQLUpdateStatement) stmt, metaEventOperation);
        } else if (stmt instanceof SQLInsertStatement) {
            // This is a Insert SQL
            getInsertPlan(context, (SQLInsertStatement) stmt, metaEventOperation);
        } else if (stmt instanceof SQLDeleteStatement) {
            // This is a Delete SQL
            getDeletePlan(context, (SQLDeleteStatement) stmt, metaEventOperation);
        }

        else {
            throw new UnsupportedException("Unsupported SQLStatement " + stmt);
        }
    }

    /**
     * Process Insert Statement and generate QueryPlan
     * 
     */
    private void getInsertPlan(ParseContext context, SQLInsertStatement sqlInsertStatement,
            MetaEventOperation metaEventOperation) throws IOException {

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

        // Parse The columns
        LinkedHashSet<String> insertColumns = parseInsertColumns(sqlInsertStatement.getColumns());
        LOG.debug("INSERT SQL Insert columns " + sqlInsertStatement.getColumns());
        // check if table has this columns
        metaEventOperation.checkAndGetFields(table, insertColumns);
        metaEventOperation.checkRequiredFields(table, insertColumns);

        // Before insert a row , it should not be exists, we do not check exists
        // here but check it in execution.
        List<InsertAction> actions = new ArrayList<InsertAction>();
        if (sqlInsertStatement instanceof WaspSqlInsertStatement) {
            // For MySQL, INSERT statements that use VALUES syntax can insert multiple
            // rows. see http://dev.mysql.com/doc/refman/5.5/en/insert.html
            List<ValuesClause> valuesList = ((WaspSqlInsertStatement) sqlInsertStatement).getValuesList();
            if (valuesList.size() > 1) {
                throw new UnsupportedException("Insert multi value unsupported now");
            }
            LOG.debug("INSERT SQL Insert Values List " + valuesList);
            for (ValuesClause values : valuesList) {
                actions.add(genInsertAction(metaEventOperation, table, insertColumns, values));
            }
        } else if (sqlInsertStatement instanceof SQLInsertStatement) {
            // INSERT statements that use VALUES syntax insert one row
            ValuesClause values = sqlInsertStatement.getValues();
            LOG.debug("INSERT SQL Insert Values " + values);
            actions.add(genInsertAction(metaEventOperation, table, insertColumns, values));
        }

        if (context.isGenWholePlan()) {
            for (InsertAction insertAction : actions) {
                // Get entityGroupLocation according to entity group key
                EntityGroupLocation entityGroupLocation = this.connection.locateEntityGroup(
                        Bytes.toBytes(table.getTableName()), insertAction.getValueOfEntityGroupKey());
                insertAction.setEntityGroupLocation(entityGroupLocation);
            }
        }
        InsertPlan insertPlan = new InsertPlan(actions);
        context.setPlan(insertPlan);
        LOG.debug("InsertPlan " + insertPlan.toString());
    }

    private InsertAction genInsertAction(MetaEventOperation metaEventOperation, FTable table,
            LinkedHashSet<String> insertColumns, ValuesClause values) throws IOException {
        Pair<List<Pair<String, byte[]>>, List<ColumnStruct>> pair = buildFieldsPair(metaEventOperation, table,
                insertColumns, values);
        byte[] primaryKey = RowBuilder.build().genRowkey(pair.getFirst());
        return new InsertAction(table.getTableName(), pair.getFirst().get(0).getSecond(), primaryKey,
                pair.getSecond());
    }

    /**
     * Process Delete Statement and generate QueryPlan
     * 
     */
    private void getDeletePlan(ParseContext context, SQLDeleteStatement sqlDeleteStatement,
            MetaEventOperation metaEventOperation) throws IOException {
        // DELETE FROM users WHERE id = 123;

        // Parse The FROM clause
        String wtableName = parseFromClause(sqlDeleteStatement.getTableSource());
        LOG.debug("UPDATE SQL From clause " + sqlDeleteStatement.getTableSource());
        // check if table exists and get Table info
        FTable table = metaEventOperation.checkAndGetTable(wtableName, false);

        // Parse The WHERE clause
        SQLExpr where = sqlDeleteStatement.getWhere();
        LOG.debug("UPDATE SQL where " + where);
        LinkedHashMap<String, Condition> eqConditions = new LinkedHashMap<String, Condition>();
        LinkedHashMap<String, Condition> ranges = new LinkedHashMap<String, Condition>();
        ParserUtils.parse(where, eqConditions, ranges);
        if (ranges.size() > 0) {
            throw new UnsupportedException("RANGE is not supported!");
        }

        // check if table has this columns
        metaEventOperation.checkAndGetFields(table, eqConditions.keySet());

        List<Pair<String, byte[]>> primaryKeyPairs = metaEventOperation.getPrimaryKeyPairList(table, eqConditions,
                null);
        if (primaryKeyPairs == null) {
            throw new NotMatchPrimaryKeyException("Not match primary key.");
        }

        byte[] primayKey = RowBuilder.build().genRowkey(primaryKeyPairs);
        DeleteAction action = new DeleteAction(wtableName, primayKey);
        if (context.isGenWholePlan()) {
            Condition entityGroupKeyCondition = ParserUtils.getCondition(table.getEntityGroupKey().getName(),
                    eqConditions);
            // 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);
        }
        List<DeleteAction> actions = new ArrayList<DeleteAction>();
        actions.add(action);
        DeletePlan deletePlan = new DeletePlan(actions);
        context.setPlan(deletePlan);
        LOG.debug("DeletePlan " + deletePlan.toString());
    }

    /**
     * Process UPDATE Statement and generate UpdatePlan
     * 
     */
    private void getUpdatePlan(ParseContext context, SQLUpdateStatement sqlUpdateStatement,
            MetaEventOperation metaEventOperation) throws IOException {
        // UPDATE users SET age = 24, name = 'Mike' WHERE id = 123;

        // Parse The FROM clause
        String fTableName = parseFromClause(sqlUpdateStatement.getTableSource());
        LOG.debug("UPDATE SQL From clause " + sqlUpdateStatement.getTableSource());
        // check if table exists and get Table info
        FTable table = metaEventOperation.checkAndGetTable(fTableName, false);

        // Parse The WHERE clause
        SQLExpr where = sqlUpdateStatement.getWhere();
        LOG.debug("UPDATE SQL where " + where);

        LinkedHashMap<String, Condition> conditions = new LinkedHashMap<String, Condition>();
        LinkedHashMap<String, Condition> ranges = new LinkedHashMap<String, Condition>();
        ParserUtils.parse(where, conditions, ranges);
        if (ranges.size() > 0) {
            throw new UnsupportedException("RANGE is not supported by update operation!");
        }

        Set<String> conditionColumns = ParserUtils.getColumns(conditions);
        // check if table has this columns
        metaEventOperation.checkAndGetFields(table, conditionColumns);
        // check if where clause is primary keys
        metaEventOperation.checkIsPrimaryKey(table, conditionColumns);

        List<Pair<String, byte[]>> primaryKeyPairs = metaEventOperation.getPrimaryKeyPairList(table, conditions,
                null);
        if (primaryKeyPairs == null) {
            throw new IOException("Not match primary key.");
        }

        byte[] primayKey = RowBuilder.build().genRowkey(primaryKeyPairs);

        UpdateAction action = new UpdateAction(fTableName, primayKey);
        // Parse Update Item
        List<SQLUpdateSetItem> updateItems = sqlUpdateStatement.getItems();
        for (SQLUpdateSetItem updateItem : updateItems) {
            String columnName = parseColumn(updateItem.getColumn());
            // check this FTable has the column and not pk
            metaEventOperation.checkFieldNotInPrimaryKeys(table, columnName);
            Field field = table.getColumn(columnName);
            // Check the input is the same as DataType
            checkType(field, updateItem.getValue());
            byte[] value = convert(field, updateItem.getValue());
            String familyName = metaEventOperation.getColumnFamily(fTableName, columnName);
            action.addEntityColumn(fTableName, familyName, columnName, field.getType(), value);
        }
        if (context.isGenWholePlan()) {
            Condition entityGroupKeyCondition = ParserUtils.getCondition(table.getEntityGroupKey().getName(),
                    conditions);
            // 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);
        }
        action.setSessionId(context.getSessionId());
        List<UpdateAction> actions = new ArrayList<UpdateAction>();
        actions.add(action);
        UpdatePlan plan = new UpdatePlan(actions);
        context.setPlan(plan);
        LOG.debug("UpdatePlan " + plan);
    }

    private LinkedHashSet<String> parseInsertColumns(List<SQLExpr> columns) throws UnsupportedException {
        LinkedHashSet<String> columnsItem = new LinkedHashSet<String>(columns.size());
        for (SQLExpr item : columns) {
            String columnName = parseColumn(item);
            LOG.debug(" SQLInsertItem " + columnName);
            columnsItem.add(columnName);
        }
        return columnsItem;
    }

    public Pair<List<Pair<String, byte[]>>, List<ColumnStruct>> buildFieldsPair(
            MetaEventOperation metaEventOperation, FTable table, LinkedHashSet<String> columns, ValuesClause values)
            throws IOException {
        // Convert a row's each field into byte[] value
        List<SQLExpr> exprValues = values.getValues();
        if (exprValues.size() != columns.size()) {
            throw new IOException("Insert clause " + columns.size() + " columns " + " not match "
                    + exprValues.size() + " values ");
        }
        Pair<String, byte[]>[] array = new Pair[table.getPrimaryKeys().size()];
        // Construct all ColumnAction
        List<ColumnStruct> cols = new ArrayList<ColumnStruct>(columns.size());
        assert (columns.size() == exprValues.size());
        Iterator<String> iter = columns.iterator();
        int i = 0;
        while (iter.hasNext()) {
            String columnName = iter.next();
            // Get the column's info
            Field column = metaEventOperation.getColumnInfo(table, columnName);
            byte[] value = convert(column, exprValues.get(i));
            Iterator<Entry<String, Field>> pkIter = table.getPrimaryKeys().entrySet().iterator();
            int j = 0;
            while (pkIter.hasNext()) {
                if (pkIter.next().getKey().equalsIgnoreCase(columnName)) {
                    array[j] = new Pair<String, byte[]>(columnName, value);
                    break;
                }
                j++;
            }
            // Check the input is the same as DataType
            checkType(column, exprValues.get(i));
            ColumnStruct columnAction = new ColumnStruct(table.getTableName(), column.getFamily(), columnName,
                    column.getType(), value);
            cols.add(columnAction);
            i++;
        }

        return new Pair<List<Pair<String, byte[]>>, List<ColumnStruct>>(Arrays.asList(array), cols);
    }
}