Java Code Examples for com.alibaba.druid.sql.parser.SQLStatementParser#parseStatement()

private static boolean isSupportSelect(String stmt) {
    SQLStatementParser parser = new MySqlStatementParser(stmt);
    SQLStatement statement = parser.parseStatement();
    if (!(statement instanceof SQLSelectStatement)) {
        return false;
    }

    SQLSelectQuery sqlSelectQuery = ((SQLSelectStatement) statement).getSelect().getQuery();
    if (!(sqlSelectQuery instanceof MySqlSelectQueryBlock)) {
        return false;
    }
    MySqlSelectQueryBlock selectQueryBlock = (MySqlSelectQueryBlock) sqlSelectQuery;
    SQLTableSource mysqlFrom = selectQueryBlock.getFrom();
    if (mysqlFrom != null) {
        return false;
    }
    for (SQLSelectItem item : selectQueryBlock.getSelectList()) {
        SQLExpr selectItem = item.getExpr();
        if (!isVariantRef(selectItem)) {
            return false;
        }
    }
    return true;
}
 

private static boolean isInsertSeq(ServerConnection c, String stmt, SchemaConfig schema) throws SQLException {
    SQLStatementParser parser = new MySqlStatementParser(stmt);
    MySqlInsertStatement statement = (MySqlInsertStatement) parser.parseStatement();
    String schemaName = schema == null ? null : schema.getName();
    SQLExprTableSource tableSource = statement.getTableSource();
    SchemaUtil.SchemaInfo schemaInfo = SchemaUtil.getSchemaInfo(c.getUser(), schemaName, tableSource);
    String tableName = schemaInfo.getTable();
    schema = schemaInfo.getSchemaConfig();
    TableConfig tableConfig = schema.getTables().get(tableName);
    if (tableConfig == null) {
        return false;
    } else if (tableConfig.isAutoIncrement()) {
        return true;
    }
    return false;
}
 

private PlanNode getPlanNode(String sql) {
    SQLStatementParser parser = new MySqlStatementParser(sql);
    SQLSelectStatement ast = (SQLSelectStatement) parser.parseStatement();
    MySQLPlanNodeVisitor visitor = new MySQLPlanNodeVisitor("TESTDB", 33, null, true, null);
    visitor.visit(ast);
    return visitor.getTableNode();
}
 

private List<WhereUnit> getAllWhereUnit(String sql) {
    SQLStatementParser parser = new MySqlStatementParser(sql);

    ServerSchemaStatVisitor visitor;
    //throw exception
    try {
        SQLStatement statement = parser.parseStatement();
        visitor = new ServerSchemaStatVisitor();
        statement.accept(visitor);
        return visitor.getAllWhereUnit();
    } catch (Exception e) {
        e.printStackTrace();
    }
    return  null;
}
 

private static String[] getColumns(String sql) {
    String[] columnNames;
    try {
        SQLStatementParser sqlStatementParser = SQLParserUtils.createSQLStatementParser(sql, JdbcUtils.MYSQL);
        SQLStatement statement = sqlStatementParser.parseStatement();
        if (statement instanceof SQLSelectStatement) {
            SQLSelect select = ((SQLSelectStatement) statement).getSelect();
            com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock query = (com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock) select.getQuery();
            int size = query.getSelectList().size();
            if (size == 1){
                if("*".equalsIgnoreCase(   query.getSelectList().get(0).toString())){
                    throw new Exception("unsupport * in select items:"+sql);
                }
            } {
                columnNames = new String[size];
                for (int i = 0; i < size; i++) {
                    columnNames[i] = query.getSelectList().get(i).toString();
                }
                return columnNames;
            }

        }
    }catch (Exception e){
        LOGGER.error("can not get column count",e);
    }
    return new String[]{};
}
 

public static String changeInsertAddSlot(String sql,int slotValue)
{
    SQLStatementParser parser = new MycatStatementParser(sql);
    MySqlInsertStatement insert = (MySqlInsertStatement) parser.parseStatement();
    insert.getColumns().add(new SQLIdentifierExpr("_slot") );
    insert.getValues().getValues().add(new SQLIntegerExpr(slotValue))  ;
    return insert.toString();
}
 

private LoadData parseLoadDataPram(String sql , String connectionCharset)
{
	SQLStatementParser parser = new MycatStatementParser(sql);
	MySqlLoadDataInFileStatement statement = (MySqlLoadDataInFileStatement) parser.parseStatement();

	LoadData loadData = new LoadData();
	SQLTextLiteralExpr rawLineEnd = (SQLTextLiteralExpr) statement.getLinesTerminatedBy();
	String lineTerminatedBy = rawLineEnd == null ? "\n" : rawLineEnd.getText();
	loadData.setLineTerminatedBy(lineTerminatedBy);

	SQLTextLiteralExpr rawFieldEnd = (SQLTextLiteralExpr) statement.getColumnsTerminatedBy();
	String fieldTerminatedBy = rawFieldEnd == null ? "\t" : rawFieldEnd.getText();
	loadData.setFieldTerminatedBy(fieldTerminatedBy);

	SQLTextLiteralExpr rawEnclosed = (SQLTextLiteralExpr) statement.getColumnsEnclosedBy();
	String enclose = rawEnclosed == null ? null : rawEnclosed.getText();
	loadData.setEnclose(enclose);

	SQLTextLiteralExpr escapseExpr =  (SQLTextLiteralExpr)statement.getColumnsEscaped() ;
	String escapse=escapseExpr==null?"\\":escapseExpr.getText();
	loadData.setEscape(escapse);
	String charset = statement.getCharset() != null ? statement.getCharset() : connectionCharset;
	loadData.setCharset(charset);

	String fileName = parseFileName(sql);
	if(StringUtils.isBlank(fileName)){
		throw new RuntimeException(" file name is null !");
	}

	loadData.setFileName(fileName);

	return loadData ;
}
 

