spring boot 多数据源动态切换

1. 配置文件

spring:
  datasource:
    driver-class-name: com.mysql.jdbc.Driver
    password: root
    username: root
    url: jdbc:mysql://localhost:3306/demo?useUnicode=true&characterEncoding=utf8&useSSL=true
  jpa:
    hibernate:
      ddl-auto: none
    show-sql: true
    properties:
      hibernate:
       format_sql: true
    database-platform: org.hibernate.dialect.MySQL57Dialect
customer:
  datasource:
    ds1:
      driver-class-name: com.mysql.jdbc.Driver
      password: root
      username: root
      url: jdbc:mysql://localhost:3306/demo1?useUnicode=true&characterEncoding=utf8&useSSL=true
    ds2:
      driver-class-name: com.mysql.jdbc.Driver
      password: root
      username: root
      url: jdbc:mysql://localhost:3306/demo2?useUnicode=true&characterEncoding=utf8&useSSL=true
logging:
  level:
    com.meng.dynamic1: debug

spring.datasource为默认数据库配置，customer.datasource为其他数据库配置，其中ds1，ds2为数据库标识名称，可根据自己要求更改

package com.meng.dynamic1.properties;

import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;

import java.util.HashMap;
import java.util.Map;

public class CustomerDataSourceProperties {

    private Map<String, DataSourceProperties> datasource = new HashMap<>();
    
    // 省略了set get方法
}

创建CustomerDataSourceProperties用于方便接受数据库配置

2. 通过ThreadLocal存放当前使用哪个数据源

package com.meng.dynamic1.datasource;

import java.util.ArrayList;
import java.util.List;

public class DynamicDataSourceContextHolder {

    private static final ThreadLocal<String> CONTEXT_HOLDER;

    private static final List<String> DATASOURCES;

    static {
        CONTEXT_HOLDER = new ThreadLocal<>();
        DATASOURCES = new ArrayList<>();
    }

    public static void set(String datasource) {
        CONTEXT_HOLDER.set(datasource);
    }

    public static String get() {
        return CONTEXT_HOLDER.get();
    }

    public static void remove() {
        CONTEXT_HOLDER.remove();
    }

    public static boolean contains(String datasource) {
        return DATASOURCES.contains(datasource);
    }

    public static void add(String name) {
        DATASOURCES.add(name);
    }

}

CONTEXT_HOLDER存放当前指定数据源，当前无存放时，使用默认数据源
DATASOURCES管理着所有数据源，包括默认数据源

3. 通过注解方式动态切换数据源

(1). 注解

package com.meng.dynamic1.annotation;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface TargetDataSource {

    String value();

}

(2). AOP

package com.meng.dynamic1.aop;

import com.meng.dynamic1.annotation.TargetDataSource;
import com.meng.dynamic1.datasource.DynamicDataSourceContextHolder;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.After;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.stereotype.Component;

@Aspect
@Component
@Slf4j
public class DynamicDataSourceAspect {

    @Before("@annotation(targetDataSource)")
    public void before(JoinPoint joinPoint, TargetDataSource targetDataSource) {
        String datasource = targetDataSource.value();
        if (!DynamicDataSourceContextHolder.contains(datasource)) {
            log.error("数据源[{}]不存在，使用默认数据源 > {}", datasource, joinPoint.getSignature());
        } else {
            log.debug("Use datasource: {} > {}", datasource, joinPoint.getSignature());
        }
        DynamicDataSourceContextHolder.set(datasource);
    }

    @After("@annotation(targetDataSource)")
    public void after(JoinPoint joinPoint, TargetDataSource targetDataSource) {
        log.debug("Revert datasource: {} > {}", targetDataSource.value(), joinPoint.getSignature());
        DynamicDataSourceContextHolder.remove();
    }

}

4. AbstractRoutingDataSource

package com.meng.dynamic1.datasource;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.get();
    }

}

AbstractRoutingDataSource继承了数据javax.sql.DataSource，本质上他也是一个数据源
继承AbstractRoutingDataSource，实现determineCurrentLookupKey方法

5. 初始化并注册DynamicDataSource

package com.meng.dynamic1.datasource;

import com.meng.dynamic1.properties.CustomerDataSourceProperties;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.MutablePropertyValues;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
import org.springframework.boot.context.properties.bind.Bindable;
import org.springframework.boot.context.properties.bind.Binder;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.env.Environment;
import org.springframework.core.type.AnnotationMetadata;

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

@Slf4j
public class DynamicDataSourceRegister implements ImportBeanDefinitionRegistrar, EnvironmentAware {

    private DataSource defaultDataSource;

    private Map<String, DataSource> customerDataSources = new HashMap<>();

    @Override
    public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry) {
        Map<String, Object> targetDataSources = new HashMap<>();
        targetDataSources.put("defaultDataSource", defaultDataSource);
        targetDataSources.putAll(customerDataSources);
        targetDataSources.forEach((k, v) -> DynamicDataSourceContextHolder.add(k));
        GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
        beanDefinition.setBeanClass(DynamicDataSource.class);
        beanDefinition.setSynthetic(true);
        MutablePropertyValues propertyValues = beanDefinition.getPropertyValues();
        propertyValues.addPropertyValue("defaultTargetDataSource", defaultDataSource);
        propertyValues.addPropertyValue("targetDataSources", targetDataSources);
        registry.registerBeanDefinition("dynamicDataSource", beanDefinition);
    }

    @Override
    public void setEnvironment(Environment environment) {
        Binder binder = Binder.get(environment);
        DataSourceProperties defaultDataSourceProperties = binder.bind("spring.datasource", Bindable.of(DataSourceProperties.class)).get();
        defaultDataSource = defaultDataSourceProperties.initializeDataSourceBuilder().build();
        CustomerDataSourceProperties customerDataSourceProperties = binder.bind("customer", Bindable.of(CustomerDataSourceProperties.class)).get();
        customerDataSourceProperties.getDatasource().forEach((k, v) -> customerDataSources.put(k, v.initializeDataSourceBuilder().build()));
    }

}

最后需要在main方法上，使用@Import注解引入DynamicDataSourceRegister类才会注册到Spring容器中