基于AbstractRoutingDataSource实现数据源的动态添加与动态切换

基于AbstractRoutingDataSource实现数据源的动态添加与动态切换

实现目的：在项目运行后动态的管理和切换数据源 包括设置默认数据源、初始化数据源、数据源的添加、删除、切换等

1. 继承 AbstractRoutingDataSource 并重写其 determineCurrentLookupKey() 方法来确定当前应该使用的数据源的标识（通常是数据源的名称或标识符），然后在数据访问时根据这个标识来选择实际的数据源。该类创建了两个添加数据源的方法，可在程序运行时动态的添加数据源

/**
 * 继承 AbstractRoutingDataSource 并重写其 determineCurrentLookupKey()方法
 * 来确定当前应该使用的数据源的标识（通常是数据源的名称或标识符），
 * 然后在数据访问时根据这个标识来选择实际的数据源。
 */
public class DynamicRoutingDataSource extends AbstractRoutingDataSource {

    /**
     * 存储存入到DynamicRoutingDataSource的数据源
     */
    private Map<Object, Object> targetDataSources = new HashMap<>();

    public DynamicRoutingDataSource() {
        super.setTargetDataSources(targetDataSources);
    }

    @Resource
    private DynamicRoutingDataSource dynamicRoutingDataSource;

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

    /**
     * 添加多个数据源
     * @param map key:数据源的key value:数据源
     */
    public void addDataSources(Map<Object, Object> map){
        targetDataSources.putAll(map);
        dynamicRoutingDataSource.setTargetDataSources(targetDataSources);
        dynamicRoutingDataSource.afterPropertiesSet();
    }

    /**
     * 添加数据源
     */
    public void addDataSource(String datasourceKey, DataSource dataSource){
        targetDataSources.put(datasourceKey, dataSource);
        dynamicRoutingDataSource.setTargetDataSources(targetDataSources);
        dynamicRoutingDataSource.afterPropertiesSet();
    }

		/**
     * 删除数据源
     */
    public void removeDataSource(String datasourceKey){
        targetDataSources.remove(datasourceKey);
        dynamicRoutingDataSource.setTargetDataSources(targetDataSources);
        dynamicRoutingDataSource.afterPropertiesSet();
    }
}

2. 切换数据源类DataSourceContextHolder，该类通过ThreadLocal存储当前线程的数据源对应的key，切换数据源是调用setDataSourceKey（key）即可完成数据源的动态切换

/**
 * 该类通过ThreadLocal存储当前线程的数据源对应的key
 */
public class DataSourceContextHolder {

    private static final ThreadLocal<String> dataSourceKey = new ThreadLocal<>();

    public static void setDataSourceKey(String key) {
        dataSourceKey.set(key);
    }

    public static String getDataSourceKey() {
        return dataSourceKey.get();
    }

    public static void clearDataSourceKey() {
        dataSourceKey.remove();
    }
}

3. 设置DynamicDataSourceConfiguration配置类，初始换DynamicRoutingDataSource并将其交给spring IOC容器管理

@Configuration
public class DynamicDataSourceConfiguration {

    @Value("${spring.datasource.username}")
    private String username;
    @Value("${spring.datasource.password}")
    private String password;
    @Value("${spring.datasource.url}")
    private String url;
    @Value("${spring.datasource.driver-class-name}")
    private String driverClassName;

    @Bean
    public DynamicRoutingDataSource dynamicRoutingDataSource() {
        DynamicRoutingDataSource dynamicRoutingDataSource = new DynamicRoutingDataSource();
        DataSource dataSource = DataSourceBuilder.create()
                .url(url)
                .driverClassName(driverClassName)
                .username(username)
                .password(password)
                .build();
        // 设置默认数据源
        dynamicRoutingDataSource.setDefaultTargetDataSource(dataSource);
        return dynamicRoutingDataSource;
    }
}

上述基于AbstractRoutingDataSource的动态数据源管理的基础框架搭建完毕，具体使用示例如下：

1. 初始化数据源

   实现ApplicationListener在项目启动后自动执行，通过该类自动注入的dynamicRoutingDataSource自动加载初始数据源

@Component
public  class  DataSourceInitListener   implements ApplicationListener<ContextRefreshedEvent> {

    @Resource
    private DatabaseNodeMapper databaseNodeMapper;
    @Resource
    private DynamicRoutingDataSource dynamicRoutingDataSource;
    protected  static  final Logger LOGGER = LoggerFactory.getLogger(DataSourceInitListener. class );

    @SuppressWarnings ( "unchecked" )
    @Override
    public  void  onApplicationEvent(ContextRefreshedEvent ev) {
        //防止重复执行。
        if (ev.getApplicationContext().getParent() ==  null ){
            List<DatabaseNode> databaseNodes = databaseNodeMapper.selectList(null);
            Map<Object, Object> map = new HashMap<>();
            for (DatabaseNode databaseNode : databaseNodes) {
                DataSource dataSource = DataSourceBuilder.create()
                        .url(databaseNode.getUrl())
                        .driverClassName(databaseNode.getDriverClassName())
                        .username(databaseNode.getUsername())
                        .password(databaseNode.getPassword())
                        .build();
                // key由一个统一的自定义规则定义，方便获取
                map.put("dataSource" + databaseNode.getId(), dataSource);
            }
            dynamicRoutingDataSource.addDataSources(map);
        }
    }
}

2. 使用示例

@SpringBootTest(classes = DatabaseConfigurationTool.class)
@RunWith(SpringRunner.class)
public class DynamicTest {
    
    @Resource
    private TestService testService;

    @Test
    public void test1(){
        // 切换数据源
        DataSourceContextHolder.setDataSourceKey("dataSource7");
        List<Map<String, String>> s = testService.test01();
        // 清空数据源，防止该线程结束后未及时清理回到线程池，下一次该线程执行初始数据源错误
        DataSourceContextHolder.clearDataSourceKey();
    }
}