Mybatis 拦截器实现单数据源内多数据库切换 | 京东物流技术团队

2024-01-08 10:09 由京东云开发者发表于 #后端开发

物流的分拣业务在某些分拣场地只有一个数据源，因为数据量比较大，将所有数据存在一张表内查询速度慢，也为了做不同设备数据的分库管理，便在这个数据源内创建了多个不同库名但表完全相同的数据库，如下图所示：

现在需要上线报表服务来查询所有数据库中的数据进行统计，那么现在的问题来了，该如何 满足在配置一个数据源的情况下来查询该数据源下不同数据库的数据 呢，借助搜索引擎查到的分库实现大多是借助 Sharding-JDBC 框架，配置多个数据源根据分库算法实现数据源的切换，但是对于只有一个数据源的系统来说，我觉得引入框架再将单个数据源根据不同的库名配置成多个不同的数据源来实现分库查询的逻辑我觉得并不好。

如果我们能在 SQL 执行前将 SQL 中所有的表名前拼接上对应的库名的话，那么就能够实现数据源的切换了，下面我们讲一下使用 JSqlParser 和 Mybatis拦截器实现该逻辑，借助 JSqlParser 主要是为了解析SQL，找到其中所有的表名进行拼接，如果大家有更好的实现方式，该组件并不是必须的。

实现逻辑

SqlSource 是读取 XML 中 SQL 内容并将其发送给数据库执行的对象，如果我们在执行前能拦截到该对象，并将其中的 SQL 替换掉便达成了我们的目的。 SqlSource 有多种实现，包括常见的DynamicSqlSource。其中包含着必要的执行逻辑，我们需要做的工作便是在这些逻辑执行完之后，对 SQL 进行改造，所以这次实现我们使用了 装饰器模式，在原来的 SqlSource 上套一层，执行完 SqlSource 本身的方法之后对其进行增强，代码如下：

public abstract class AbstractDBNameInterceptor {

    /**
     * SqlSource 的装饰器，作用是增强了 getBoundSql 方法，在基础上增加了动态分库的逻辑
     */
    static class SqlSourceDecorator implements SqlSource {

        /**
         * SQL 字段名称
         */
        private static final String SQL_FIELD_NAME = "sql";

        /**
         * 原本的 sql source
         */
        private final SqlSource sqlSource;

        /**
         * 装饰器进行封装
         */
        public SqlSourceDecorator(SqlSource sqlSource) {
            this.sqlSource = sqlSource;
        }

        @Override
        public BoundSql getBoundSql(Object parameterObject) {
            try {
                // 先生成出未修改前的 SQL
                BoundSql boundSql = sqlSource.getBoundSql(parameterObject);
                // 获取数据库名
                String dbName = getSpecificDBName(parameterObject);
                // 有效才修改
                if (isValid(dbName)) {
                    // 生成需要修改完库名的 SQL
                    String targetSQL = getRequiredSqlWithSpecificDBName(boundSql, dbName);
                    // 更新 SQL
                    updateSql(boundSql, targetSQL);
                }

                return boundSql;
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }

        /**
         * 校验是否为有效库名
         */
        private boolean isValid(String dbName) {
            return StringUtils.isNotEmpty(dbName) !"null".equals(dbName);
        }

        /**
         * 获取到我们想要的库名的 SQL
         */
        private String getRequiredSqlWithSpecificDBName(BoundSql boundSql, String dbName) throws JSQLParserException {
            String originSql = boundSql.getSql();
            // 获取所有的表名
            Set<String> tables = TablesNamesFinder.findTables(originSql);
            for (String table : tables) {
                originSql = originSql.replaceAll(table, dbName + "." + table);
            }
            return originSql;
        }

        /**
         * 修改 SQL
         */
        private void updateSql(BoundSql boundSql, String sql) throws NoSuchFieldException, IllegalAccessException {
            // 通过反射修改sql语句
            Field field = boundSql.getClass().getDeclaredField(SQL_FIELD_NAME);
            field.setAccessible(true);
            field.set(boundSql, sql);
        }
    }
    
    // ... 
}

定义了 AbstractDBNameInterceptor 抽象类是为了实现复用，并将 SqlSourceDecorator 装饰器定义为静态内部类，这样的话，将所有逻辑都封装在抽象类内部，之后这部分实现好后研发直接实现抽象类的通用方法即可，不必关注它的内部实现。

结合注释我们解释一下 SqlSourceDecorator 的逻辑，其中用到了 Java 反射相关的操作。首先通过反射获取到 SQL，getSpecificDBName 方法是需要自定义实现的，其中 parameterObject 对象是传到 DAO 层执行查询时的参数，在我们的业务中是能够根据其中的设备相关参数拿到对应的所在库名的，而设备和具体库名的映射关系需要提前初始化好。在获取到具体的库名后执行 getRequiredSqlWithSpecificDBName 方法来将其拼接到表名前，在这里我们使用到了 JSqlParser 的工具类，解析出来所有的表名，执行字符串的替换，最后一步同样是使用反射操作将该参数值再写回去，这样便完成了指定库名的任务。

接下来我们需要看下抽象拦截器中供拦截器复用的方法，如下：

public abstract class AbstractDBNameInterceptor {

    /**
     * SqlSource 字段名称
     */
    private static final String SQL_SOURCE_FIELD_NAME = "sqlSource";

    /**
     * 执行修改数据库名的逻辑
     */
    protected Object updateDBName(Invocation invocation) throws Throwable {
        // 装饰器装饰 SqlSource
        decorateSqlSource((MappedStatement) invocation.getArgs()[0]);
        return invocation.proceed();
    }

    /**
     * 装饰 SqlSource
     */
    private void decorateSqlSource(MappedStatement statement) throws NoSuchFieldException, IllegalAccessException {
        if (!(statement.getSqlSource() instanceof SqlSourceDecorator)) {
            Field sqlSource = statement.getClass().getDeclaredField(SQL_SOURCE_FIELD_NAME);
            sqlSource.setAccessible(true);
            sqlSource.set(statement, new SqlSourceDecorator(statement.getSqlSource()));
        }
    }
}

这个还是比较简单的，只是借助反射机制做了一层“装饰”，查询拦截器实现如下：

@Intercepts({
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class})
})
public class SelectDBNameInterceptor extends AbstractDBNameInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        return updateDBName(invocation);
    }
}

将其配置到 Mybatis 拦截器中，便能实现数据库动态切换了。

作者：京东物流王奕龙

来源：京东云开发者社区自猿其说 Tech 转载请注明来源