Spring Boot 2.x动态数据源切换

背景

有一个使用C#的老项目，原来设计的时候是根据一个区域一个数据库的设计，现在需要将这个项目重构成Java的项目。并且新增动态添加区域，不用重启数据库的需求。

PS：原来添加区域步骤
1.手动新建数据库
2.手动录入该区域的静态数据
3.手动将数据库链接复制到主库
4.重新运行程序

需求

1.可以动态的通过接口添加数据库
2.可以根据规则动态的切换数据源

解决方案

Spring boot提供了AbstractRoutingDataSource 根据用户定义的规则选择当前的数据源，这样我们可以在执行查询之前，设置使用的数据源。而切换的数据源存在于targetDataSources对象中，并且提供了setTargetDataSources让我们来动态设置数据源。所以解决方案就有了，

1.新建一个单例的动态数据源类

继承AbstractRoutingDataSource，并且重写其determineCurrentLookupKey（切换时指定数据源）和setTargetDataSources方法。新增createDataSource来动态添加数据源，linkDatasource判断数据源是否可用。

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.fastjson.JSONObject;
import org.omg.PortableInterceptor.SYSTEM_EXCEPTION;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import rfpatrolsystem.entity.Substation;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.sql.Connection;
import java.sql.DriverManager;
import java.util.Map;

/**
 * 动态数据源
 *
 * @author admin
 */
public class DynamicDataSource extends AbstractRoutingDataSource {

    private static DynamicDataSource instance;

    static Map<Object, Object> dynamicTargetDataSources;

    public static DynamicDataSource getInstance() {
        if (instance == null) {
            synchronized (DynamicDataSource.class) {
                if (instance == null) {
                    instance = new DynamicDataSource();
                }
            }
        }
        return instance;
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return DbContextHolder.getDataSourceKey();
    }

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

