[Springboot] Implement DynamicDataSource to configure multiple data sources – Part1

Main requirements

In a Springboot project, multiple data sources (connected to multiple different databases) need to be used. MybatisPlus is used as the persistence layer solution in the project. It needs to be similar to adding different annotations to a certain Mybatis Mapper, so that the mapper can connect to different data sources. As shown below:

@Mapper
public interface TestMapper {
    @Select("XXXXXX")
    void Test();

    @Insert("XXXXXX")
    @Datasource("db-test1")
    void Test1();

    @Select("XXXXXX")
    @Datasource("db-test2")
    void test2();

    @Select("XXXXXX")
    @Datasource("db-test3")
    void test3();
}

There is an annotation called @Datasource. This annotation can be configured on the class or on the method. When configured on a class, it means that all methods in the current mapper use the data source configured by the class. The configuration above the method indicates that the current method uses the configured data source. If both class and method are configured, the method will prevail. If neither class nor method has @Datasource, it means the default data source connection is used.

Main implementation

Let’s first clarify the above requirements and build the framework. Implement a simple multi-data source query. Here we do not consider the annotation method to switch the data source, because using the annotation method to switch the data source means that we need an aspect. Before each call to the Mapper method, we must check the value of the annotation in the aspect, and then Use this value as the key of the data source to obtain the data source. Let’s make a simple implementation, which is to manually switch to the data source required by the current method every time before the Service calls the Mapper method. Similar to the structure below.

@Service
public class TestServiceImpl implements TestService {

    @Autowired
    private TestMapper testMapper;
    
    @Override
    public void test(User user) {
        testMapper.insert(user);
    }

    @Override
    public void test1(User user) {
        XXXXX.setDataSourceName("db-test1");
        testMapper.updateById(user);
    }

    @Override
    public void test2(User user) {
        XXXXX.setDataSourceName("db-test2");
        testMapper.updateById(user);
    }
}

When we need to use the default data source to execute mapper, we can use it directly. When we want to use some other data source, we need to set the data source name to an identifier, and the dynamic data source receives this identifier. will be switched.

Detailed implementation

Now we will start to implement the above description, and explain the meaning of each step in detail during the implementation process.

1. Create DynamicDataSourceContext thread-safe class

Let’s first create the following class. This class is the basis of the entire implementation. Let’s first take a look at what this class does.

public class DynamicDataSourceContext {

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

    /**
     * Set the name of the current data source
     */
    public static void setDataSourceName(String dataSourceName) {
        HOLDER.set(dataSourceName);
    }

    /**
     * Get the name of the current data source
     */
    public static String getDataSourceName() {
        return HOLDER.get();
    }

    /**
     * Clear the name of the current data source
     */
    public static void clearDataSourceName() {
        HOLDER.remove();
    }

}

The code is very simple. By analyzing the code, you can see that a static ThreadLocal constant object is defined in the class. This ThreadLocal generic is of type String. The class contains three methods, all of which operate on the current static constant ThreadLocal. Three of the methods are very simple,

• The setDataSourceName() method is used to set the current data source name to ThreadLocal
• The getDataSourceName() method is used to obtain the name of the data source currently existing in ThreadLocal
• The clearDataSourceName() method is used to clear the data source name currently existing in ThreadLocal

After reading the code, let’s talk about why the book needs this class, and this class is very important. This is to consider the issue of data source switching in a multi-threaded environment. In a multi-threaded environment, if multiple threads access the same method at the same time, and each thread uses a different data source, then each data source under each thread needs to be dynamically switched. The meaning of using ThreadLocal is to use ThreadLocal to access the data source name to be used by the current thread. ThreadLocal is a thread-local storage mechanism that provides each thread with an independent copy of variables. Each thread can independently operate its own copy of variables without affecting the copies of variables in other threads. So when we switch to multiple data sources, each thread can know which data source the current thread needs to use through the variable copy in this ThreadLocal, thus avoiding the problem of multiple threads calling a method to access different data sources at the same time. The data source confusion caused by the time. At the same time, using ThreadLocal can also avoid memory leaks. If you are not familiar with ThreadLocal yet, you can first search and learn the relevant knowledge of ThreadLocal.

Generally speaking, the above method is mainly responsible for when we switch between multiple data sources:

1. Before each thread calls the Mapper method, use setDataSourceName() to confirm the data source called by the current thread.
2. During the actual call, obtain the data source name through getDataSourceName() to obtain the data source acquisition link.
3. After use, clear the data source name in ThreadLocal through clearDataSourceName() to ensure safe release.

2. Create the DynamicDataSource class

If the previous class is the basis of the implementation, this class is the core of the implementation. This class needs to inherit AbstractRoutingDataSource

Let’s take a look at the contents of this class:

public class DynamicDataSource extends AbstractRoutingDataSource {
    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContext.getDataSourceName();
    }

    public DynamicDataSource(DataSource defaultTargetDataSource, Map<Object, Object> targetDataSources) {
        super.setDefaultTargetDataSource(defaultTargetDataSource);
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }
}

This class is the multiple data source we will use in springboot. When initializing, it needs to pass in a default main data source (defaultTargetDataSource) and a collection of other required data sources (targetDataSources).

To know how to use it, we still have to look at its parent class AbstractRoutingDataSource. Let’s take a brief look at some of its main parameters and methods. (*The code is not complete, just a few necessary methods and parameters are selected)