private static boolean isMycatSeq(String stmt, SchemaConfig schema)
  {
      if(pattern.matcher(stmt).find()) {
	return true;
}
      SQLStatementParser parser =new MySqlStatementParser(stmt);
      MySqlInsertStatement statement = (MySqlInsertStatement) parser.parseStatement();
      String tableName=   statement.getTableName().getSimpleName();
      TableConfig tableConfig= schema.getTables().get(tableName.toUpperCase());
      if(tableConfig==null) {
	return false;
}
      if(tableConfig.isAutoIncrement())
      {
          boolean isHasIdInSql=false;
          String primaryKey = tableConfig.getPrimaryKey();
          List<SQLExpr> columns = statement.getColumns();
          for (SQLExpr column : columns)
          {
              String columnName = column.toString();
              if(primaryKey.equalsIgnoreCase(columnName))
              {
                  isHasIdInSql = true;
                  break;
              }
          }
          if(!isHasIdInSql) {
		return true;
	}
      }


      return false;
  }
 

private static SQLStatement parseSQL(String stmt) throws SQLSyntaxErrorException {
    SQLStatementParser parser = new MySqlStatementParser(stmt);
    try {
        return parser.parseStatement();
    } catch (Exception t) {
        if (t.getMessage() != null) {
            throw new SQLSyntaxErrorException(t.getMessage());
        } else {
            throw new SQLSyntaxErrorException(t);
        }
    }
}
 

private static boolean isSupportShow(String stmt) {
    SQLStatementParser parser = new MySqlStatementParser(stmt);
    SQLStatement statement = parser.parseStatement();
    return statement instanceof MySqlShowWarningsStatement || statement instanceof MySqlShowVariantsStatement;
}
 

private SQLSelectStatement getSelectStatement(String sql) {
    SQLStatementParser parser = new MySqlStatementParser(sql);
    return (SQLSelectStatement) parser.parseStatement();
}
 

private RouteResultset constructByQuery(String sql, Map<String, String> mapTableToSimple, SchemaConfig schemaConfig) throws SQLException {
    SQLStatementParser parser = new MySqlStatementParser(sql);
    SQLSelectStatement select = (SQLSelectStatement) parser.parseStatement();
    return constructByStatement(sql, mapTableToSimple, select, schemaConfig);
}
 

private SQLStatement parseStmt(String sql) {
    SQLStatementParser statParser = SQLParserUtils.createSQLStatementParser(sql, "mysql");
    SQLStatement stmt = statParser.parseStatement();
    return stmt;
}
 

private SQLStatement parseStmt(String sql) {
	SQLStatementParser statParser = SQLParserUtils.createSQLStatementParser(sql, "mysql");
	SQLStatement stmt = statParser.parseStatement();
	return stmt;		
}
 

@Override
public boolean checkDmlPrivilege(String user, String schema, String sql) {

	if ( schema == null ) {
		return true;
	}
	
	boolean isPassed = false;

	MycatConfig conf = MycatServer.getInstance().getConfig();
	UserConfig userConfig = conf.getUsers().get(user);
	if (userConfig != null) {
		
		UserPrivilegesConfig userPrivilege = userConfig.getPrivilegesConfig();
		if ( userPrivilege != null && userPrivilege.isCheck() ) {				
		
			UserPrivilegesConfig.SchemaPrivilege schemaPrivilege = userPrivilege.getSchemaPrivilege( schema );
			if ( schemaPrivilege != null ) {
	
				String tableName = null;
				int index = -1;
				
				//TODO 此处待优化，寻找更优SQL 解析器
				
				//修复bug
				// https://github.com/alibaba/druid/issues/1309
				//com.alibaba.druid.sql.parser.ParserException: syntax error, error in :'begin',expect END, actual EOF begin
				if ( sql != null && sql.length() == 5 && sql.equalsIgnoreCase("begin") ) {
					return true;
				}
				
				SQLStatementParser parser = new MycatStatementParser(sql);			
				SQLStatement stmt = parser.parseStatement();

				if (stmt instanceof MySqlReplaceStatement || stmt instanceof SQLInsertStatement ) {
					index = 0;
				} else if (stmt instanceof SQLUpdateStatement ) {
					index = 1;
				} else if (stmt instanceof SQLSelectStatement ) {
					index = 2;
				} else if (stmt instanceof SQLDeleteStatement ) {
					index = 3;
				}

				if ( index > -1) {
					
					SchemaStatVisitor schemaStatVisitor = new MycatSchemaStatVisitor();
					stmt.accept(schemaStatVisitor);
					String key = schemaStatVisitor.getCurrentTable();
					if ( key != null ) {
						
						if (key.contains("`")) {
							key = key.replaceAll("`", "");
						}
						
						int dotIndex = key.indexOf(".");
						if (dotIndex > 0) {
							tableName = key.substring(dotIndex + 1);
						} else {
							tableName = key;
						}							
						
						//获取table 权限, 此处不需要检测空值, 无设置则自动继承父级权限
						UserPrivilegesConfig.TablePrivilege tablePrivilege = schemaPrivilege.getTablePrivilege( tableName );
						if ( tablePrivilege.getDml()[index] > 0 ) {
							isPassed = true;
						}
						
					} else {
						//skip
						isPassed = true;
					}
					
					
				} else {						
					//skip
					isPassed = true;
				}
				
			} else {					
				//skip
				isPassed = true;
			}
			
		} else {
			//skip
			isPassed = true;
		}

	} else {
		//skip
		isPassed = true;
	}
	
	if( !isPassed ) {
		 ALARM.error(new StringBuilder().append(Alarms.DML_ATTACK ).append("[sql=").append( sql )
                    .append(",user=").append(user).append(']').toString());
	}
	
	return isPassed;
}
 

