Spring手动注册/销毁多数据源

package org.demo.spring.mysql;

import com.alibaba.druid.pool.DruidDataSource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.demo.spring.mysql.dto.DataSourceDTO;
import org.demo.spring.mysql.dto.DataSourceService;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.List;
import java.util.concurrent.locks.ReentrantReadWriteLock;

@Component
@Slf4j
public class SimpleJdbcTemplate implements InitializingBean {

    @Resource
    private ApplicationContext applicationContext;

    @Resource
    private DataSourceService dataSourceService;

    /**
     * 可重入读写锁
     */
    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();

    @Override
    public void afterPropertiesSet() {
        registryAllDataSourceBean();
        log.info("初始化数据源配置数据成功");
    }


    /**
     * 获取数据源对应的JdbcTemplateBean
     *
     * @param dataSourceName 数据源配置名
     * @return JdbcTemplate
     */
    public JdbcTemplate getJdbcTemplate(String dataSourceName) {
        lock.readLock().lock();
        try {
            return this.applicationContext.getBean(buildJdbcTemplateBeanName(dataSourceName), JdbcTemplate.class);
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * 获取SpringBean注册器
     *
     * @return BeanDefinitionRegistry
     */
    private DefaultListableBeanFactory getBeanFactory() {
        ConfigurableApplicationContext configurableApplicationContext = (ConfigurableApplicationContext) applicationContext;
        return (DefaultListableBeanFactory) configurableApplicationContext.getBeanFactory();
    }

    /**
     * 启动的时候加载所有已经定义的DataSourceBean、JdbcTemplateBean
     */
    private void registryAllDataSourceBean() {
        lock.writeLock().lock();
        try {
            List<DataSourceDTO> tagDataSourceList = this.dataSourceService.findAllDataSource();
            DefaultListableBeanFactory beanFactory = getBeanFactory();
            if (CollectionUtils.isNotEmpty(tagDataSourceList)) {
                tagDataSourceList.forEach(tagDataSource -> {

                    //注册DataSourceBean
                    GenericBeanDefinition dataSourceBeanDefinition = this.buildDataSourceBeanDefinition(tagDataSource);
                    beanFactory.registerBeanDefinition(buildDataSourceBeanName(tagDataSource.getDataSourceName()),
                            dataSourceBeanDefinition);

                    //注册JdbcTemplateBean
                    GenericBeanDefinition jdbcTemplateBeanDefinition = this.buildJdbcTemplateBeanDefinition(tagDataSource);
                    beanFactory.registerBeanDefinition(buildJdbcTemplateBeanName(tagDataSource.getDataSourceName()),
                            jdbcTemplateBeanDefinition);
                });
            }
        } finally {
            lock.writeLock().unlock();
        }
    }

    /**
     * 注册新的数据源Bean
     * 先对老的数据源Bean进行注销，再创建新的数据源
     *
     * @param dataSourceName 数据源配置名
     */
    private void registryDataSourceBean(String dataSourceName) {
        log.info("开始注册新的数据源[{}]。", dataSourceName);
        lock.writeLock().lock();
        try {
            DataSourceDTO tagDataSource = this.dataSourceService.findByDataSourceName(dataSourceName);
            if (tagDataSource != null) {
                DefaultListableBeanFactory beanFactory = getBeanFactory();

                //注册DataSourceBean
                GenericBeanDefinition dataSourceBeanDefinition = this.buildDataSourceBeanDefinition(tagDataSource);
                beanFactory.registerBeanDefinition(buildDataSourceBeanName(dataSourceName), dataSourceBeanDefinition);

                //注册JdbcTemplateBean
                GenericBeanDefinition jdbcTemplateBeanDefinition = this.buildJdbcTemplateBeanDefinition(tagDataSource);
                beanFactory.registerBeanDefinition(buildJdbcTemplateBeanName(dataSourceName), jdbcTemplateBeanDefinition);
                log.info("注册新的数据源[{}]成功！", dataSourceName);
            }
        } finally {
            lock.writeLock().unlock();
        }
    }

    /**
     * 注销数据源Bean
     * 先注销对于的JdbcTemplateBean,再注销DataSourceBean
     * 注销可由事件触发
     * @param dataSourceName 数据源配置名
     */
    private void unRegistryDataSourceBean(String dataSourceName) {
        log.info("开始注销数据源[{}]。", dataSourceName);
        lock.writeLock().lock();
        try {
            DefaultListableBeanFactory beanFactory = getBeanFactory();
            String jdbcTemplateBeanName = this.buildJdbcTemplateBeanName(dataSourceName);
            String dataSourceBeanName = this.buildDataSourceBeanName(dataSourceName);
            if (beanFactory.containsBeanDefinition(jdbcTemplateBeanName)) {
                beanFactory.destroySingleton(jdbcTemplateBeanName);
                beanFactory.removeBeanDefinition(jdbcTemplateBeanName);
                if (beanFactory.containsBeanDefinition(dataSourceBeanName)) {
                    beanFactory.destroySingleton(dataSourceBeanName);
                    beanFactory.removeBeanDefinition(dataSourceBeanName);
                }
                log.info("注销数据源[{}]成功！", dataSourceName);
            } else {
                log.info("不存在数据源[{}]，不需要注销！", dataSourceName);
            }
        } finally {
            lock.writeLock().unlock();
        }
    }

    /**
     * 构建DataSource的BeanDefinition
     *
     * @param tagDataSource 数据源配置对象
     * @return BeanDefinition
     */
    private GenericBeanDefinition buildDataSourceBeanDefinition(DataSourceDTO tagDataSource) {
        BeanDefinitionBuilder builder = BeanDefinitionBuilder.genericBeanDefinition(DruidDataSource.class);
        GenericBeanDefinition definition = (GenericBeanDefinition) builder.getRawBeanDefinition();
        definition.setInitMethodName("init");
        definition.setDestroyMethodName("close");
        definition.setBeanClass(DruidDataSource.class);
        definition.setAutowireMode(GenericBeanDefinition.AUTOWIRE_BY_NAME);
        definition.getPropertyValues()
                .add("driverClassName", "com.mysql.jdbc.Driver")
                .add("url", "jdbc:mysql://" + tagDataSource.getDbIp() + ":" + tagDataSource.getDbPort() +
                        "?useUnicode=true&characterEncoding=UTF-8&allowMultiQueries=true&zeroDateTimeBehavior=convertToNull")
                .add("username", tagDataSource.getDbUserName())
                .add("password", tagDataSource.getDbUserPwd())
                .add("initialSize", 20)
                .add("minIdle", 10)
                .add("maxActive", 100)
                .add("maxWait", 30)
                .add("timeBetweenEvictionRunsMillis", 60000)
                .add("minEvictableIdleTimeMillis", 300000)
                .add("validationQuery", "select 1 from dual")
                .add("testWhileIdle", "true")
                .add("testOnBorrow", "false")
                .add("testOnReturn", "false")
                .add("poolPreparedStatements", "false");
        return definition;
    }

    /**
     * 构建JdbcTemplate的BeanDefinition
     *
     * @param tagDataSource 数据源配置对象
     * @return BeanDefinition
     */
    private GenericBeanDefinition buildJdbcTemplateBeanDefinition(DataSourceDTO tagDataSource) {
        BeanDefinitionBuilder builder = BeanDefinitionBuilder.genericBeanDefinition(JdbcTemplate.class);
        GenericBeanDefinition definition = (GenericBeanDefinition) builder.getRawBeanDefinition();
        definition.setBeanClass(JdbcTemplate.class);
        definition.setAutowireMode(GenericBeanDefinition.AUTOWIRE_BY_NAME);
        definition.getPropertyValues().add("dataSource", this.applicationContext
                .getBean(buildDataSourceBeanName(tagDataSource.getDataSourceName()), DruidDataSource.class));
        return definition;
    }

    /**
     * 构建DataSourceBean的名称
     *
     * @param dataSourceName 数据源配置名
     * @return BeanName
     */
    private String buildDataSourceBeanName(String dataSourceName) {
        return dataSourceName + "DataSource";
    }

    /**
     * 构建JdbcTemplateBean的名称
     *
     * @param dataSourceName 数据源配置名
     * @return BeanName
     */
    private String buildJdbcTemplateBeanName(String dataSourceName) {
        return dataSourceName + "JdbcTemplate";
    }


}

/**
 * Description: 查询数据库配置
 */
@Component
public class DataSourceService {

    public List<DataSourceDTO> findAllDataSource() {
        List<DataSourceDTO> dataSource = new ArrayList<>();
        DataSourceDTO dataSourceDTO = new DataSourceDTO();
        dataSourceDTO.setDataSourceName("app");
        dataSourceDTO.setDbUserName("root");
        dataSourceDTO.setDbIp("127.0.0.1");
        dataSourceDTO.setDbUserPwd("root");
        dataSourceDTO.setDbPort("3306");
        dataSource.add(dataSourceDTO);
        return dataSource;
    }

    public DataSourceDTO findByDataSourceName(String dataSourceName) {
        return null;
    }
}