基於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切換為需要切換的數據源。