@Override
public void start(String sql)
{
    clear();
    this.sql = sql;


    SQLStatementParser parser = new MycatStatementParser(sql);
    statement = (MySqlLoadDataInFileStatement) parser.parseStatement();
    fileName = parseFileName(sql);

    if (fileName == null)
    {
        serverConnection.writeErrMessage(ErrorCode.ER_FILE_NOT_FOUND, " file name is null !");
        clear();
        return;
    }
    schema = MycatServer.getInstance().getConfig()
            .getSchemas().get(serverConnection.getSchema());
    if (schema == null){
        throw new RuntimeException("please sql:use schema  before load data");
    }
    tableId2DataNodeCache = (LayerCachePool) MycatServer.getInstance().getCacheService().getCachePool("TableID2DataNodeCache");
    tableName = statement.getTableName().getSimpleName().toUpperCase();
    tableConfig = schema.getTables().get(tableName);
  if(  tableConfig.getRule() != null && tableConfig.getRule().getRuleAlgorithm() instanceof SlotFunction){
      shoudAddSlot=true;
  }
    tempPath = SystemConfig.getHomePath() + File.separator + "temp" + File.separator + serverConnection.getId() + File.separator;
    tempFile = tempPath + "clientTemp.txt";
    tempByteBuffer = new ByteArrayOutputStream();

    List<SQLExpr> columns = statement.getColumns();
    if(tableConfig!=null)
    {
        String pColumn = getPartitionColumn();
        if (pColumn != null && columns != null && columns.size() > 0) {
            for (int i = 0, columnsSize = columns.size(); i < columnsSize; i++) {
                String column = StringUtil.removeBackquote(columns.get(i).toString());
                if (pColumn.equalsIgnoreCase(column)) {
                    partitionColumnIndex = i;
                }
                if("_slot".equalsIgnoreCase(column)){
                    shoudAddSlot=false;
                }
            }

        }
    }
        if(shoudAddSlot){
            columns.add(new SQLIdentifierExpr("_slot"));
        }
    parseLoadDataPram();
    if (statement.isLocal())
    {
        isStartLoadData = true;
        //向客户端请求发送文件
        ByteBuffer buffer = serverConnection.allocate();
        RequestFilePacket filePacket = new RequestFilePacket();
        filePacket.fileName = fileName.getBytes();
        filePacket.packetId = 1;
        filePacket.write(buffer, serverConnection, true);
    } else
    {
        if (!new File(fileName).exists())
        {
            serverConnection.writeErrMessage(ErrorCode.ER_FILE_NOT_FOUND, fileName + " is not found!");
            clear();
        } else
        {
            parseFileByLine(fileName, loadData.getCharset(), loadData.getLineTerminatedBy());
            RouteResultset rrs = buildResultSet(routeResultMap);
            if (rrs != null)
            {
                flushDataToFile();
                isStartLoadData = false;
                serverConnection.getSession2().execute(rrs, ServerParse.LOAD_DATA_INFILE_SQL);
            }

        }
    }
}
 

private Object[] getParseTables(String sql) throws Exception{
	
	SQLStatementParser parser = new MySqlStatementParser(sql);
	SQLStatement statement = parser.parseStatement();
       MycatSchemaStatVisitor visitor = new MycatSchemaStatVisitor();
       
       
       LayerCachePool cachePool = mock(LayerCachePool.class);
       RouteResultset rrs = new RouteResultset(sql, ServerParse.SELECT);
	
	druidParser.parser(schema, rrs, statement, sql, cachePool, visitor);
	
	DruidShardingParseInfo ctx = druidParser.getCtx();
	return ctx.getTables().toArray();
}