基于AbstractRoutingDataSource实现动态数据源

  在分库分表的时候用到了sharding-jdbc实现，有时候需要实现动态切换数据源。比如根据用户的公司，每个公司分配不同的DB。spring-jdbc的包org.springframework.jdbc.datasource.lookup提供了AbstractRoutingDataSource，可以实现动态数据源切换。MybatisPlus也有提供的多数据源，简单了解下是基于注解切换数据源。

　　下面的动态数据源切换方式是基于ThreadLocal进行切换，如果需要注解方式实现的话可以用AOP自定义注解来实现。

　　使用的数据库连接池是springboot2自带的HikariCP。

1.  查看AbstractRoutingDataSource类源码入下

/** <a href="http://www.cpupk.com/decompiler">Eclipse Class Decompiler</a> plugin, Copyright (c) 2017 Chen Chao. */
package org.springframework.jdbc.datasource.lookup;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;

import javax.sql.DataSource;

import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.datasource.AbstractDataSource;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;

public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {

    @Nullable
    private Map<Object, Object> targetDataSources;

    @Nullable
    private Object defaultTargetDataSource;

    private boolean lenientFallback = true;

    private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();

    @Nullable
    private Map<Object, DataSource> resolvedDataSources;

    @Nullable
    private DataSource resolvedDefaultDataSource;


    /**
     * Specify the map of target DataSources, with the lookup key as key.
     * The mapped value can either be a corresponding {@link javax.sql.DataSource}
     * instance or a data source name String (to be resolved via a
     * {@link #setDataSourceLookup DataSourceLookup}).
     * <p>The key can be of arbitrary type; this class implements the
     * generic lookup process only. The concrete key representation will
     * be handled by {@link #resolveSpecifiedLookupKey(Object)} and
     * {@link #determineCurrentLookupKey()}.
     */
    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        this.targetDataSources = targetDataSources;
    }

    /**
     * Specify the default target DataSource, if any.
     * <p>The mapped value can either be a corresponding {@link javax.sql.DataSource}
     * instance or a data source name String (to be resolved via a
     * {@link #setDataSourceLookup DataSourceLookup}).
     * <p>This DataSource will be used as target if none of the keyed
     * {@link #setTargetDataSources targetDataSources} match the
     * {@link #determineCurrentLookupKey()} current lookup key.
     */
    public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
        this.defaultTargetDataSource = defaultTargetDataSource;
    }

    /**
     * Specify whether to apply a lenient fallback to the default DataSource
     * if no specific DataSource could be found for the current lookup key.
     * <p>Default is "true", accepting lookup keys without a corresponding entry
     * in the target DataSource map - simply falling back to the default DataSource
     * in that case.
     * <p>Switch this flag to "false" if you would prefer the fallback to only apply
     * if the lookup key was {@code null}. Lookup keys without a DataSource
     * entry will then lead to an IllegalStateException.
     * @see #setTargetDataSources
     * @see #setDefaultTargetDataSource
     * @see #determineCurrentLookupKey()
     */
    public void setLenientFallback(boolean lenientFallback) {
        this.lenientFallback = lenientFallback;
    }

    /**
     * Set the DataSourceLookup implementation to use for resolving data source
     * name Strings in the {@link #setTargetDataSources targetDataSources} map.
     * <p>Default is a {@link JndiDataSourceLookup}, allowing the JNDI names
     * of application server DataSources to be specified directly.
     */
    public void setDataSourceLookup(@Nullable DataSourceLookup dataSourceLookup) {
        this.dataSourceLookup = (dataSourceLookup != null ? dataSourceLookup : new JndiDataSourceLookup());
    }


    @Override
    public void afterPropertiesSet() {
        if (this.targetDataSources == null) {
            throw new IllegalArgumentException("Property 'targetDataSources' is required");
        }
        this.resolvedDataSources = new HashMap<>(this.targetDataSources.size());
        this.targetDataSources.forEach((key, value) -> {
            Object lookupKey = resolveSpecifiedLookupKey(key);
            DataSource dataSource = resolveSpecifiedDataSource(value);
            this.resolvedDataSources.put(lookupKey, dataSource);
        });
        if (this.defaultTargetDataSource != null) {
            this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
        }
    }

    /**
     * Resolve the given lookup key object, as specified in the
     * {@link #setTargetDataSources targetDataSources} map, into
     * the actual lookup key to be used for matching with the
     * {@link #determineCurrentLookupKey() current lookup key}.
     * <p>The default implementation simply returns the given key as-is.
     * @param lookupKey the lookup key object as specified by the user
     * @return the lookup key as needed for matching
     */
    protected Object resolveSpecifiedLookupKey(Object lookupKey) {
        return lookupKey;
    }

    /**
     * Resolve the specified data source object into a DataSource instance.
     * <p>The default implementation handles DataSource instances and data source
     * names (to be resolved via a {@link #setDataSourceLookup DataSourceLookup}).
     * @param dataSource the data source value object as specified in the
     * {@link #setTargetDataSources targetDataSources} map
     * @return the resolved DataSource (never {@code null})
     * @throws IllegalArgumentException in case of an unsupported value type
     */
    protected DataSource resolveSpecifiedDataSource(Object dataSource) throws IllegalArgumentException {
        if (dataSource instanceof DataSource) {
            return (DataSource) dataSource;
        }
        else if (dataSource instanceof String) {
            return this.dataSourceLookup.getDataSource((String) dataSource);
        }
        else {
            throw new IllegalArgumentException(
                    "Illegal data source value - only [javax.sql.DataSource] and String supported: " + dataSource);
        }
    }


    @Override
    public Connection getConnection() throws SQLException {
        return determineTargetDataSource().getConnection();
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return determineTargetDataSource().getConnection(username, password);
    }

    @Override
    @SuppressWarnings("unchecked")
    public <T> T unwrap(Class<T> iface) throws SQLException {
        if (iface.isInstance(this)) {
            return (T) this;
        }
        return determineTargetDataSource().unwrap(iface);
    }

    @Override
    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return (iface.isInstance(this) || determineTargetDataSource().isWrapperFor(iface));
    }

    /**
     * Retrieve the current target DataSource. Determines the
     * {@link #determineCurrentLookupKey() current lookup key}, performs
     * a lookup in the {@link #setTargetDataSources targetDataSources} map,
     * falls back to the specified
     * {@link #setDefaultTargetDataSource default target DataSource} if necessary.
     * @see #determineCurrentLookupKey()
     */
    protected DataSource determineTargetDataSource() {
        Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
        Object lookupKey = determineCurrentLookupKey();
        DataSource dataSource = this.resolvedDataSources.get(lookupKey);
        if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
            dataSource = this.resolvedDefaultDataSource;
        }
        if (dataSource == null) {
            throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
        }
        return dataSource;
    }

    /**
     * Determine the current lookup key. This will typically be
     * implemented to check a thread-bound transaction context.
     * <p>Allows for arbitrary keys. The returned key needs
     * to match the stored lookup key type, as resolved by the
     * {@link #resolveSpecifiedLookupKey} method.
     */
    @Nullable
    protected abstract Object determineCurrentLookupKey();

}

