使用mybatis插件完成数据库的读写分离

上篇文章完成了使用SpringAOP的方式的数据源读写分离.先看一张图

要点

Spring
Mybatis plugins
HikariCP

整体思路

使用一个代理数据源来管理其他的数据源，比如我有三个数据源，一主两从。那么我就生成了四个数据源，第四个数据源是前三个数据整合起来组成的，在Mybatis插件中进行判断使用哪个数据源。

为什么mybatis可以判断什么时候使用什么数据源呢?

因为Mybatis提供几种SqlCommandType(UNKNOWN, INSERT, UPDATE, DELETE, SELECT, FLUSH).可以根据这个进行判断。

上代码

pom文件内容

 <!-- 数据库支持 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>
<dependency>
    <groupId>org.mybatis.spring.boot</groupId>
    <artifactId>mybatis-spring-boot-starter</artifactId>
    <version>2.0.1</version>
</dependency>
<dependency>
    <groupId>mysql</groupId>
    <artifactId>mysql-connector-java</artifactId>
    <scope>runtime</scope>
</dependency>
<dependency>
    <groupId>com.baomidou</groupId>
    <artifactId>mybatis-plus-boot-starter</artifactId>
    <version>3.1.1</version>
</dependency>

<!--lombok-->
<dependency>
    <groupId>org.projectlombok</groupId>
    <artifactId>lombok</artifactId>
</dependency>

<!--自动配置类-->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-configuration-processor</artifactId>
    <optional>true</optional>
</dependency>

<!-- HikariCP 数据库连接池-->
<!-- https://mvnrepository.com/artifact/com.zaxxer/HikariCP -->
<dependency>
    <groupId>com.zaxxer</groupId>
    <artifactId>HikariCP</artifactId>
    <version>3.3.0</version>
</dependency>

自定义数据源属性

package com.cyou.ashura.player.datasource;

import com.zaxxer.hikari.HikariConfig;
import lombok.Getter;
import lombok.Setter;
import org.springframework.boot.context.properties.ConfigurationProperties;

import java.util.Map;
import java.util.TreeMap;

/**
 * 动态数据源属性
 *
 * @author fangjiaxiaobai@gmail.com
 * @since 2019-06-17
 */
@ConfigurationProperties(prefix = "ashura.dynamic")
@Setter
@Getter
public class DynamicDataSourceProperties {

    private String master = "master";

    private Map<String, HikariConfig> dataSources = new TreeMap<>();

    private String slavePrefix = "slave";
}

自动配置数据源类

package com.cyou.ashura.player.datasource;

import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * 代理动态数据源
 *
 * @author fangjiaxiaobai@gmail.com
 * @since 2019-06-17
 */
@Configuration
@EnableConfigurationProperties(DynamicDataSourceProperties.class)
@Slf4j
public class DynamicDataSourceProxyConfiguration {

    private DynamicDataSourceProperties dataSourceProperties;

    @Bean
    public Map<String, DataSource> initDataSource() {
        Map<String, HikariConfig> dataSources = dataSourceProperties.getDataSources();
        Map<String, DataSource> dataSourceMaps = new HashMap<>(dataSources.size());
        dataSources.forEach((key, value) -> {
            HikariDataSource hikariDataSource = new HikariDataSource(value);
            if (key.startsWith(dataSourceProperties.getSlavePrefix()) || key.startsWith(dataSourceProperties.getMaster())) {
                dataSourceMaps.put(key, hikariDataSource);
            } else {
                log.error("数据源名称不符合规定,必须以[{}/{}]开头", dataSourceProperties.getMaster(), dataSourceProperties.getSlavePrefix());
                throw new RuntimeException("数据源名称不符合规定");
            }
        });
        return dataSourceMaps;
    }

