JPA动态注册多数据源

背景

目前已经是微服务的天下，但是随着业务需求的日益增长，部分应用还是出现了需要同时连接多个数据源操作数据的技术诉求。

需要对现有的技术架构进行优化升级，查阅了下网上的文章，基本都是照搬的同一篇文章，通过代码的方式同时注册primary和second两个数据源。这种实现方案的技术成本比较低，但是维护成本非常高的，如果我需要同时连接4个、5个甚至更多的数据源，需要不断增加代码注册数据源。

实现方案

比较理想的方式，当然是增加些许配置，就能实现数据源的拓展。本文只讨论如果动态注册JPA的数据源，涉及数据源切换等功能大家可以自行实现。

JPA多数据源需要在Spring IOC中注册DataSource、EntityManagerFactory和JpaTransactionManger，具体实现通过BeanDefinitionRegistryPostProcessor进行bean的动态注册，源码如下：

配置项：

spring.datasource.multiple.test.username=root
spring.datasource.multiple.test.password=123456
spring.datasource.multiple.test.url=jdbc:mysql://localhost:3306/tester?useUnicode=true&characterEncoding=UTF-8&useSSL=false&allowMultiQueries=true&serverTimezone=GMT%2B8
spring.datasource.multiple.test.driver-class-name=com.mysql.cj.jdbc.Driver
spring.datasource.multiple.test.dialect=org.hibernate.dialect.MySQLDialect
spring.datasource.multiple.test.hikari.pool-name=test-datasource
spring.datasource.multiple.test.hikari.maximum_pool_size=12
#连接超时时间:毫秒，小于250毫秒，否则被重置为默认值30秒
spring.datasource.multiple.test.hikari.connection_timeout=6000
#最小空闲连接，默认值10，小于0或大于maximum-pool-size，都会重置为maximum-pool-size
spring.datasource.multiple.test.hikari.minimum_idle=10
#空闲连接超时时间，默认值600000（10分钟），大于等于max-lifetime且max-lifetime>0，会被重置为0；不等于0且小于10秒，会被重置为10秒。
# 只有空闲连接数大于最大连接数且空闲时间超过该值，才会被释放
spring.datasource.multiple.test.hikari.idle_timeout=600000
#连接最大存活时间.不等于0且小于30秒，会被重置为默认值30分钟.设置应该比mysql设置的超时时间短
spring.datasource.multiple.test.hikari.max_lifetime=1800000
spring.datasource.multiple.test.hikari.login_timeout=500
spring.datasource.multiple.test.hikari.validation_timeout=1000
spring.datasource.multiple.test.hikari.initialization_fail_timeout=1000
#连接测试查询
spring.datasource.multiple.test.hikari.connection_test_query=SELECT 1

配置类：

public class MultipleDataSource {

    private String url;

    private String username;

    private String password;

    private String driverClassName;

    private String dialect;

    private HikariConfig hikari;

    public String getUrl() {
        return url;
    }

    public void setUrl(String url) {
        this.url = url;
    }

    public String getUsername() {
        return username;
    }

    public void setUsername(String username) {
        this.username = username;
    }

    public String getPassword() {
        return password;
    }

    public void setPassword(String password) {
        this.password = password;
    }

    public String getDriverClassName() {
        return driverClassName;
    }

    public void setDriverClassName(String driverClassName) {
        this.driverClassName = driverClassName;
    }

    public String getDialect() {
        return dialect;
    }

    public void setDialect(String dialect) {
        this.dialect = dialect;
    }

    public HikariConfig getHikari() {
        return hikari;
    }

    public void setHikari(HikariConfig hikari) {
        this.hikari = hikari;
    }

    /**
     * 数据源配置参数校验
     */
    public void validate() {
        Assert.hasText(url, "数据库连接地址不能为空");
        Assert.hasText(username, "数据库用户名不能为空");
        Assert.hasText(password, "数据库密码不能为空");
        Assert.hasText(driverClassName, "数据库驱动类不能为空");
        Assert.hasText(dialect, "数据库方言不能为空");
        Assert.notNull(hikari, "数据库连接池配置不能为空");
    }

}

public class MultipleDataSourceProperties {

    private Map<String, MultipleDataSource> multiple;

    public Map<String, MultipleDataSource> getMultiple() {
        return multiple;
    }

    public void setMultiple(Map<String, MultipleDataSource> multiple) {
        this.multiple = multiple;
    }

    /**
     * 数据源配置参数校验
     */
    public void validate() {
        for (Map.Entry<String, MultipleDataSource> entry : multiple.entrySet()) {
            entry.getValue().validate();
        }
    }

}

动态注册：

@Slf4j
@Configuration
public class MultipleDataSourceRegistryPostProcessor implements BeanDefinitionRegistryPostProcessor, EnvironmentAware, ApplicationContextAware {

    private static final String DATASOURCE_PREFIX = "spring.datasource";

    public static final String BEAN_DATASOURCE = "%sDataSource";

