SpringBoot+AOP构建多数据源的切换实践

针对微服务架构中常用的设计模块，通常我们都会需要使用到druid作为我们的数据连接池，当架构发生扩展的时候 ，通常面对的数据存储服务器也会渐渐增加，从原本的单库架构逐渐扩展为复杂的多库架构。

当在业务层需要涉及到查询多种同数据库的场景下，我们通常需要在执行sql的时候动态指定对应的datasource。

而Spring的AbstractRoutingDataSource则正好为我们提供了这一功能点，下边我将通过一个简单的基于springboot+aop的案例来实现如何通过自定义注解切换不同的数据源进行读数据操作，同时也将结合部分源码的内容进行讲解。

首先我们需要自定义一个专门用于申明当前java应用程序所需要使用到哪些数据源信息：

package mutidatasource.annotation;

import mutidatasource.config.DataSourceConfigRegister;
import mutidatasource.enums.SupportDatasourceEnum;
import org.springframework.context.annotation.Import;
import org.springframework.stereotype.Component;

import java.lang.annotation.*;

/**
 * 注入数据源
 *
 * @author idea
 * @data 2020/3/7
 */
@Target({ElementType.METHOD,ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Import(DataSourceConfigRegister.class)
public @interface AppDataSource {

    SupportDatasourceEnum[] datasourceType();
}

这里为了方便，我将测试中使用的数据源地址都配置在来enum里面，如果后边需要灵活处理的话，可以将这些配置信息抽取出来放在一些配置中心上边。

package mutidatasource.enums;

import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.NoArgsConstructor;

/**
 * 目前支持的数据源信息
 *
 * @author idea
 * @data 2020/3/7
 */
@AllArgsConstructor
@Getter
public enum SupportDatasourceEnum {

    PROD_DB("jdbc:mysql://localhost:3306/db-prod?useUnicode=true&characterEncoding=utf8","root","root","db-prod"),

    DEV_DB("jdbc:mysql://localhost:3306/db-dev?useUnicode=true&characterEncoding=utf8","root","root","db-dev"),

    PRE_DB("jdbc:mysql://localhost:3306/db-pre?useUnicode=true&characterEncoding=utf8","root","root","db-pre");

    String url;
    String username;
    String password;
    String databaseName;

    @Override
    public String toString() {
        return super.toString().toLowerCase();
    }
}

之所以要创建这个@AppDataSource注解，是要在springboot的启动类上边进行标注：

package mutidatasource;

import mutidatasource.annotation.AppDataSource;
import mutidatasource.enums.SupportDatasourceEnum;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;

/**
 * @author idea
 * @data 2020/3/7
 */
@SpringBootApplication
@AppDataSource(datasourceType = {SupportDatasourceEnum.DEV_DB, SupportDatasourceEnum.PRE_DB, SupportDatasourceEnum.PROD_DB})
public class SpringApplicationDemo {

    public static void main(String[] args) {
        SpringApplication.run(SpringApplicationDemo.class);
    }

}

借助springboot的ImportSelector 自定义一个注册器来获取启动类头部的注解所指定的数据源类型：

package mutidatasource.config;

import lombok.extern.slf4j.Slf4j;
import mutidatasource.annotation.AppDataSource;
import mutidatasource.core.DataSourceContextHolder;
import mutidatasource.enums.SupportDatasourceEnum;
import org.springframework.context.annotation.ImportSelector;
import org.springframework.core.annotation.AnnotationAttributes;
import org.springframework.core.type.AnnotationMetadata;
import org.springframework.stereotype.Component;

/**
 * @author idea
 * @data 2020/3/7
 */
@Slf4j
@Component
public class DataSourceConfigRegister implements ImportSelector {

    @Override
    public String[] selectImports(AnnotationMetadata annotationMetadata) {
        AnnotationAttributes attributes = AnnotationAttributes.fromMap(annotationMetadata.getAnnotationAttributes(AppDataSource.class.getName()));
        System.out.println("#######  datasource import #######");
        if (null != attributes) {
            Object object = attributes.get("datasourceType");
            SupportDatasourceEnum[] supportDatasourceEnums = (SupportDatasourceEnum[]) object;
            for (SupportDatasourceEnum supportDatasourceEnum : supportDatasourceEnums) {
                DataSourceContextHolder.addDatasource(supportDatasourceEnum);
            }
        }
        return new String[0];
    }


}

好的，现在我们已经能够获取到对应的数据源类型信息了，这里你会看到一个叫做DataSourceContextHolder的角色。这个对象主要是用于对每个请求线程的数据源信息做统一的分配和管理。

在多并发场景下，为了防止不同线程请求的数据源出现“互窜”情况，通常我们都会使用到threadlocal来做处理。为每一个线程都分配一个指定的，属于其内部的副本变量，当当前线程结束之前，记得将对应的线程副本也进行销毁。

package mutidatasource.core;

import mutidatasource.enums.SupportDatasourceEnum;

import java.util.HashSet;

/**
 * @author idea
 * @data 2020/3/7
 */
public class DataSourceContextHolder {

    private static final HashSet dataSourceSet = new HashSet<>();

    private static final ThreadLocal databaseHolder = new ThreadLocal<>();

    public static void setDatabaseHolder(SupportDatasourceEnum supportDatasourceEnum) {
        databaseHolder.set(supportDatasourceEnum.toString());
    }

    /**
     * 取得当前数据源
     *
     * @return
     */
    public static String getDatabaseHolder() {
        return databaseHolder.get();
    }

    /**
     * 添加数据源
     *
     * @param supportDatasourceEnum
     */
    public static void addDatasource(SupportDatasourceEnum supportDatasourceEnum) {
        dataSourceSet.add(supportDatasourceEnum);
    }

    /**
     * 获取当期应用所支持的所有数据源
     *
     * @return
     */
    public static HashSet getDataSourceSet() {
        return dataSourceSet;
    }

    /**
     * 清除上下文数据
     */
    public static void clear() {
        databaseHolder.remove();
    }

}

spring内部的AbstractRoutingDataSource动态路由数据源里面有一个抽象方法叫做
determineCurrentLookupKey，这个方法适用于提供给开发者自定义对应数据源的查询key。

package mutidatasource.core;

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

/**
 * @author idea
 * @data 2020/3/7
 */
public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        String dataSource = DataSourceContextHolder.getDatabaseHolder();
        return dataSource;
    }
}