public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {

    @Nullable
    private Map<Object, Object> targetDataSources;

    @Nullable
    private Object defaultTargetDataSource;

    @Nullable
    private Map<Object, DataSource> resolvedDataSources;

    @Nullable
    private DataSource resolvedDefaultDataSource;


    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        this.targetDataSources = targetDataSources;
    }

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

    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 & amp; & amp; (this.lenientFallback || lookupKey == null)) {
            dataSource = this.resolvedDefaultDataSource;
        }

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

    @Nullable
    protected abstract Object determineCurrentLookupKey();

}

First of all, the AbstractRoutingDataSource class inherits AbstractDataSource, and AbstractDataSource implements javax.sql’s DataSource, The DataSourceclass is a standard factory for obtaining database connections. It is a standard and all common database connection pools implement this interface (C3P0, DBCP, Hikari, Druid, etc.). So AbstractRoutingDataSource is also such a class for obtaining database links. Let’s take a look at its parameters first:

• defaultTargetDataSource: identifies the project’s default database connection
• targetDataSources: other database connections required by the project
• resolvedDefaultDataSource: The actual default data source obtained through defaultTargetDataSource.
• resolvedDataSources: Other data sources actually used obtained through targetDataSources.

The AbstractRoutingDataSource class implements an InitializingBean interface. This interface has only one abstract method that needs to be implemented by AbstractRoutingDataSource, which is the afterPropertiesSet method.

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

The afterPropertiesSet method is executed when the bean is initialized. In this method, we can see the process of converting targetDataSources and defaultTargetDataSource into resloved data sources. This means that when we initialize our dynamic data source object, all the data sources we need are ready.

Here comes the key point. Let’s take a look at the determineTargetDataSource() method. This method will be called wherever the data source needs to be obtained, such as when Mybatis performs addition, deletion, modification and query in mapper, or jdbcTemplate calls the database connection. method. Its main function is to tell the project which data source we currently need to use. There is a determineCurrentLookupKey() method, which is used to determine the name of the data source. We get the data source through this method and obtain the data source in resolveDataSources. If not, And the obtained data source name is also empty, then the default data source is returned. And this determineCurrentLookupKey() needs us to implement it in the implementation class. This implementation class is the DynamicDataSource we mentioned above

 protected DataSource determineTargetDataSource() {
        Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
        //Get the current data source name. This method needs to be implemented by yourself in the implementation class.
        Object lookupKey = this.determineCurrentLookupKey();
        //Get the data source from resolveDataSources through the data source name
        DataSource dataSource = (DataSource)this.resolvedDataSources.get(lookupKey);
        //If not, it means it is the default data source, and then directly return to the default data source.
        if (dataSource == null & amp; & amp; (this.lenientFallback || lookupKey == null)) {
            dataSource = this.resolvedDefaultDataSource;
        }

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

    @Nullable
    protected abstract Object determineCurrentLookupKey();

Let’s take a look at DynamicDataSource at the end and it will become clear.

public class DynamicDataSource extends AbstractRoutingDataSource {
    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContext.getDataSourceName();
    }

    public DynamicDataSource(DataSource defaultTargetDataSource, Map<Object, Object> targetDataSources) {
        super.setDefaultTargetDataSource(defaultTargetDataSource);
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }
}

When the method is initialized, the configured default data source and the set of other data sources are assigned values. Then the afterPropertiesSet() method is called to convert them into resolved data sources. (Or hand over the Bean to Spring management, and let the bean that changes the dynamic data source call the afterPropertiesSet() method when building)

Then when a database connection is needed, the system will call the determineTargetDataSource() method of the parent class AbstractRoutingDataSource, and the call to obtain the database connection name is implemented in the subclass . The implementation method in the subclass is determineCurrentLookupKey(), which is the thread name obtained directly from ThreadLocal.

3. Initialize each data source

DynamicDataSource requires two parameters when initializing, which are the default data source and a collection of other data sources.

This step is actually relatively simple. Directly make a Configuration class (maybe called the DynamicDataConfig class). This class reads the data source information configured in yml and converts this information into targetDataSources and defalutTargetDataSource< /strong>, and then initialize a DynamicDataSource through these two parameters and assemble it into the Bean. This part will not be described in detail. The relevant contents of yml are as follows:

spring:
  datasource:
    driver-class-name: com.mysql.cj.jdbc.Driver
    url: jdbc:mysql://****:3306/test
    username: root
    password: root
mutil-datasource:
  connection:
    - dbName: db-test1
      dbDriver: com.microsoft.sqlserver.jdbc.SQLServerDriver
      dbUrl: jdbc:sqlserver://****:1433;DatabaseName=db_test1
      dbUsername: root
      dbPassword: root
    - dbName: db-test2
      dbDriver: com.sybase.jdbc4.jdbc.SybDriver
      dbUrl: jdbc:sybase:Tds:****:5000/db_test2
      dbUsername: root
      dbPassword: root

@Service
public class TestServiceImpl implements TestService {

    @Autowired
    private TestMapper testMapper;
    
    @Override
    public void test(User user) {
        testMapper.insert(user);
    }

    @Override
    public void test1(User user) {
        DynamicDataSourceContext.setDataSourceName("db-test1");
        testMapper.updateById(user);
        DynamicDataSourceContext.clearDataSourceName()
    }

    @Override
    public void test2(User user) {
        DynamicDataSourceContext.setDataSourceName("db-test2");
        testMapper.updateById(user);
        DynamicDataSourceContext.clearDataSourceName()
    }
}