PGSQLSelectVisitor.java

package io.anyway.sherlock.sqlparser.visitor.pgsql;

import java.util.LinkedList;
import java.util.List;

import io.anyway.sherlock.exception.SQLParserException;
import io.anyway.sherlock.merger.Limit;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLOrderBy;
import com.alibaba.druid.sql.ast.SQLSetQuantifier;
import com.alibaba.druid.sql.ast.expr.SQLAggregateExpr;
import com.alibaba.druid.sql.ast.expr.SQLAllColumnExpr;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.expr.SQLIntegerExpr;
import com.alibaba.druid.sql.ast.expr.SQLNumericLiteralExpr;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause;
import com.alibaba.druid.sql.ast.statement.SQLSelectItem;
import com.alibaba.druid.sql.ast.statement.SQLSelectOrderByItem;
import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock;
import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.PGLimit;
import com.alibaba.druid.sql.dialect.postgresql.visitor.PGOutputVisitor;
import com.google.common.base.Optional;

import io.anyway.sherlock.sqlparser.bean.AggregationColumn;
import io.anyway.sherlock.sqlparser.bean.AggregationColumn.AggregationType;
import io.anyway.sherlock.sqlparser.bean.GroupColumn;
import io.anyway.sherlock.sqlparser.bean.OrderColumn;
import io.anyway.sherlock.sqlparser.bean.OrderColumn.OrderType;

/**
 * PostgreSQL查询访问收集器，主要收集表名、条件项、列元数据、distinct、聚合表达式、groupby、orderby和limit等内容
 * @author yangzz
 *
 */
public class PGSQLSelectVisitor extends AbstractPGSQLVisitor {
	
	final private static String AUTO_GEN_COL= "auto_gen_col_";
	
	final private static String AUTO_GEN_COL_COUNT= AUTO_GEN_COL + "count";
	
	private int selectLayer= 0;
	
	private boolean finishCollectMetadata= false;
	
	private boolean attachCountExpression= false;
	
	private List<String> missOrderbyColumns;
	
    //遍历表名
    @Override
    public boolean visit(final PGSelectQueryBlock x) {
    	selectLayer++;
        if (x.getFrom() instanceof SQLExprTableSource) {
            SQLExprTableSource tableExpr = (SQLExprTableSource) x.getFrom();
            setCurrentTable(tableExpr.getExpr().toString(), Optional.fromNullable(tableExpr.getAlias()));
        }
        //处理distinct
        if(isEnableCollectMetadata()){
        	if(SQLSetQuantifier.DISTINCT==x.getDistionOption()){
        		parseResult.markDistinct();
        	}
        }
        return super.visit(x);
    }
    
    @Override
    public void endVisit(final PGSelectQueryBlock x) {
    	//把缺失的orderby列补上,union太复杂暂时不支持
    	if(isMasterSelect()){
	    	if(!CollectionUtils.isEmpty(missOrderbyColumns)){
	    		String orderby_columns="";
	    		for(String columnName: missOrderbyColumns){
	    			orderby_columns+= ", "+ columnName;
	    		}
	    		parseResult.getSqlBuilder().buildSQL("select_missing_columns", orderby_columns);
	    	}
	    	//增加一个limit东西限制查询大小
	    	if(parseResult.getLimit()== null){
	    		//print(" limit 9");
	    	}
    	}
    	
    	if(isEnableCollectMetadata()){
    		finishCollectMetadata= true;
    	}
    	selectLayer--;
    }
    
    /*
     * 主要屏蔽嵌套子查询和UNIO后边查询的东西
     * 
     * 场景一：select a ,(select b from ...) from ... where xx= (select xx from ...); 不处理嵌套子查询
     * 场景二：(select xxx from yyy) UNION [ALL] (select xxx from yyy2); 不处理UNION 后边的查询
     */
    private boolean isEnableCollectMetadata(){
    	return !finishCollectMetadata && selectLayer == 1;
    }
    
    private boolean isMasterSelect(){
    	return selectLayer == 1;
    }
    
