For the common design modules in the microservice architecture, we usually need to use druid as our data connection pool. When the architecture is expanded, the data storage servers that we usually face will gradually increase, from the original single database architecture to the complex multi database architecture.

When we need to query multiple databases in the business layer, we usually need to specify the corresponding data source dynamically when executing sql.

Spring's AbstractRoutingDataSource just provides this function point for us. Next, I will use a simple case based on springboot+aop to realize how to switch different data sources for data reading operation through custom annotation. At the same time, I will also explain the content of some source codes.

First of all, we need to customize a data source information specifically used to declare which data source information the current java application needs to use:

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.*;

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

    SupportDatasourceEnum[] datasourceType();
}

For convenience, I have configured all the data source addresses used in the test in lai enum. If you need to deal with them flexibly later, you can extract these configuration information and put it on some configuration centers.

package mutidatasource.enums;

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

/**
 * Currently supported data source information
 *
 * @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();
    }
}

The reason to create the @ AppDataSource annotation is to annotate it on the springboot startup class:

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);
    }

}

With the help of springboot's ImportSelector, you can customize a register to get the data source type specified by the annotation of the start class header:

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];
    }


}

OK, now we can get the corresponding data source type information. Here you will see a role called DataSourceContextHolder. This object is mainly used for the unified distribution and management of the data source information of each request thread.

In the multi concurrency scenario, in order to prevent the data source of different thread requests from "channeling" to each other, we usually use threadlocal for processing. Each thread is assigned a specified internal copy variable. Before the current thread ends, remember to destroy the corresponding thread copy.

package mutidatasource.core;

import mutidatasource.enums.SupportDatasourceEnum;

import java.util.HashSet;

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

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

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

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

    /**
     * Get current data source
     *
     * @return
     */
    public static String getDatabaseHolder() {
        return databaseHolder.get();
    }

    /**
     * add data source
     *
     * @param supportDatasourceEnum
     */
    public static void addDatasource(SupportDatasourceEnum supportDatasourceEnum) {
        dataSourceSet.add(supportDatasourceEnum);
    }

    /**
     * Get all data sources supported by current application
     *
     * @return
     */
    public static HashSet<SupportDatasourceEnum> getDataSourceSet() {
        return dataSourceSet;
    }

    /**
     * Clear context data
     */
    public static void clear() {
        databaseHolder.remove();
    }

}

There is an abstract method called determineCurrentLookupKey in the AbstractRoutingDataSource dynamic routing data source in spring. This method is suitable for developers to customize the query key of the corresponding data source.

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;
    }
}

Here I use the druid data source, so the configuration class of the configuration data source is as follows: in this case, I use the application configuration class PROD data source by default for testing.

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();
        //Set original data source
        HashMap<Object, Object> dataSourcesMap = new HashMap<>();
        HashSet<SupportDatasourceEnum> 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());
        //Specific configuration
        druidDataSource.setMaxActive(100);
        druidDataSource.setInitialSize(5);
        druidDataSource.setMinIdle(1);
        druidDataSource.setMaxWait(30000);
        //How often is the interval detected? It is used to detect idle connections that need to be closed, in milliseconds
        druidDataSource.setTimeBetweenConnectErrorMillis(60000);
        return druidDataSource;
    }


}

Now that a basic data source injection is ready, how can we use annotations to dynamically switch data sources?

For this reason, I designed a annotation called UsingDataSource to identify the data source operations needed by the current thread:

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()  ;
}

Then, with the help of spring's aop, we do the section interception:

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();
    }

}

The test classes are as follows:

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()));
    }


}

Finally, start the spring boot service, and test the corresponding functions by using annotations.

About the injection principle of AbstractRoutingDataSource dynamic routing datasource,

You can see that this inner class contains a variety of map data structures for data source mapping.

 

At the bottom of the class, there is a determineCurrentLookupKey function, which is the method we mentioned above to query the current data source key.

The specific code is as follows:

 /**
     * 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");
        //It injects the data source used by our current thread
        Object lookupKey = determineCurrentLookupKey();
        //When initializing the data source, we need to give resolvedDataSources Injection
        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.
     * <p>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();

In the afterproperties set of this class, there are injection operations for initializing the data source. The targetDataSources in this is exactly the information we injected when initializing the data source.

 @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);
        }
    }