　　determineTargetDataSource()方法决定采用的数据源，方法内部调用determineCurrentLookupKey()方法获取对应的datasource的key，需要子类实现该方法，因此子类需要做的就是维护targetDataSources以及重写determineCurrentLookupKey方法决定采用的key。

2.简单实现

1.动态数据源类

package cn.qlq.config.datasource;

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

import javax.sql.DataSource;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

public class DynamicDataSource extends AbstractRoutingDataSource {

    public static Map<Object, Object> dynamicTargetDataSources = new HashMap<>();

    private static DynamicDataSource dynamicDataSource = null;

    private DynamicDataSource() {

    }

    public static synchronized DynamicDataSource getInstance() {
        if (dynamicDataSource == null) {
            dynamicDataSource = new DynamicDataSource();
        }

        return dynamicDataSource;
    }

    /**
     * 如果不希望数据源在启动配置时就加载好，可以定制这个方法，从任何你希望的地方读取并返回数据源
     * 比如从数据库、文件、外部接口等读取数据源信息，并最终返回一个DataSource实现类对象即可
     * 
     * @return
     */
    @Override
    protected DataSource determineTargetDataSource() {
        return super.determineTargetDataSource();
    }

    /**
     * 如果希望所有数据源在启动配置时就加载好，这里通过设置数据源Key值来切换数据，定制这个方法
     * 
     * @return
     */
    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.getDataSourceKey();
    }

    @Override
    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        super.setTargetDataSources(targetDataSources);
        this.dynamicTargetDataSources = targetDataSources;
    }

    /**
     * 新增数据源
     * 
     * @param key
     *            数据源标识
     * @param dataSource
     *            数据源
     */
    public void addTargetDataSources(Object key, Object dataSource) {
        dynamicTargetDataSources.put(key, dataSource);
        super.setTargetDataSources(dynamicTargetDataSources);
        super.afterPropertiesSet();
    }

    /**
     * 设置默认数据源
     * 
     * @param defaultDataSource
     */
    public void setDefaultDataSource(Object defaultDataSource) {
        super.setDefaultTargetDataSource(defaultDataSource);
    }

    public void setDataSources(Map<Object, Object> dataSources) {
        setTargetDataSources(dataSources);
    }
}

 　　determineCurrentLookupKey方法决定数据源是调用DynamicDataSourceContextHolder.getDataSourceKey()获取当前线程的数据源key。

