spring AbstractRoutingDataSource类、AOP实现动态数据源切换

spring boot中AbstractRoutingDataSource 类会根据用户定义的规则选择当前的数据源，这样我们可以在执行查询之前，设置使用的数据源。实现可动态路由的数据源，在每次数据库查询操作前执行。它的抽象方法 determineCurrentLookupKey() 决定使用哪个数据源。

import com.zaxxer.hikari.HikariDataSource;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class DynamicDataSource extends AbstractRoutingDataSource {
    private Map<Object, Object> hikariDataSourceMap = new HashMap<>();
	
	//根据key获取指定数据源
    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceHolder.getDataSourceKey();
    }
	
	//设置targetDataSources保存key和数据库连接的映射关系
    public DynamicDataSource() {
        this.initDataSources();
        this.setTargetDataSources(hikariDataSourceMap);
    }
    
	//初始化数据源，在这里配置两个数据源ds0,ds1
    private void initDataSources() {
        List<String> keys = new ArrayList<>();
        keys.add("ds0");
        keys.add("ds1");
        keys.forEach(this::singleDataSource);
    }
	
	//配置单个数据源并加入Map
    private void singleDataSource(String key) {
        HikariDataSource hikariDataSource = new HikariDataSource();
        hikariDataSource.setJdbcUrl(MessageFormat.format("jdbc:mysql://localhost:3306/{0}?useSSL=false&serverTimezone=UTC&characterEncoding=UTF8", key));
        hikariDataSource.setUsername("root");
        hikariDataSource.setPassword("123456");
        hikariDataSource.setDriverClassName("com.mysql.cj.jdbc.Driver");
        hikariDataSourceMap.put(key, hikariDataSource);
    }
}

public class DataSourceHolder {
    static final ThreadLocal<String> dataSourceKey = new ThreadLocal<>();
	
	//默认连接ds0
    public static String getDataSourceKey() {
        String key = dataSourceKey.get();
        return key == null ? "ds0" : key;
    }

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

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

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class DataSourceConfiguration {

    @Bean
    public DynamicDataSource dynamicDataSource() {
        return new DynamicDataSource();
    }
}

import org.aspectj.lang.annotation.AfterReturning;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

@Aspect
@Component
@Order(-1) //保证此AOP在@Transactional之前执行
public class DataSourceAspect {
	//也可以选择注解层，这里是所有Controller层
    @Pointcut("execution(public * 项目路径.controller..*.*(..))")
    public void handlePlaceholder() {
    }
	
	//根据自定义条件指定key
    @Before("handlePlaceholder()")
    public void doBefore() {
    	//自定义，这里简单指定为ds1
        DataSourceHolder.setDataSourceKey("ds1");
    }

    @AfterReturning(pointcut = "handlePlaceholder()")
    public void doAfterReturning() {
        DataSourceHolder.clear();
    }
}