    @Bean
    public DataSource dynamicDataSource(@Autowired Map<String, DataSource> dataSourceMap) {

        if (dataSourceMap.isEmpty()) {
            throw new RuntimeException("未配置任何数据源");
        }

        Map<Object, Object> targetDataSources = new HashMap<>(dataSourceMap.size());
        boolean b = dataSourceMap.containsKey(dataSourceProperties.getMaster());
        if (!b) {
            // 不包含主数据源
            log.warn("未设置主数据源,将从其他数据源中随机选取一个设置为主数据源");
            Map.Entry<String, DataSource> stringDataSourceEntry = dataSourceMap.entrySet().stream().findFirst().get();
            DataSource masterDataSource = stringDataSourceEntry.getValue();
            targetDataSources.put("master", masterDataSource);
            dataSourceMap.entrySet().stream().skip(1).forEach(entry ->
                    targetDataSources.put(entry.getKey(), entry.getValue()));
        } else {
            dataSourceMap.entrySet().forEach(entry ->
                    targetDataSources.put(entry.getKey(), entry.getValue()));
        }

        DynamicRoutingDataSource dynamicDataSource = new DynamicRoutingDataSource();
        // 将写库设置为默认的数据源
        dynamicDataSource.setDefaultTargetDataSource(targetDataSources.get(dataSourceProperties.getMaster()));
        dynamicDataSource.setTargetDataSources(targetDataSources);
        return dynamicDataSource;
    }


    @Autowired
    public void setDataSourceProperties(DynamicDataSourceProperties dataSourceProperties) {
        this.dataSourceProperties = dataSourceProperties;
    }

}

动态数据源

package com.cyou.ashura.player.datasource;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 动态数据源
 *
 * @author fangjiaxiaobai@gmail.com
 * @since 2019-06-14
 */
@Slf4j
public class DynamicRoutingDataSource extends AbstractRoutingDataSource {

    private DynamicDataSourceProperties dataSourceProperties;

    private AtomicInteger count = new AtomicInteger(0);

    private static final List<String> SLAVE_DATA_SOURCE_KEYS = new ArrayList<>();

    @Override
    protected Object determineCurrentLookupKey() {
        DynamicDataSourceType dataSourceType = DynamicDataSourceHolder.getDataSource();
        String datasourceId;
        // 指定走从库，否则走主库
        if (dataSourceType == DynamicDataSourceType.Slave) {
            // 从库
            int size = dataSourceProperties.getDataSources().size() - 1;
            if (0 == size) {
                datasourceId = dataSourceProperties.getMaster();
            } else if (1 == size) {
                datasourceId = dataSourceProperties.getDataSources().keySet().parallelStream().findFirst().get();
            } else {
                int i = count.addAndGet(1);
                int index = i % size;
                if (SLAVE_DATA_SOURCE_KEYS.isEmpty()) {
                    List<String> keys = dataSourceProperties.getDataSources().keySet()
                            .parallelStream()
                            .filter(key -> !key.equals(dataSourceProperties.getMaster()))
                            .collect(Collectors.toList());
                    SLAVE_DATA_SOURCE_KEYS.addAll(keys);
                }
                datasourceId = SLAVE_DATA_SOURCE_KEYS.get(index);
            }
        } else {
            datasourceId = dataSourceProperties.getMaster();
        }
        log.debug("使用的数据源是:{}", datasourceId);
        DynamicDataSourceHolder.clearDataSource();
        return datasourceId;
    }

    @Autowired
    public void setDataSourceProperties(DynamicDataSourceProperties dataSourceProperties) {
        this.dataSourceProperties = dataSourceProperties;
    }

}

全局动态数据源的持有类

package com.cyou.ashura.player.datasource;

/**
 * 动态数据源的全局持有类
 *
 * @author fangjiaxiaobai@gmail.com
 * @since 2019-06-14
 */
public final class DynamicDataSourceHolder {


    private static final ThreadLocal<DynamicDataSourceType> HOLDER = new ThreadLocal<>();

    private DynamicDataSourceHolder() {

    }

    public static void setDataSource(DynamicDataSourceType dataSource) {
        HOLDER.set(dataSource);
    }

    static DynamicDataSourceType getDataSource() {
        return HOLDER.get();
    }

    static void clearDataSource() {
        HOLDER.remove();
    }
}

编写Mybatis插件

package com.cyou.ashura.player.datasource.plugins;

import com.cyou.ashura.player.datasource.DynamicDataSourceHolder;
import com.cyou.ashura.player.datasource.DynamicDataSourceType;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.keygen.SelectKeyGenerator;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import java.util.Locale;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 动态数据源的插件
 *
 * @author fangjiaxiaobai@gmail.com
 * @since 2019-06-17
 */