    /*
     * 只遍历第一级SELECT选择项，收集Metadata信息，同时收集MAX、MIN、COUNT、SUM和AVG操作；建议子查询和表达式使用alias别名
     */
    @Override
    protected void printSelectList(List<SQLSelectItem> selectList) {
        super.printSelectList(selectList);
        //如果是嵌套子查询直接滤过
        if(!isMasterSelect()){
        	return;
        }
        int columnIndex= 0;
        //遍历第一级SELECT选择项
        for(SQLSelectItem each: selectList){
        	columnIndex++;
        	String alias= each.getAlias();
        	SQLExpr expr= each.getExpr();
        	//如果是表达式
        	if(expr instanceof SQLAggregateExpr){
        		SQLAggregateExpr x= (SQLAggregateExpr)expr;
        		StringBuilder expression = new StringBuilder();
        		x.accept(new PGOutputVisitor(expression));
        		try{
        			//如果不是AggregationType枚举中的内容，抛异常直接跳过不做处理
	        		AggregationType aggregationType = AggregationType.valueOf(x.getMethodName().toUpperCase());
	        		//解析过一遍就不在解析，此段逻辑主要处理如下场景
	        		//(select avg(salary) from emp where userid<500) union (select avg(salary) from emp where userid>100)
	        		//解析完第一个select的元数据就不在解析，但是对于union后面的select语句，avg聚合字段也要追加 count(1)列
	        		if(isEnableCollectMetadata()){
		        		AggregationColumn aggregationColumn= new AggregationColumn(expression.toString(),columnIndex,aggregationType);
		        		parseResult.addAggregationColumn(aggregationColumn);
		        		if(logger.isInfoEnabled()){
		        			logger.info("AggregationColumn: "+aggregationColumn);
		        		}
	        		}
	        		//在有表达式的SQLSelectItem后面加上只追加一个COUNT(1) 即可，然后通过(avg1*count1+avg2*count2)/(count1+count2)求最终平均值
	        		//原始SQL：select avg(salary) as avg_salary, (select avg(bonus) from user) from user  
	        	    //转换SQL：select avg(salary) as avg_salary, (select avg(bonus) from user), count(1) as auto_gen_col_count from user
	        		if(AggregationType.AVG.equals(aggregationType) && !attachCountExpression){
	        			print(", COUNT(1) AS " + AUTO_GEN_COL_COUNT);
	        			attachCountExpression= true;
	        			if(logger.isInfoEnabled()){
	            			logger.info("auto create count(1) AS "+AUTO_GEN_COL_COUNT);
	            		}
	        		}
        		}catch(IllegalArgumentException ex){
        			if(logger.isInfoEnabled()){
            			logger.info(expression.toString()+" expression not in (MAX、MIN、COUNT、SUM and AVG)");
            		}
        		}
        	}
        	
        	//可以收集列元数据信息
        	if(isEnableCollectMetadata()){
	        	//如果别名不为空，则存储别名
	        	if(!StringUtils.isEmpty(alias)){
	        		parseResult.addMetadataColumn(alias);
	        		continue;
	        	}
	        	//如果是其他列名
	        	if(expr instanceof SQLIdentifierExpr){
	        		parseResult.addMetadataColumn(((SQLIdentifierExpr)expr).getName());
	        	}
	        	else if(expr instanceof SQLPropertyExpr){
	        		parseResult.addMetadataColumn(((SQLPropertyExpr)expr).getName());
	        	}
	        	else if(expr instanceof SQLAllColumnExpr){
	        		throw new SQLParserException("not support select * for sql: "+sql+",please enumerate every column");
	        	}
	        	else{
	        		//可能还是子查询SQLQueryExpr
	        		String columnName= AUTO_GEN_COL + columnIndex;
	        		parseResult.addMetadataColumn(columnName);
	        		StringBuilder expression = new StringBuilder();
	        		expr.accept(new PGOutputVisitor(expression));
	        		logger.warn(expression.toString()+" miss alias, will use temporary column name: "+columnName);
	        	}
        	}
        }
        
        /*
    	 * 预留一个位置给那些不在select选项里的oderby列，否则结果集无法按给定的orderby合并排序
    	 * 如 select name from emp order by deptno，需要把deptno 追加到select选项中
    	 * 最终sql：select name ,deptno from emp order by deptno
    	 */
    	parseResult.getSqlBuilder().appendToken("select_missing_columns", false);
    }
    
    //解析Groupby节点
    @Override
    public boolean visit(SQLSelectGroupByClause x){
    	//如果嵌套子查询或第一个SQL解析完毕
    	if(!isEnableCollectMetadata()){
    		return super.visit(x); 
    	}
    	for(SQLExpr expr:x.getItems()){
    		String columnName= null;
    		if(expr instanceof SQLPropertyExpr) {
    			columnName= ((SQLPropertyExpr)expr).getName();
    		} else if (expr instanceof SQLIdentifierExpr) {
    			columnName = ((SQLIdentifierExpr)expr).getName();
    		}
    		if(!StringUtils.isEmpty(columnName)){
				int index= parseResult.getMetadataColumns().indexOf(columnName);
				//不应该存在groupby列不在select选项里的情况
				if(index!= -1){
					parseResult.addGroupColumn(new GroupColumn(columnName,++index));
					if(logger.isInfoEnabled()){
	        			logger.info("group column ["+columnName+"] index is: "+index);
	        		}
				}
				else{
					logger.warn("cannot found group column ["+columnName+"] in metadatacolumns: "+parseResult.getMetadataColumns());
				}
    		}
    	}
    	return super.visit(x);
    }
    
