SpringBoot实现动态数据源切换

前文通过docker实现了MySQL的主从同步，那么在应用层可以对主从数据库来实现读写分离，即主库用于写/读操作，从库只用于读操作。通过读写分离，可以有效利用资源提升服务器吞吐量。

以下将介绍在SpringBoot项目下通过注解和AOP来切换数据源。
注解类CurrentDataSource：

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface CurrentDataSource {
    String value() default "MASTER";
}

AOP切面类

@Component
@Aspect
public class DynamicDataSourceAspect {

    @Before("@annotation(currentDataSource)")
    public void beforeSwitchDS(JoinPoint joinPoint, CurrentDataSource currentDataSource) {
        //获取当前访问的class
        Class clazz = joinPoint.getTarget().getClass();
        //获得访问的方法名
        String methodName = joinPoint.getSignature().getName();
        //得到方法的参数类型
        Class[] argClazz = ((MethodSignature)joinPoint.getSignature()).getParameterTypes();

        String dbType = DataSourceContextHoler.DEFAULT_DATA_SOURCE_TYPE;
        try {
            Method method = clazz.getMethod(methodName, argClazz);
            CurrentDataSource annotation = method.getAnnotation(CurrentDataSource.class);
            if (annotation != null) {
                dbType = annotation.value();
            }
        } catch (NoSuchMethodException e) {
            e.printStackTrace();
        }
        DataSourceContextHoler.setDB(dbType);
    }

    @After("@annotation(currentDataSource)")
    public void afterSwitchDS(JoinPoint joinPoint, CurrentDataSource currentDataSource) {
        DataSourceContextHoler.cleanDB();
    }
}

配置两个数据源注入到DataSource中

@Slf4j
@Configuration
public class DataSourceConfig {

    @Bean
    @ConfigurationProperties(prefix = "spring.datasource.master")
    public DataSource masterDataSource() {
        return DataSourceBuilder.create().build();
    }

    @Bean
    @ConfigurationProperties(prefix = "spring.datasource.slave")
    public DataSource slaveDataSource() {
        return DataSourceBuilder.create().build();
    }

    @Bean
    public SqlSessionFactoryBean sqlSessionFactoryBean() throws IOException {
        SqlSessionFactoryBean sqlSessionFactoryBean = new SqlSessionFactoryBean();
        sqlSessionFactoryBean.setTypeAliasesPackage("dev.springbootmasterslavedemo.entity");
        sqlSessionFactoryBean.setMapperLocations(new PathMatchingResourcePatternResolver().getResources("classpath:mybatis/mapper/*Mapper.xml"));
        sqlSessionFactoryBean.setDataSource(dataSource());
        return sqlSessionFactoryBean;
    }

    @Bean
    public SqlSessionTemplate sqlSessionTemplate () throws Exception {
        return new SqlSessionTemplate(sqlSessionFactoryBean().getObject());
    }

    @Bean
    public DataSource dataSource() {
        DynamicDataSource dynamicDataSource = new DynamicDataSource();

        //设置默认DataSource
        dynamicDataSource.setDefaultTargetDataSource(masterDataSource());

        Map<Object, Object> dsMap = new HashMap<>();
        dsMap.put("MASTER", masterDataSource());
        dsMap.put("SLAVE", slaveDataSource());

        dynamicDataSource.setTargetDataSources(dsMap);

        return dynamicDataSource;
    }
}

在Spring中，有一个AbstractRoutingDataSource抽象类，通过继承这个抽象类设置dynamicDataSource.setTargetDataSources(dsMap)可以配置多个数据源。
动态DynamicDataSource类：

@Slf4j
public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        log.info("【数据源】当前数据源为: {}", DataSourceContextHoler.getDB());
        return DataSourceContextHoler.getDB();
    }
}

在bean中注入主从的数据源后，设置dynamicDataSource.setTargetDataSources(dsMap)时，继续执行afterPropertiesSet()。在AbstractRoutingDataSource.class中，有：

AbstractRoutingDataSource.class
public void afterPropertiesSet() {
        if (this.targetDataSources == null) {
            throw new IllegalArgumentException("Property 'targetDataSources' is required");
        } else {
            this.resolvedDataSources = new HashMap(this.targetDataSources.size());
            Iterator var1 = this.targetDataSources.entrySet().iterator();

            while(var1.hasNext()) {
                Entry<Object, Object> entry = (Entry)var1.next();
                Object lookupKey = this.resolveSpecifiedLookupKey(entry.getKey());
                DataSource dataSource = this.resolveSpecifiedDataSource(entry.getValue());
                this.resolvedDataSources.put(lookupKey, dataSource);
            }

            if (this.defaultTargetDataSource != null) {
                this.resolvedDefaultDataSource = this.resolveSpecifiedDataSource(this.defaultTargetDataSource);
            }

        }
    }

可以看到，在AbstractRoutingDataSource.class中，将targetDataSources中注入的Datasource数据源赋值给了resolvedDataSources。
这里切换数据源的关键代码是AbstractRoutingDataSource.class中determineTargetDataSource()方法。

AbstractRoutingDataSource.class
protected DataSource determineTargetDataSource() {
        Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
        Object lookupKey = this.determineCurrentLookupKey();
        DataSource dataSource = (DataSource)this.resolvedDataSources.get(lookupKey);
        if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
            dataSource = this.resolvedDefaultDataSource;
        }

        if (dataSource == null) {
            throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
        } else {
            return dataSource;
        }
    }

上面方法中Object lookupKey = this.determineCurrentLookupKey()自己继承的AbstractRoutingDataSource类重写了该方法，通过统一数据源管理类DataSourceContextHoler的值来对数据源进行切换。

public class DataSourceContextHoler {

    private static final ThreadLocal<String> CONTEXT_HOLDER = new ThreadLocal<>();
    public static final String DEFAULT_DATA_SOURCE_TYPE = "MASTER";

    //设置数据源名
    public static void setDB(String dbType) {
        log.info("【设置数据源】切换到数据源: {}", dbType);
        CONTEXT_HOLDER.set(dbType);
    }

    //获取数据源名
    public static String getDB() {
        return CONTEXT_HOLDER.get();
    }

    //清除数据源
    public static void cleanDB() {
        CONTEXT_HOLDER.remove();
    }
}

至此，动态数据源的实现已经完成。

测试
在service层通过注解实现动态数据源

@Service
public class DevUserServiceImpl implements DevUserService {

    private final DevUserMapper devUserMapper;

    @Autowired
    public DevUserServiceImpl(DevUserMapper devUserMapper) {
        this.devUserMapper = devUserMapper;
    }


    @CurrentDataSource(value = "SLAVE")
    @Override
    public DevUser findByUserName(String username) {
        DevUser devUser = devUserMapper.findByUserName(username);
        return devUser;
    }

    @CurrentDataSource
    @Transactional
    @Override
    public void save(DevUser devUser) {
        devUserMapper.save(devUser);
    }
}

通过web调用/select可以看到控制台日志：

调用/save可以看到控制台日志：

通过Debug可以看到一整个流程，到99行执行实现该抽象类里重写的方法获取lookupkey值，在100行根据lookupkey值获取对应的数据源，然后返回数据源，然后去进行连接数据库，执行后续逻辑，最终返回数据：