@Slf4j
@Intercepts({
        @Signature(type = Executor.class, method = "update", args = {
                MappedStatement.class, Object.class
        }),
        @Signature(type = Executor.class, method = "query", args = {
                MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class
        })
})
public class DynamicDataSourcePlugin implements Interceptor {

    private static final String REGEX = ".insert\\\\u0020.*|.*delete\\\\u0020.*|.*update\\\\u0020.*";

    /***
     * 缓存什么?
     */
    private static final Map<String, DynamicDataSourceType> CACHE_MAP = new ConcurrentHashMap<>();

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        boolean synchronizationActive = TransactionSynchronizationManager.isSynchronizationActive();

        if (!synchronizationActive) {
            Object[] args = invocation.getArgs();
            MappedStatement ms = (MappedStatement) args[0];
            DynamicDataSourceType dataSourceType;

            if ((dataSourceType = CACHE_MAP.get(ms.getId())) == null) {

                // 读方法
                if (ms.getSqlCommandType().equals(SqlCommandType.SELECT)) {
                    //!selectKey 为自增id查询主键(SELECT 1581623 )方法，使用主库
                    if (ms.getId().contains(SelectKeyGenerator.SELECT_KEY_SUFFIX)) {
                        dataSourceType = DynamicDataSourceType.Master;
                    } else {
                        BoundSql boundSql = ms.getSqlSource().getBoundSql(args[1]);
                        String sql = boundSql.getSql().toLowerCase(Locale.CANADA).replaceAll("[\\t\\n\\r]", "");
                        if (sql.matches(REGEX)) {
                            dataSourceType = DynamicDataSourceType.Master;
                        } else {
                            dataSourceType = DynamicDataSourceType.Slave;
                        }
                    }
                } else {
                    dataSourceType = DynamicDataSourceType.Master;
                }
                log.debug("设置方法[{}] 使用[{}]策略, Sql的命令类型是[{}]", ms.getId(), dataSourceType.name(),
                        ms.getSqlCommandType().name());
                CACHE_MAP.put(ms.getId(), dataSourceType);
            }
            DynamicDataSourceHolder.setDataSource(dataSourceType);
        }
        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof Executor) {
            return Plugin.wrap(target, this);
        } else {
            return target;
        }
    }

    @Override
    public void setProperties(Properties properties) {

    }
}

配置mybatis插件

<plugins>
    <plugin interceptor="com.cyou.ashura.player.datasource.plugins.DynamicDataSourcePlugin"/>
</plugins>

最后一步,配置yml属性

ashura:
  dynamic:
    datasources:
      master:
        username: mysql
        jdbc-url: jdbc:mysql://********:3306/ashura?serverTimezone=GMT%2B8&useUnicode=true&characterEncoding=UTF-8&autoReconnect=true&useSSL=false&zeroDateTimeBehavior=convertToNull
        password: password
        minimum-idle: 5
        maximum-pool-size: 15
        auto-commit: true
        idle-timeout: 3000
        pool-name: player-dataSource-pool
        max-lifetime: 1800000
        connection-timeout: 30000
        connection-test-query: SELECT 1
      slave1:
        username: mysql
        jdbc-url: jdbc:mysql://********:3306/ashura_slave1?serverTimezone=GMT%2B8&useUnicode=true&characterEncoding=UTF-8&autoReconnect=true&useSSL=false&zeroDateTimeBehavior=convertToNull
        password: password
        minimum-idle: 5
        maximum-pool-size: 15
        auto-commit: true
        idle-timeout: 3000
        pool-name: player-dataSource-pool
        max-lifetime: 1800000
        connection-timeout: 30000
        connection-test-query: SELECT 1
      slave2:
        username: mysql
        jdbc-url: jdbc:mysql://********:3306/ashura_slave2?serverTimezone=GMT%2B8&useUnicode=true&characterEncoding=UTF-8&autoReconnect=true&useSSL=false&zeroDateTimeBehavior=convertToNull
        password: password
        minimum-idle: 5
        maximum-pool-size: 15
        auto-commit: true
        idle-timeout: 3000
        pool-name: player-dataSource-pool
        max-lifetime: 1800000
        connection-timeout: 30000
        connection-test-query: SELECT 1