    public static final String BEAN_ENTITY_MANAGER_FACTORY = "%sEntityManagerFactory";

    public static final String BEAN_TRANSACTION_MANAGER = "%sTransactionManager";

    protected static final String SCAN_PACKAGE = "com.cestc.*.bean.%s";

    protected static final String PERSISTENCE_NAME = "%sPersistenceUnit";

    protected final Map<String, Object> jpaProperties = new HashMap<>(8);

    private Environment environment;

    private ApplicationContext context;

    public MultipleDataSourceRegistryPostProcessor() {
        jpaProperties.put("current_session_context_class", "org.springframework.orm.hibernate5.SpringSessionContext");
        jpaProperties.put("hibernate.show_sql", "true");
        jpaProperties.put("hibernate.format_sql", "true");
        jpaProperties.put("hibernate.hbm2ddl.auto", "none");
    }

    @Override
    public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry) throws BeansException {
    }

    /**
     * 注册Hikari数据源
     * @param dataSource    数据源配置信息
     */
    protected BeanDefinition registryDataSource(MultipleDataSource dataSource) {

        return BeanDefinitionBuilder.genericBeanDefinition(HikariDataSource.class, () -> {
            dataSource.getHikari().setJdbcUrl(dataSource.getUrl());
            dataSource.getHikari().setUsername(dataSource.getUsername());
            dataSource.getHikari().setPassword(dataSource.getPassword());
            dataSource.getHikari().setDriverClassName(dataSource.getDriverClassName());

            return new HikariDataSource(dataSource.getHikari());
        }).getBeanDefinition();

    }

    /**
     * 注册EntityManagerFactory
     * @param name          数据源名称
     * @param dataSource    数据源配置信息
     * @param beanFactory
     */
    protected BeanDefinition registryEntityManagerFactory(String name, MultipleDataSource dataSource, DefaultListableBeanFactory beanFactory) {

        HikariDataSource hikariDataSource = beanFactory.getBean(String.format(BEAN_DATASOURCE, name), HikariDataSource.class);
        jpaProperties.put("hibernate.dialect", dataSource.getDialect());

        return BeanDefinitionBuilder.genericBeanDefinition(LocalContainerEntityManagerFactoryBean.class, () -> {
            LocalContainerEntityManagerFactoryBean entityManagerFactoryBean = new LocalContainerEntityManagerFactoryBean();
            entityManagerFactoryBean.setDataSource(hikariDataSource);
            entityManagerFactoryBean.setPackagesToScan(String.format(SCAN_PACKAGE, name));
            entityManagerFactoryBean.setJpaVendorAdapter(new HibernateJpaVendorAdapter());
            entityManagerFactoryBean.setJpaPropertyMap(jpaProperties);
            entityManagerFactoryBean.setPersistenceUnitName(String.format(PERSISTENCE_NAME, name));

            return entityManagerFactoryBean;
        }).getBeanDefinition();

    }

    /**
     * 注册数据源事务管理器
     * @param name
     * @param beanFactory
     */
    protected BeanDefinition registryTransactionManager(String name, DefaultListableBeanFactory beanFactory) {

        BeanDefinitionBuilder builder = BeanDefinitionBuilder.rootBeanDefinition(JpaTransactionManager.class);
        builder.addConstructorArgReference(String.format(BEAN_ENTITY_MANAGER_FACTORY, name));
        return builder.getBeanDefinition();

    }

    @Override
    public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory) throws BeansException {

        DefaultListableBeanFactory defaultListableBeanFactory = (DefaultListableBeanFactory) beanFactory;

        MultipleDataSourceProperties multipleDataSourceProperties = Binder.get(this.environment).bind(DATASOURCE_PREFIX, MultipleDataSourceProperties.class).get();

        if (Optional.ofNullable(multipleDataSourceProperties).isPresent()
                && !multipleDataSourceProperties.getMultiple().isEmpty()) {

            // 校验数据源参数
            multipleDataSourceProperties.validate();

            for (Map.Entry<String, MultipleDataSource> entry : multipleDataSourceProperties.getMultiple().entrySet()) {
                log.info("注册数据源[{}]", entry.getKey());

                // 注册数据源
                defaultListableBeanFactory.registerBeanDefinition(String.format(BEAN_DATASOURCE, entry.getKey()), registryDataSource(entry.getValue()));
                // 注册entityManagerFactory
                defaultListableBeanFactory.registerBeanDefinition(String.format(BEAN_ENTITY_MANAGER_FACTORY, entry.getKey()), registryEntityManagerFactory(entry.getKey(), entry.getValue(), defaultListableBeanFactory));
                // 注册事务管理器
                defaultListableBeanFactory.registerBeanDefinition(String.format(BEAN_TRANSACTION_MANAGER, entry.getKey()), registryTransactionManager(entry.getKey(), defaultListableBeanFactory));
            }
        }

    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.context = applicationContext;
    }

    @Override
    public void setEnvironment(Environment environment) {
        this.environment = environment;
    }
}

欢迎交流