    /**
     * 创建数据源
     *
     * @param key        数据源名称
     * @param driveClass 数据源类型
     * @param url        数据源连接字符串
     * @param username   用户名
     * @param password   密码
     * @return 添加是否成功
     */
    public boolean createDataSource(String key, String driveClass, String url, String username, String password) {
        try {
            if (!linkDatasource(driveClass, url, username, password)) {
                return false;
            }
            DruidDataSource druidDataSource = new DruidDataSource();
            druidDataSource.setName(key);
            druidDataSource.setDriverClassName(driveClass);
            druidDataSource.setUrl(url);
            druidDataSource.setUsername(username);
            druidDataSource.setPassword(password);
            druidDataSource.setMaxWait(60000);
            druidDataSource.setFilters("stat");
            druidDataSource.init();
            Map<Object, Object> dynamicTargetDataSources2 = dynamicTargetDataSources;
            // 加入map
            dynamicTargetDataSources2.put(key, druidDataSource);
            // 将map赋值给父类的TargetDataSources
            this.setTargetDataSources(dynamicTargetDataSources2);
            // 将TargetDataSources中的连接信息放入resolvedDataSources管理
            super.afterPropertiesSet();
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 判断数据库是否能连接
     *
     * @param driveClass 连接驱动
     * @param url        连接地址
     * @param username   用户名
     * @param password   密码
     * @return 连接是否成功
     */
    private boolean linkDatasource(String driveClass, String url, String username, String password) {
        Connection connection = null;
        try {
            // 排除连接不上的错误
            Class.forName(driveClass);
            //
            DriverManager.setLoginTimeout(1);
            // 相当于连接数据库
            connection = DriverManager.getConnection(url, username, password);
        } catch (Exception e) {
            return false;
        } finally {
            if (connection != null) {
                try {
                    connection.close();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
        return true;
    }

}

2.动态数据源切换，使用ThreadLocal和线程绑定

/**
 * 动态数据源
 *
 * @author admin
 */
public class DbContextHolder {

    /**
     * 绑定当前线程
     */
    private static ThreadLocal<String> dataSourceKey = new ThreadLocal<>();


    /**
     * 设置当前线程使用的数据源key
     *
     * @param dataSourceType 数据源key
     */
    public static void setDataSourceKey(String dataSourceType) throws BusinessException {
        // 数据源中没有这个key
        if (DynamicDataSource.dynamicTargetDataSources == null || DynamicDataSource.dynamicTargetDataSources.get(dataSourceType) == null) {
            throw new BusinessException(ExceptionMsg.DATABASE_ERROR);
        }
        dataSourceKey.set(dataSourceType);
    }

    /**
     * 获取当前线程使用的数据源key
     *
     * @return 数据源key
     */
    static String getDataSourceKey() {
        return dataSourceKey.get();
    }

    /**
     * 清楚数据源
     */
    static void clearDataSourceKey() {
        dataSourceKey.remove();
    }

}

3.数据源配置

import com.alibaba.druid.pool.DruidDataSource;
import org.apache.ibatis.session.SqlSessionFactory;
import org.mybatis.spring.SqlSessionFactoryBean;
import org.mybatis.spring.SqlSessionTemplate;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.env.Environment;

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

/**
 * 数据源配置
 *
 * @author admin
 */
@Configuration
public class DataSourceConfig {

    @Autowired
    private Environment env;

    @Value("${spring.datasource.names}")
    private String datasourceNames;

    /**
     * 默认数据源
     *
     * @return 默认数据源
     */
    private DataSource getDefaultDataSource() {
        DruidDataSource defaultDS = new DruidDataSource();
        defaultDS.setUrl(env.getProperty("spring.datasource.base.url"));
        defaultDS.setUsername(env.getProperty("spring.datasource.base.username"));
        defaultDS.setPassword(env.getProperty("spring.datasource.base.password"));
        defaultDS.setDriverClassName(env.getProperty("spring.datasource.base.driver-class-name"));
        return defaultDS;
    }

    /**
     * 已经配置的动态数据源。
     */
    private Map<Object, Object> getDataSources() {
        Map<Object, Object> map = new HashMap<>(2);
        if (datasourceNames != null && datasourceNames.length() > 0) {
            String[] names = datasourceNames.split(",");
            for (String name : names) {
                DruidDataSource dataSource = new DruidDataSource();
                dataSource.setUrl(env.getProperty("spring.datasource." + name + ".url"));
                dataSource.setUsername(env.getProperty("spring.datasource." + name + ".username"));
                dataSource.setPassword(env.getProperty("spring.datasource." + name + ".password"));
                dataSource.setDriverClassName(env.getProperty("spring.datasource." + name + ".driver-class-name"));
                dataSource.setBreakAfterAcquireFailure(true);
                dataSource.setConnectionErrorRetryAttempts(0);
                map.put(name, dataSource);
            }
        }
        return map;
    }

    @Bean
    public DynamicDataSource dynamicDataSource() {
        DynamicDataSource dynamicDataSource = DynamicDataSource.getInstance();
        Map<Object, Object> dataSources = getDataSources();
        if (dataSources.size() > 0) {
            dynamicDataSource.setTargetDataSources(dataSources);
        }
        DataSource ds = getDefaultDataSource();
        if (ds != null) {
            dynamicDataSource.setDefaultTargetDataSource(ds);
        }
        return dynamicDataSource;
    }

    @Bean(name = "sqlSessionTemplate")
    public SqlSessionTemplate sqlSessionTemplate(@Qualifier("sqlSessionFactory") SqlSessionFactory sqlSessionFactory) {
        return new SqlSessionTemplate(sqlSessionFactory);
    }

    @Bean
    public SqlSessionFactory sqlSessionFactory(@Qualifier("dynamicDataSource") DataSource dynamicDataSource) throws Exception {
        SqlSessionFactoryBean bean = new SqlSessionFactoryBean();
        bean.setDataSource(dynamicDataSource);
        return bean.getObject();
    }
}

4.使用AOP实现数据源切换

import org.aspectj.lang.annotation.After;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import rfpatrolsystem.contant.Const;
import rfpatrolsystem.result.BusinessException;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpSession;

/**
 * 利用aop原理实现切换数据源
 *
 * @author admin
 */
@Aspect
@Component
public class TargetDataSourceAspect {

    /**
     * 根据session的name值设置不同的DataSource
     */
    @Before("execution(* com.yugioh.mapper.*.*(..))")
    public void changeDataSource() throws BusinessException {
        if (DbContextHolder.getDataSourceKey() == null) {
            ServletRequestAttributes servletRequestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
            HttpServletRequest request = null;
            HttpSession session = null;
            String dbCode = null;
            if (servletRequestAttributes != null) {
                request = servletRequestAttributes.getRequest();
            }
            if (request != null) {
                dbCode = (String) request.getAttribute(Const.SESSION_DB_CODE);
                session = request.getSession();
            }
            if (dbCode == null && session != null) {
                dbCode = (String) session.getAttribute(Const.SESSION_DB_CODE);
            }
            if (dbCode != null) {
                DbContextHolder.setDataSourceKey(dbCode);
            }
        }
    }

    /**
     * 方法执行完之后清楚当前数据源，让其使用默认数据源
     */
    @After("execution(* com.yugioh.mapper.*.*(..))")
    public void restoreDataSource() {
        DbContextHolder.clearDataSourceKey();
    }
}

后续使用

1.添加数据源使用：DynamicDataSource.getInstance().createDatasourece
2.在切面中自定义数据源的切换规则