2.DynamicDataSourceContextHolder 类用ThreadLocal记录当前线程的数据源以及切换数据源

package cn.qlq.config.datasource;

import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class DynamicDataSourceContextHolder {
    private static final Logger log = LoggerFactory.getLogger(DynamicDataSourceContextHolder.class);

    private static final ThreadLocal<String> contextHolder = new ThreadLocal<String>() {
        /**
         * 将 master 数据源的 key作为默认数据源的 key
         */
        @Override
        protected String initialValue() {
            return "master";
        }
    };

    /**
     * 切换数据源
     * 
     * @param key
     *            数据源
     */
    public static void setDataSourceKey(String key) {
        if (!StringUtils.isEmpty(key)) {
            // 当数据源不存在时，添加数据源
            DatabaseUtils.addTargetDataSources(key);

            contextHolder.set(key);
        }
    }

    /**
     * 获取数据源
     * 
     * @return
     */
    public static String getDataSourceKey() {
        return contextHolder.get();
    }

    /**
     * 重置数据源
     */
    public static void clearDataSourceKey() {
        contextHolder.remove();
    }
}

 　　使用了ThreadLocal来实现数据在线程中的传输。

3.DatabaseUtils 工具类可以增加数据源

package cn.qlq.config.datasource;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Map;
import java.util.Set;

import com.zaxxer.hikari.HikariDataSource;

public class DatabaseUtils {

    static final String JDBC_DRIVER = "com.mysql.jdbc.Driver";
    static final String DB_URL = "jdbc:mysql://localhost:3306/";

    // Database credentials
    static final String USER = "root";
    static final String PASS = "123456";

    public synchronized static void addTargetDataSources(String database) {
        DynamicDataSource dynamicDataSource = DynamicDataSource.getInstance();
        Map<Object, Object> dynamicTargetDataSources = DynamicDataSource.dynamicTargetDataSources;
        Set<Object> keySet = dynamicTargetDataSources.keySet();
        if (!keySet.contains(database)) {
            HikariDataSource druidDataSource = new HikariDataSource();
            druidDataSource.setUsername(DB_URL + database
                    + "?useUnicode=true&amp;characterEncoding=utf8&amp;autoReconnect=true&amp;autoReconnectForPools=true&amp;failOverReadOnly=false");
            druidDataSource.setUsername(USER);
            druidDataSource.setPassword(PASS);
            druidDataSource.setDriverClassName(JDBC_DRIVER);
            dynamicDataSource.addTargetDataSources(database, druidDataSource);
        }
    }

    public static void createDatabase(String database) {
        Connection conn = null;
        Statement stmt = null;
        try {
            Class.forName(JDBC_DRIVER);
            conn = DriverManager.getConnection(DB_URL, USER, PASS);
            stmt = conn.createStatement();
            String sql = "CREATE DATABASE IF NOT EXISTS " + database
                    + " DEFAULT CHARACTER SET utf8 DEFAULT COLLATE utf8_general_ci";
            stmt.executeUpdate(sql);

            // 如果需要执行一些数据库初始化的脚本可以放到这里

        } catch (SQLException se) {
            se.printStackTrace();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                if (stmt != null)
                    stmt.close();
            } catch (SQLException se2) {
                try {
                    if (conn != null)
                        conn.close();
                } catch (SQLException se) {
                    se.printStackTrace();
                }
            }
        }
    }
}

4. Mybatis 会话工厂以及数据源的配置

package cn.qlq.mybatisplus.config;

import java.util.HashMap;
import java.util.Map;
import java.util.Properties;

import javax.sql.DataSource;

import org.apache.commons.lang3.ArrayUtils;
import org.apache.ibatis.plugin.Interceptor;
import org.mybatis.spring.annotation.MapperScan;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;

import com.baomidou.mybatisplus.extension.spring.MybatisSqlSessionFactoryBean;
import com.github.pagehelper.PageInterceptor;
import com.zaxxer.hikari.HikariDataSource;

import cn.qlq.config.datasource.DynamicDataSource;

@Configuration
@MapperScan({ "cn.qlq.mapper", "cn.qlq.sharedjdbc.client", "cn.qlq.mybatisplus" })
public class MybatisConfiguration {

