【基于druid 支持select sql库名及表名映射】

hackcoder

于 2023-09-26 11:35:38 发布

阅读量193

点赞数 1

文章标签： sql 数据库

本文链接：https://blog.csdn.net/hackcoder/article/details/133302310

版权

基于druid 支持select sql库名及表名映射

package org.example;

import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.ast.statement.SQLSelectItem;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitorAdapter;

import java.util.HashMap;
import java.util.List;
import java.util.Map;


public class DruidTest {

    public static class MySqlExportTableAliasVisitor extends MySqlASTVisitorAdapter {
        Map<String, String> schemaMapping = new HashMap<>();
        Map<String, String> tableMapping = new HashMap<>();

        public MySqlExportTableAliasVisitor(Map<String, String> schemaMapping, Map<String, String> tableMapping) {
            this.schemaMapping = schemaMapping;
            this.tableMapping = tableMapping;
        }

        @Override
        public boolean visit(MySqlSelectQueryBlock x) {
            for (SQLSelectItem selectItem : x.getSelectList()) {
                SQLExpr sqlExpr = ((SQLPropertyExpr) selectItem.getExpr()).getOwner();

                if (sqlExpr instanceof SQLPropertyExpr) {
                    String tableName = ((SQLPropertyExpr) sqlExpr).getName();
                    ((SQLPropertyExpr) sqlExpr).setName(tableMapping.getOrDefault(tableName, tableName));

                    SQLIdentifierExpr schema = (SQLIdentifierExpr) ((SQLPropertyExpr) sqlExpr).getOwner();
                    if (schema != null) {
                        schema.setName(schemaMapping.getOrDefault(schema.getName(), schema.getName()));
                    }
                } else if (sqlExpr instanceof SQLIdentifierExpr) {
                    SQLIdentifierExpr sqlIdentifierExpr = ((SQLIdentifierExpr) sqlExpr);
                    sqlIdentifierExpr.setName(tableMapping.getOrDefault(sqlIdentifierExpr.getName(), sqlIdentifierExpr.getName()));
                }
            }
            return true;
        }

        @Override
        public boolean visit(SQLExprTableSource x) {
            x.setSchema(schemaMapping.getOrDefault(x.getSchema(), x.getSchema()));
            if (x.getExpr() instanceof SQLPropertyExpr) {
                SQLPropertyExpr sqlPropertyExpr = (SQLPropertyExpr) x.getExpr();
                sqlPropertyExpr.setName(tableMapping.getOrDefault(x.getTableName(), x.getTableName()));
            } else if (x.getExpr() instanceof SQLIdentifierExpr) {
                SQLIdentifierExpr sqlIdentifierExpr = (SQLIdentifierExpr) x.getExpr();
                sqlIdentifierExpr.setName(tableMapping.getOrDefault(x.getTableName(), x.getTableName()));
            }
            return true;
        }
    }


    public static void main(String[] args) {
//        String sql = "select `test_01`.`id` AS `id`,`test_01`.`a` AS `a` from `test_01`";
        String sql = "select `hy_test_0`.`test_01`.`id` AS `id`,`hy_test_0`.`test_01`.`a` AS `a` from `hy_test_0`.`test_01`";

        Map<String, String> schemaMapping = new HashMap<>();
        schemaMapping.put("`hy_test_0`", "`hy_test_1`");

        Map<String, String> tableMapping = new HashMap<>();
        tableMapping.put("`test_01`", "`new_test_01`");
        List<SQLStatement> stmtList = SQLUtils.parseStatements(sql, DbType.mysql);
        MySqlExportTableAliasVisitor visitor = new MySqlExportTableAliasVisitor(schemaMapping, tableMapping);
        for (SQLStatement stmt : stmtList) {
            stmt.accept(visitor);
        }
        String handleSQL = SQLUtils.toSQLString(stmtList, DbType.mysql);

        System.out.println(handleSQL);
    }
}