    /*
     * 遍历orderby内容
     */
    @Override
    public boolean visit(SQLOrderBy x) {
    	//如果嵌套子查询或第一个SQL解析完毕
    	if(!isEnableCollectMetadata()){
    		return visit(x);
    	}
    	for (SQLSelectOrderByItem each : x.getItems()) {
            SQLExpr expr = each.getExpr();
            //默认是升序
            OrderType orderType= each.getType()== null? OrderType.ASC: ("DESC".equalsIgnoreCase(each.getType().toString())? OrderType.DESC : OrderType.ASC);
            if (expr instanceof SQLIntegerExpr) {
            	int index= ((SQLIntegerExpr)expr).getNumber().intValue();
            	if(index > parseResult.getMetadataColumns().size()){
            		logger.warn("order index ["+index+"] is outboundary of metadatacolumns: "+parseResult.getMetadataColumns());
            	}
            	else{
            		parseResult.addOrderColumn(new OrderColumn(parseResult.getMetadataColumns().get(index-1),orderType,index));
            	}
            	continue;
            } 
            String columnName= null;
            String owner= null;
            if (expr instanceof SQLIdentifierExpr) {
            	columnName= ((SQLIdentifierExpr)expr).getName();
            }
            else if (expr instanceof SQLPropertyExpr) {
                columnName= ((SQLPropertyExpr)expr).getName();
                SQLExpr exp= ((SQLPropertyExpr)expr).getOwner();
                if(exp!=null){
                	owner= exp.toString();
                }
            }
            if(!StringUtils.isEmpty(columnName)){
            	int index= parseResult.getMetadataColumns().indexOf(columnName);
            	//如果orderby的列不在select选项里，则需要追加处理
            	if(index!= -1){
            		parseResult.addOrderColumn(new OrderColumn(columnName,orderType,++index));
            		if(logger.isInfoEnabled()){
            			logger.info("order column ["+columnName+"] index is: "+index);
            		}
            	}
            	else{
            		//只有当grouby为空的时候才能追加orderby否则sql执行会出粗
            		if(CollectionUtils.isEmpty(parseResult.getGroupColumns())){
	            		index= parseResult.getMetadataColumns().size() + (attachCountExpression?1:0);
	            		if(StringUtils.isEmpty(missOrderbyColumns)){
	            			missOrderbyColumns= new LinkedList<String>();
	            		}
	            		if(!missOrderbyColumns.contains(columnName)){
	            			missOrderbyColumns.add((owner==null?"":owner+".")+columnName);
	            			index= missOrderbyColumns.size() + index;
	            			parseResult.addOrderColumn(new OrderColumn(columnName,orderType,index));
	            			if(logger.isInfoEnabled()){
	                			logger.info("will append missing orderby column ["+columnName+"] index is: "+index);
	                		}
	            		}
	            		else{
	            			logger.warn("duplicated orderby column ["+columnName+"] "+orderType +",it doesnt work");
	            		}
            		}
            	}
            }
        }
        return super.visit(x);
    }
    
    //解析Limit节点
    @Override
    public boolean visit(PGLimit x) {
    	//如果嵌套子查询或第一个SQL解析完毕
    	if(!isEnableCollectMetadata()){
    		return super.visit(x); 
    	}
    	//第一次解析
        int offset = 0;
        String offsetFragment= "";
        if (null != x.getOffset()) {
            if (x.getOffset() instanceof SQLNumericLiteralExpr) {
                offset = ((SQLNumericLiteralExpr) x.getOffset()).getNumber().intValue();
                offsetFragment= " OFFSET 0";
            } else {
                offset = ((Number) getParameters().get(((SQLVariantRefExpr) x.getOffset()).getIndex())).intValue();
                getParameters().set(((SQLVariantRefExpr) x.getOffset()).getIndex(), 0);
                offsetFragment=" OFFSET ?";
            }
        }
        int rowCount;
        if (x.getRowCount() instanceof SQLNumericLiteralExpr) {
            rowCount = ((SQLNumericLiteralExpr) x.getRowCount()).getNumber().intValue();
            print(rowCount + offset);
        } else {
            rowCount = ((Number) getParameters().get(((SQLVariantRefExpr) x.getRowCount()).getIndex())).intValue();
            getParameters().set(((SQLVariantRefExpr) x.getRowCount()).getIndex(), rowCount + offset);
            print("?");
        }
        //最后在再打印offset字段
        print(offsetFragment);
        
        parseResult.setLimit(new Limit(offset, rowCount));
        if(logger.isInfoEnabled()){
			logger.info("Limit [offset: "+offset+",rowCount: "+rowCount+"]");
		}
        return false;
    }

}