这里我使用的druid数据源，所以配置数据源的配置类如下：这里面我默认该应用配置类PROD数据源，用于测试使用。

package mutidatasource.core;

import com.alibaba.druid.pool.DruidDataSource;
import lombok.extern.slf4j.Slf4j;
import mutidatasource.enums.SupportDatasourceEnum;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Component;

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

/**
 * @author idea
 * @data 2020/3/7
 */
@Slf4j
@Component
public class DynamicDataSourceConfiguration {


    @Bean
    @Primary
    @ConditionalOnMissingBean
    public DataSource dataSource() {
        System.out.println("init datasource");
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        //设置原始数据源
        HashMap dataSourcesMap = new HashMap<>();
        HashSet dataSet = DataSourceContextHolder.getDataSourceSet();
        for (SupportDatasourceEnum supportDatasourceEnum : dataSet) {
            DataSource dataSource = this.createDataSourceProperties(supportDatasourceEnum);
            dataSourcesMap.put(supportDatasourceEnum.toString(), dataSource);
        }
        dynamicDataSource.setTargetDataSources(dataSourcesMap);
        dynamicDataSource.setDefaultTargetDataSource(createDataSourceProperties(SupportDatasourceEnum.PRE_DB));
        return dynamicDataSource;
    }

    private synchronized DataSource createDataSourceProperties(SupportDatasourceEnum supportDatasourceEnum) {
        DruidDataSource druidDataSource = new DruidDataSource();
        druidDataSource.setUrl(supportDatasourceEnum.getUrl());
        druidDataSource.setUsername(supportDatasourceEnum.getUsername());
        druidDataSource.setPassword(supportDatasourceEnum.getPassword());
        //具体配置
        druidDataSource.setMaxActive(100);
        druidDataSource.setInitialSize(5);
        druidDataSource.setMinIdle(1);
        druidDataSource.setMaxWait(30000);
        //间隔多久才进行一次检测，检测需要关闭的空闲连接，单位是毫秒
        druidDataSource.setTimeBetweenConnectErrorMillis(60000);
        return druidDataSource;
    }


}

好了现在一个基础的数据源注入已经可以了，那么我们该如何借助注解来实现动态切换数据源的操作呢？

为此，我设计了一个叫做UsingDataSource的注解，通过利用该注解来识别当前线程所需要使用的数据源操作：

