Spring boot + AOP build multi data source switching practice

Time:2020-7-27

For the commonly used design modules in microservice architecture, we usually need to use Druid as our data connection pool. When the architecture expands, the data storage servers we usually face will gradually increase, from the original single database architecture to the complex multi database architecture.

When the business layer needs to query multiple databases, we usually need to specify the corresponding datasource dynamically when executing SQL.

And spring’sAbstractRoutingDataSourceThis function point is just provided for us. Next, I will use a simple case based on springboot + AOP to realize how to switch different data sources to read data through user-defined annotations. At the same time, I will explain the content of part of the source code.

First of all, we need to customize a data source information that is specifically used to state what data source information is needed by current Java applications

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

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

 

Here, for convenience, I configure all the data source addresses used in the test in Lai enum. If you need to handle it flexibly, you can extract the 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 why you want to create this @ appdatasource annotation is to mark it on the startup class of 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);
    }

}

 

Use the importselector of springboot to customize a register to obtain the data source type specified by the annotation in the header of the boot class

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 that we can get the corresponding data source type information, you can see a role called datasourcecontextholder. This object is mainly used for the unified allocation and management of data source information of each request thread.

In the multi concurrency scenario, in order to prevent the data sources requested by different threads from running across each other, we usually use ThreadLocal for processing. Each thread is assigned a specified copy variable that belongs to its internal copy. 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 dataSourceSet = new HashSet<>();

    private static final ThreadLocal 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 getDataSourceSet() {
        return dataSourceSet;
    }

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

}

 

Spring’s internal abstractroutingdatasource dynamic routing data source contains an abstract method called
Determine current lookupkey. This method is suitable for developers to customize the query key corresponding to the 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: I default the application configuration class prod data source 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 up raw data source
        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());
        //Specific configuration
        druidDataSource.setMaxActive(100);
        druidDataSource.setInitialSize(5);
        druidDataSource.setMinIdle(1);
        druidDataSource.setMaxWait(30000);
        //How often is the detection interval? It detects the idle connections that need to be closed, in milliseconds
        druidDataSource.setTimeBetweenConnectErrorMillis(60000);
        return druidDataSource;
    }


}

 

Well, now that a basic data source injection is ready, how can we dynamically switch data sources with annotations?

To this end, I designed an annotation called usingdatasource to identify the data source operations that the current thread needs to use

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, it intercepts all aspects

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 springboot service and test the corresponding functions by using annotations.

About the injection principle of abstractroutingdatasource dynamic routing data source,

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

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

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

The specific codes are 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");
        //This injects the data source used by our current thread
        Object lookupKey = determineCurrentLookupKey();
        //When initializing the data source, we need to inject the resolved data sources
        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();

 

In the afterpropertiesset of this class, there is the injection operation for initializing the data source. The targetdatasources in this class are 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);
        }
    }