/*** @author leo_zl  2022/2/28*/publicclassMysqlParserVisitorextendsMySqlASTVisitorAdapter {
/*** 限制访问的表名* [parser_table,demo_table]*/privatefinalSettableNameSet;
/*** 每个表的约束条件** {*    "parser_table":"id=1" ,*    "demo_table":"id=2"* }*/privatefinalMaptableConditionMap;
publicMysqlParserVisitor(SettableNameSet, MaptableConditionMap) {
this.tableNameSet=tableNameSet;
this.tableConditionMap=tableConditionMap;
    }
@Overridepublicbooleanvisit(SQLExprTableSourcereference) {
StringtableName=reference.getName().toString();
// 符合目标表源名if (isTargetTable(tableName)) {
SQLObjectparent=reference.getParent();
Stringalias=reference.getAlias();
// 回溯到选择语句while (!(parentinstanceofMySqlSelectQueryBlock) &&parent!=null) {
parent=parent.getParent();
            }
// 插入行控制条件Stringconditions=formCondition(tableName, alias);
if (StringUtils.isNotEmpty(conditions) &&parent!=null) {
                ((MySqlSelectQueryBlock) parent).addCondition(conditions);
            }
        }
returnfalse;
    }
privateStringformCondition(StringtableName, Stringalias) {
Stringcondition=tableConditionMap.get(tableName);
if (StringUtils.isEmpty(condition)) {
returnnull;
        }
returnalias==null?condition : alias+"."+condition;
    }
privatebooleanisTargetTable(StringtableName) {
returntableNameSet.contains(tableName);
    }
}

/*** @author leo_zl  2022/2/28*/@RunWith(PowerMockRunner.class)
@PrepareForTest(MysqlParserVisitor.class)
publicclassMysqlParserVisitorTest {
@Testpublicvoidvisit() {
SettableNameSet=newHashSet<>();
tableNameSet.add("parser_table");
MaptableConditionMap=newHashMap<>();
tableConditionMap.put("parser_table", "id=1");
StringinputSql="SELECT * FROM parser_table";
SQLStatementsqlStatement=SQLUtils.parseSingleStatement(inputSql, JdbcConstants.MYSQL);
MysqlParserVisitorvisitor=newMysqlParserVisitor(tableNameSet, tableConditionMap);
sqlStatement.accept(visitor);
StringexpectSql="SELECT *\n"+"FROM parser_table\n"+"WHERE id = 1";
Assert.assertEquals(expectSql, sqlStatement.toString());
    }
}

/*** @author leo_zl  2022/2/9*/publicclassPhoenixSqlParser {
privatestaticfinalLoggerLOG=LoggerFactory.getLogger(PhoenixSqlParser.class);
publicStringprivilegeControl(Stringsql, MapprivilegeInfo) {
LOG.info("original sql :\n{}", sql);
SQLParsersqlParser=newSQLParser(sql);
try {
SelectStatementselectStatement=sqlParser.parseQuery();
selectStatement=parserSelect(selectStatement, privilegeInfo);
if (selectStatement.isUnion()) {
Listselects=selectStatement.getSelects();
returnStringUtils.join(selects, " \n UNION ALL \n");
            }
returnselectStatement.toString();
        } catch (SQLExceptione) {
LOG.info("解析SQL 异常:{}", e.getMessage());
        }
returnsql;
    }
privateSelectStatementparserSelect(SelectStatementselectStatement, MapprivilegeInfo) throwsSQLException {
if (selectStatement.isUnion()) {
returnparseUnion(selectStatement, privilegeInfo);
        }
// 获取from nodereturnparseTable(selectStatement, privilegeInfo);
    }
privateSelectStatementparseUnion(SelectStatementselectStatement, MapprivilegeInfo) throwsSQLException {
Listselects=selectStatement.getSelects();
ListnewStatements=newArrayList<>();
for (SelectStatementselect : selects) {
SelectStatementnewStatement=parserSelect(select, privilegeInfo);
newStatements.add(newStatement);
        }
StringdestSql=StringUtils.join(newStatements, " \n UNION ALL ");
SQLParsersqlParser=newSQLParser(destSql);
returnsqlParser.parseQuery();
    }
privateSelectStatementparseTable(SelectStatementselectStatement, MapprivilegeInfo) throwsSQLException {
TableNodetableNode=selectStatement.getFrom();
if (tableNodeinstanceofJoinTableNode) {
// join 查询JoinTableNodejoinTableNode= (JoinTableNode) tableNode;
TableNodelhs=joinTableNode.getLHS();
SelectStatementleftSelect=parseJoinTable(selectStatement, lhs, privilegeInfo);
TableNoderhs=joinTableNode.getRHS();
returnparseJoinTable(leftSelect, rhs, privilegeInfo);
        } else {
returnparseJoinTable(selectStatement, tableNode, privilegeInfo);
        }
    }
privateSelectStatementparseJoinTable(SelectStatementselectStatement, TableNodetableNode, MapprivilegeInfo) throwsSQLException {
if (tableNodeinstanceofDerivedTableNode) {
// 子查询DerivedTableNodesubQuery= (DerivedTableNode) tableNode;
SelectStatementsubQuerySelect=subQuery.getSelect();
returnparserSelect(subQuerySelect, privilegeInfo);
        } elseif (tableNodeinstanceofNamedTableNode) {
// 普通表NamedTableNodenamedTableNode= (NamedTableNode) tableNode;
TableNamename=namedTableNode.getName();
Stringalias=namedTableNode.getAlias();
ParseNodeparseNode=constructParseNode(name, alias, privilegeInfo);
returnselectStatement.combine(parseNode);
        }
returnselectStatement;
    }
privateParseNodeconstructParseNode(TableNametableName, StringaliasName, MapprivilegeInfo) throwsSQLException {
Stringconditions=privilegeInfo.get(tableName.getTableName());
if (Objects.isNull(conditions)) {
returnnull;
        }
conditions=aliasName==null?conditions : aliasName+"."+conditions;
// 因为不熟悉Phoenix parser的API，所以暂时使用这种笨办法，构建对应的nodeStringBuildersqlBuilder=newStringBuilder("SELECT * FROM "+tableName+" ");
if (Objects.nonNull(aliasName)) {
sqlBuilder.append(aliasName).append(" ");
        }
Stringsql=sqlBuilder.append("WHERE ").append(conditions).toString();
SQLParsersqlParser=newSQLParser(sql);
returnsqlParser.parseQuery().getWhere();
    }
}

/*** @author leo_zl  2022/2/28*/@RunWith(PowerMockRunner.class)
@PowerMockIgnore({"javax.*"})
publicclassPhoenixSqlParserTest {
@InjectMocksprivatePhoenixSqlParserphoenixSqlParser;
@TestpublicvoidprivilegeControl() {
StringinputSql="SELECT full_name FROM SALES_PERSON WHERE ranking >= 5.0\n"+"    UNION ALL SELECT reviewer_name FROM CUSTOMER_REVIEW WHERE score >= 8.0\n"+"\n";
Mapconditions=newHashMap<>();
conditions.put("SALES_PERSON","id = 1");
StringresultSql=phoenixSqlParser.privilegeControl(inputSql, conditions);
StringexpectSql="SELECT FULL_NAME FROM SALES_PERSON  WHERE (RANKING >= 5.0 AND ID = 1) \n"+" UNION ALL \n"+"SELECT REVIEWER_NAME FROM CUSTOMER_REVIEW  WHERE SCORE >= 8.0";
Assert.assertEquals(expectSql,resultSql);
    }
}