package mutidatasource.annotation;

import mutidatasource.enums.SupportDatasourceEnum;

import java.lang.annotation.*;

/**
 * @author idea
 * @data 2020/3/7
 */
@Target({ElementType.METHOD,ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface UsingDataSource {

    SupportDatasourceEnum type()  ;
}

然后，借助了spring的aop来做切面拦截：

package mutidatasource.core;

import lombok.extern.slf4j.Slf4j;
import mutidatasource.annotation.UsingDataSource;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.*;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Arrays;

/**
 * @author idea
 * @data 2020/3/7
 */
@Slf4j
@Aspect
@Configuration
public class DataSourceAspect {

    public DataSourceAspect(){
        System.out.println("this is init");
    }



    @Pointcut("@within(mutidatasource.annotation.UsingDataSource) || " +
            "@annotation(mutidatasource.annotation.UsingDataSource)")
    public void pointCut(){

    }

    @Before("pointCut() && @annotation(usingDataSource)")
    public void doBefore(UsingDataSource usingDataSource){
        log.debug("select dataSource---"+usingDataSource.type());
        DataSourceContextHolder.setDatabaseHolder(usingDataSource.type());
    }

    @After("pointCut()")
    public void doAfter(){
        DataSourceContextHolder.clear();
    }

}

测试类如下所示：

package mutidatasource.controller;

import lombok.extern.slf4j.Slf4j;
import mutidatasource.annotation.UsingDataSource;
import mutidatasource.enums.SupportDatasourceEnum;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

/**
 * @author idea
 * @data 2020/3/8
 */
@RestController
@RequestMapping(value = "/test")
@Slf4j
public class TestController {

    @Autowired
    private JdbcTemplate jdbcTemplate;


    @GetMapping(value = "/testDev")
    @UsingDataSource(type=SupportDatasourceEnum.DEV_DB)
    public void testDev() {
        showData();
    }

    @GetMapping(value = "/testPre")
    @UsingDataSource(type=SupportDatasourceEnum.PRE_DB)
    public void testPre() {
        showData();
    }

    private void showData() {
        jdbcTemplate.queryForList("select * from test1").forEach(row -> log.info(row.toString()));
    }


}

最后 启动springboot服务，通过使用注解即可测试对应功能。

关于AbstractRoutingDataSource 动态路由数据源的注入原理，

可以看到这个内部类里面包含了多种用于做数据源映射的map数据结构。

在该类的最底部，有一个determineCurrentLookupKey函数，也就是上边我们所提及的使用于查询当前数据源key的方法。

具体代码如下：

/**
     * Retrieve the current target DataSource. Determines the
     * {@link #determineCurrentLookupKey() current lookup key}, performs
     * a lookup in the {@link #setTargetDataSources targetDataSources} map,
     * falls back to the specified
     * {@link #setDefaultTargetDataSource default target DataSource} if necessary.
     * @see #determineCurrentLookupKey()
     */
    protected DataSource determineTargetDataSource() {
        Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
        //这里面注入我们当前线程使用的数据源
        Object lookupKey = determineCurrentLookupKey();
        //在初始化数据源的时候需要我们去给resolvedDataSources进行注入
        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 + "]");
        }
        return dataSource;
    }

    /**
     * Determine the current lookup key. This will typically be
     * implemented to check a thread-bound transaction context.
     * 
Allows for arbitrary keys. The returned key needs
     * to match the stored lookup key type, as resolved by the
     * {@link #resolveSpecifiedLookupKey} method.
     */
    @Nullable
    protected abstract Object determineCurrentLookupKey();

而在该类的afterPropertiesSet里面，又有对于初始化数据源的注入操作，这里面的targetDataSources 正是上文中我们对在初始化数据源时候注入的信息。

@Override
    public void afterPropertiesSet() {
        if (this.targetDataSources == null) {
            throw new IllegalArgumentException("Property 'targetDataSources' is required");
        }
        this.resolvedDataSources = new HashMap<>(this.targetDataSources.size());
        this.targetDataSources.forEach((key, value) -> {
            Object lookupKey = resolveSpecifiedLookupKey(key);
            DataSource dataSource = resolveSpecifiedDataSource(value);
            this.resolvedDataSources.put(lookupKey, dataSource);
        });
        if (this.defaultTargetDataSource != null) {
            this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
        }
    }