    private static final Logger log = LoggerFactory.getLogger(MybatisConfiguration.class);

    @Bean("master")
    @Primary
    @ConfigurationProperties(prefix = "spring.datasource")
    public DataSource master() {
        HikariDataSource druidDataSource = new HikariDataSource();
        return druidDataSource;
    }

    @Bean("slave")
    @ConfigurationProperties(prefix = "spring.datasource.slave")
    public DataSource slave() {
        // return DataSourceBuilder.create().build();
        HikariDataSource druidDataSource = new HikariDataSource();
        return druidDataSource;
    }

    @Bean("dynamicDataSource")
    public DataSource dynamicDataSource() {
        DynamicDataSource dynamicDataSource = DynamicDataSource.getInstance();
        Map<Object, Object> dataSourceMap = new HashMap<>();
        dataSourceMap.put("master", master());
        dataSourceMap.put("slave", slave());

        log.info("dynamicDataSource dataSourceMap:{}", dataSourceMap);
        // 将 master 数据源作为默认指定的数据源
        dynamicDataSource.setDefaultDataSource(master());
        // 将 master 和 slave 数据源作为指定的数据源
        dynamicDataSource.setDataSources(dataSourceMap);
        return dynamicDataSource;
    }

    @Bean
    public MybatisSqlSessionFactoryBean sqlSessionFactoryBean() throws Exception {
        MybatisSqlSessionFactoryBean sessionFactory = new MybatisSqlSessionFactoryBean();
        // 配置数据源，此处配置为关键配置，如果没有将 dynamicDataSource作为数据源则不能实现切换
        sessionFactory.setDataSource(dynamicDataSource());

        // 扫描Model
        // sessionFactory.setTypeAliasesPackage("cn.qlq");
        PathMatchingResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
        // 扫描映射文件
        sessionFactory.setMapperLocations(resolver.getResources("classpath*:mapper/**/*Mapper.xml"));

        // 添加插件
        Interceptor[] interceptors = getPlugins();
        if (ArrayUtils.isNotEmpty(interceptors)) {
            sessionFactory.setPlugins(interceptors);
        }

        return sessionFactory;
    }

    private Interceptor[] getPlugins() {
        Interceptor[] plugins = new Interceptor[0];

        // PageHelper分页插件
        PageInterceptor pageInterceptor = new PageInterceptor();
        Properties properties = new Properties();
        properties.setProperty("helperDialect", "mysql");
        properties.setProperty("reasonable", "true");
        pageInterceptor.setProperties(properties);

        plugins = ArrayUtils.add(plugins, pageInterceptor);
        return plugins;
    }
}

5.application.properties中配置：

############################################################
#
# datasource settings
#
############################################################
spring.datasource.driver-class-name= com.mysql.jdbc.Driver
spring.datasource.jdbc-url = jdbc:mysql://localhost:3306/test1?useUnicode=true&characterEncoding=utf-8
spring.datasource.username = root
spring.datasource.password = 123456

#slave
spring.datasource.slave.driver-class-name= com.mysql.jdbc.Driver
spring.datasource.slave.jdbc-url = jdbc:mysql://localhost:3306/test2?useUnicode=true&characterEncoding=utf8&autoReconnect=true&autoReconnectForPools=true&failOverReadOnly=false
spring.datasource.slave.username = root
spring.datasource.slave.password = 123456

6.测试：

import java.util.List;

import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import cn.qlq.MySpringBootApplication;
import cn.qlq.config.datasource.DynamicDataSourceContextHolder;
import cn.qlq.mybatisplus.bean.MpUser;
import cn.qlq.mybatisplus.bean.MpUserMapper;

@RunWith(SpringRunner.class)
@SpringBootTest(classes = MySpringBootApplication.class)
public class PlainTest {

    @Autowired
    private MpUserMapper mpUserMapper;

    @Test
    public void findAll() {
        // 切换数据源，默认是master数据源
        DynamicDataSourceContextHolder.setDataSourceKey("slave");
        List<MpUser> MpUser = mpUserMapper.selectList(null);
        System.out.println(MpUser);
    }
}

总结：

(1)常见的可以是代码中切换数据源，也可以是在过滤器中根据当前用户的身份标识进行全局切换。

(2)DynamicDataSourceContextHolder.setDataSourceKey()切换数据源过程如下：

调用DatabaseUtils.addTargetDataSources(key);， 方法中在未包含需要切换的数据源的时候会尝试建立连接数据源并存到DynamicDataSource

调用contextHolder.set(key);将ThreadLocal中的key切换为需要切换的数据源。