package io.mycat.route.parser.druid.impl;
import io.mycat.route.RouteResultset;
import io.mycat.route.RouteResultsetNode;
import io.mycat.route.function.AbstractPartitionAlgorithm;
import io.mycat.route.parser.druid.MycatSchemaStatVisitor;
import io.mycat.route.parser.druid.RouteCalculateUnit;
import io.mycat.route.util.RouterUtil;
import io.mycat.server.config.node.SchemaConfig;
import io.mycat.server.config.node.TableConfig;
import io.mycat.server.executors.FetchStoreNodeOfChildTableHandler;
import io.mycat.server.parser.ServerParse;
import io.mycat.util.StringUtil;
import java.sql.SQLNonTransientException;
import java.sql.SQLSyntaxErrorException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
import com.alibaba.druid.sql.ast.expr.SQLCharExpr;
import com.alibaba.druid.sql.ast.expr.SQLIntegerExpr;
import com.alibaba.druid.sql.ast.statement.SQLInsertStatement.ValuesClause;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlInsertStatement;
public class DruidInsertParser extends DefaultDruidParser {
@Override
public void visitorParse(RouteResultset rrs, SQLStatement stmt, MycatSchemaStatVisitor visitor) throws SQLNonTransientException {
}
/**
* 考虑因素：isChildTable、批量、是否分片
*/
@Override
public void statementParse(SchemaConfig schema, RouteResultset rrs, SQLStatement stmt) throws SQLNonTransientException {
MySqlInsertStatement insert = (MySqlInsertStatement)stmt;
String tableName = StringUtil.removeBackquote(insert.getTableName().getSimpleName()).toUpperCase();
ctx.addTable(tableName);
if(RouterUtil.isNoSharding(schema,tableName)) {//整个schema都不分库或者该表不拆分
RouterUtil.routeForTableMeta(rrs, schema, tableName, rrs.getStatement());
rrs.setFinishedRoute(true);
return;
}
TableConfig tc = schema.getTables().get(tableName);
if(tc == null) {
String msg = "can't find table define in schema "
+ tableName + " schema:" + schema.getName();
LOGGER.warn(msg);
throw new SQLNonTransientException(msg);
} else {
//childTable的insert直接在解析过程中完成路由
if (tc.isChildTable()) {
parserChildTable(schema, rrs, tableName, insert);
return;
}
String partitionColumn = tc.getPartitionColumn();
if(partitionColumn != null) {//分片表
//拆分表必须给出column list,否则无法寻找分片字段的值
if(insert.getColumns() == null || insert.getColumns().size() == 0) {
throw new SQLSyntaxErrorException("partition table, insert must provide ColumnList");
}
//批量insert
if(isMultiInsert(insert)) {
// String msg = "multi insert not provided" ;
// LOGGER.warn(msg);
// throw new SQLNonTransientException(msg);
parserBatchInsert(schema, rrs, partitionColumn, tableName, insert);
} else {
parserSingleInsert(schema, rrs, partitionColumn, tableName, insert);
}
}
}
}
/**
* 寻找joinKey的索引
* @param columns
* @param joinKey
* @return -1表示没找到，>=0表示找到了
*/
private int getJoinKeyIndex(List<SQLExpr> columns, String joinKey) {
for(int i = 0; i < columns.size(); i++) {
String col = StringUtil.removeBackquote(columns.get(i).toString()).toUpperCase();
if(col.equals(joinKey)) {
return i;
}
}
return -1;
}
/**
* 是否为批量插入：insert into ...values (),()...或 insert into ...select.....
* @param insertStmt
* @return
*/
private boolean isMultiInsert(MySqlInsertStatement insertStmt) {
return (insertStmt.getValuesList() != null && insertStmt.getValuesList().size() > 1) || insertStmt.getQuery() != null;
}
private RouteResultset parserChildTable(SchemaConfig schema, RouteResultset rrs,
String tableName, MySqlInsertStatement insertStmt) throws SQLNonTransientException {
TableConfig tc = schema.getTables().get(tableName);
String joinKey = tc.getJoinKey();
int joinKeyIndex = getJoinKeyIndex(insertStmt.getColumns(), joinKey);
if(joinKeyIndex == -1) {
String inf = "joinKey not provided :" + tc.getJoinKey()+ "," + insertStmt;
LOGGER.warn(inf);
throw new SQLNonTransientException(inf);
}
if(isMultiInsert(insertStmt)) {
String msg = "ChildTable multi insert not provided" ;
LOGGER.warn(msg);
throw new SQLNonTransientException(msg);
}
String joinKeyVal = insertStmt.getValues().getValues().get(joinKeyIndex).toString();
String sql = insertStmt.toString();
// try to route by ER parent partion key
RouteResultset theRrs = RouterUtil.routeByERParentKey(null,schema, ServerParse.INSERT,sql, rrs, tc,joinKeyVal);
if (theRrs != null) {
rrs.setFinishedRoute(true);
return theRrs;
}
// route by sql query root parent's datanode
String findRootTBSql = tc.getLocateRTableKeySql().toLowerCase() + joinKeyVal;
if (LOGGER.isDebugEnabled()) {
LOGGER.debug("find root parent's node sql "+ findRootTBSql);
}
FetchStoreNodeOfChildTableHandler fetchHandler = new FetchStoreNodeOfChildTableHandler();
String dn = fetchHandler.execute(schema.getName(),findRootTBSql, tc.getRootParent().getDataNodes());
if (dn == null) {
throw new SQLNonTransientException("can't find (root) parent sharding node for sql:"+ sql);
}
if (LOGGER.isDebugEnabled()) {
LOGGER.debug("found partion node for child table to insert "+ dn + " sql :" + sql);
}
return RouterUtil.routeToSingleNode(rrs, dn, sql);
}
/**
* 单条insert（非批量）
* @param schema
* @param rrs
* @param partitionColumn
* @param tableName
* @param insertStmt
* @throws SQLNonTransientException
*/
private void parserSingleInsert(SchemaConfig schema, RouteResultset rrs, String partitionColumn,
String tableName, MySqlInsertStatement insertStmt) throws SQLNonTransientException {
boolean isFound = false;
for(int i = 0; i < insertStmt.getColumns().size(); i++) {
if(partitionColumn.equalsIgnoreCase(StringUtil.removeBackquote(insertStmt.getColumns().get(i).toString()))) {//找到分片字段
isFound = true;
String column = StringUtil.removeBackquote(insertStmt.getColumns().get(i).toString());
String value = StringUtil.removeBackquote(insertStmt.getValues().getValues().get(i).toString());
RouteCalculateUnit routeCalculateUnit = new RouteCalculateUnit();
routeCalculateUnit.addShardingExpr(tableName, column, value);
ctx.addRouteCalculateUnit(routeCalculateUnit);
//mycat是单分片键，找到了就返回
break;
}
}
if(!isFound) {//分片表的
String msg = "bad insert sql (sharding column:"+ partitionColumn + " not provided," + insertStmt;
LOGGER.warn(msg);
throw new SQLNonTransientException(msg);
}
// insert into .... on duplicateKey
//such as :INSERT INTO TABLEName (a,b,c) VALUES (1,2,3) ON DUPLICATE KEY UPDATE b=VALUES(b);
//INSERT INTO TABLEName (a,b,c) VALUES (1,2,3) ON DUPLICATE KEY UPDATE c=c+1;
if(insertStmt.getDuplicateKeyUpdate() != null) {
List<SQLExpr> updateList = insertStmt.getDuplicateKeyUpdate();
for(SQLExpr expr : updateList) {
SQLBinaryOpExpr opExpr = (SQLBinaryOpExpr)expr;
String column = StringUtil.removeBackquote(opExpr.getLeft().toString().toUpperCase());
if(column.equals(partitionColumn)) {
String msg = "partion key can't be updated: " + tableName + " -> " + partitionColumn;
LOGGER.warn(msg);
throw new SQLNonTransientException(msg);
}
}
}
}
/**
* insert into .... select .... 或insert into table() values (),(),....
* @param schema
* @param rrs
* @param insertStmt
* @throws SQLNonTransientException
*/
private void parserBatchInsert(SchemaConfig schema, RouteResultset rrs, String partitionColumn,
String tableName, MySqlInsertStatement insertStmt) throws SQLNonTransientException {
//insert into table() values (),(),....
if(insertStmt.getValuesList().size() > 1) {
//字段列数
int columnNum = insertStmt.getColumns().size();
int shardingColIndex = getSharingColIndex(insertStmt, partitionColumn);
if(shardingColIndex == -1) {
String msg = "bad insert sql (sharding column:"+ partitionColumn + " not provided," + insertStmt;
LOGGER.warn(msg);
throw new SQLNonTransientException(msg);
} else {
List<ValuesClause> valueClauseList = insertStmt.getValuesList();
Map<Integer,List<ValuesClause>> nodeValuesMap = new HashMap<Integer,List<ValuesClause>>();
TableConfig tableConfig = schema.getTables().get(tableName);
AbstractPartitionAlgorithm algorithm = tableConfig.getRule().getRuleAlgorithm();
for(ValuesClause valueClause : valueClauseList) {
if(valueClause.getValues().size() != columnNum) {
String msg = "bad insert sql columnSize != valueSize:"
+ columnNum + " != " + valueClause.getValues().size()
+ "values:" + valueClause;
LOGGER.warn(msg);
throw new SQLNonTransientException(msg);
}
SQLExpr expr = valueClause.getValues().get(shardingColIndex);
String shardingValue = null;
if(expr instanceof SQLIntegerExpr) {
SQLIntegerExpr intExpr = (SQLIntegerExpr)expr;
shardingValue = intExpr.getNumber() + "";
} else if (expr instanceof SQLCharExpr) {
SQLCharExpr charExpr = (SQLCharExpr)expr;
shardingValue = charExpr.getText();
}
Integer nodeIndex = algorithm.calculate(shardingValue);
//没找到插入的分片
if(nodeIndex == null) {
String msg = "can't find any valid datanode :" + tableName
+ " -> " + partitionColumn + " -> " + shardingValue;
LOGGER.warn(msg);
throw new SQLNonTransientException(msg);
}
if(nodeValuesMap.get(nodeIndex) == null) {
nodeValuesMap.put(nodeIndex, new ArrayList<ValuesClause>());
}
nodeValuesMap.get(nodeIndex).add(valueClause);
System.out.println();
}
RouteResultsetNode[] nodes = new RouteResultsetNode[nodeValuesMap.size()];
int count = 0;
for(Map.Entry<Integer,List<ValuesClause>> node : nodeValuesMap.entrySet()) {
Integer nodeIndex = node.getKey();
List<ValuesClause> valuesList = node.getValue();
insertStmt.setValuesList(valuesList);
nodes[count++] = new RouteResultsetNode(tableConfig.getDataNodes().get(nodeIndex),
rrs.getSqlType(),insertStmt.toString());
}
rrs.setNodes(nodes);
rrs.setFinishedRoute(true);
}
} else if(insertStmt.getQuery() != null) { // insert into .... select ....
String msg = "TODO:insert into .... select .... not supported!";
LOGGER.warn(msg);
throw new SQLNonTransientException(msg);
}
}
/**
* 寻找拆分字段在 columnList中的索引
* @param insertStmt
* @param partitionColumn
* @return
*/
private int getSharingColIndex(MySqlInsertStatement insertStmt,String partitionColumn) {
int shardingColIndex = -1;
for(int i = 0; i < insertStmt.getColumns().size(); i++) {
if(partitionColumn.equalsIgnoreCase(StringUtil.removeBackquote(insertStmt.getColumns().get(i).toString()))) {//找到分片字段
shardingColIndex = i;
return shardingColIndex;
}
}
return shardingColIndex;
}
}