Elegant Dynamic Data Source Switching in Spring Boot with ThreadLocal and AbstractRoutingDataSource

Introduction

The author needed to fetch data from different databases and write to the current database, requiring dynamic data source switching. The MyBatis-Plus dynamic-datasource-spring-boot-starter was considered but failed due to project environment issues. Instead, a custom implementation using ThreadLocal and AbstractRoutingDataSource was built to mimic the starter's thread-level data source switching.

Core Concepts

ThreadLocal

ThreadLocal

provides thread-local variables, giving each thread its own copy to avoid concurrency issues. It stores values in a map keyed by the current thread instance, achieving isolation with increased memory but reduced synchronization overhead.

Role: Shared within a thread, isolated across threads.

Principle: Uses the current thread as key in the thread's internal map.

AbstractRoutingDataSource

Spring's AbstractRoutingDataSource selects a data source based on user-defined rules. Its abstract method determineCurrentLookupKey() is called before each database operation to decide which data source to use.

Implementation Environment

Spring Boot 2.4.8

MyBatis-Plus 3.2.0

Druid 1.2.6

Lombok 1.18.20

commons-lang3 3.10

Code Implementation

2.1 ThreadLocal Holder

A DataSourceContextHolder class manages the current thread's data source name via setDataSource, getDataSource, and removeDataSource methods.

public class DataSourceContextHolder {
    private static final ThreadLocal<String> DATASOURCE_HOLDER = new ThreadLocal<>();

    public static void setDataSource(String dataSourceName) {
        DATASOURCE_HOLDER.set(dataSourceName);
    }

    public static String getDataSource() {
        return DATASOURCE_HOLDER.get();
    }

    public static void removeDataSource() {
        DATASOURCE_HOLDER.remove();
    }
}

2.2 Dynamic Data Source Routing

DynamicDataSource

extends AbstractRoutingDataSource and overrides determineCurrentLookupKey to return the data source name from DataSourceContextHolder. The constructor accepts a default data source and a map of target data sources.

public class DynamicDataSource extends AbstractRoutingDataSource {

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

    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceContextHolder.getDataSource();
    }
}

2.3 Data Source Configuration

Two data sources (master and slave) are defined in application.yml using Druid. A configuration class creates the data source beans and registers the DynamicDataSource as the primary data source.

spring:
  datasource:
    type: com.alibaba.druid.pool.DruidDataSource
    druid:
      master:
        url: jdbc:mysql://xxxxxx:3306/test1?characterEncoding=utf-8&allowMultiQueries=true&zeroDateTimeBehavior=convertToNull&useSSL=false
        username: root
        password: 123456
        driver-class-name: com.mysql.cj.jdbc.Driver
      slave:
        url: jdbc:mysql://xxxxx:3306/test2?characterEncoding=utf-8&allowMultiQueries=true&zeroDateTimeBehavior=convertToNull&useSSL=false
        username: root
        password: 123456
        driver-class-name: com.mysql.cj.jdbc.Driver
      initial-size: 15
      min-idle: 15
      max-active: 200
      max-wait: 60000
      time-between-eviction-runs-millis: 60000
      min-evictable-idle-time-millis: 300000
      validation-query: "select 1"
      test-while-idle: true
      test-on-borrow: false
      test-on-return: false
      pool-prepared-statements: false
      connection-properties: false

@Configuration
public class DateSourceConfig {

    @Bean
    @ConfigurationProperties("spring.datasource.druid.master")
    public DataSource masterDataSource() {
        return DruidDataSourceBuilder.create().build();
    }

    @Bean
    @ConfigurationProperties("spring.datasource.druid.slave")
    public DataSource slaveDataSource() {
        return DruidDataSourceBuilder.create().build();
    }

    @Bean(name = "dynamicDataSource")
    @Primary
    public DynamicDataSource createDynamicDataSource() {
        Map<Object, Object> dataSourceMap = new HashMap<>();
        DataSource defaultDataSource = masterDataSource();
        dataSourceMap.put("master", defaultDataSource);
        dataSourceMap.put("slave", slaveDataSource());
        return new DynamicDataSource(defaultDataSource, dataSourceMap);
    }
}

Note: Exclude DataSourceAutoConfiguration to avoid circular dependency.

@SpringBootApplication(exclude = DataSourceAutoConfiguration.class)

2.4 Testing

A test_user table with a user_name column is created in both databases. The master contains 'master', the slave contains 'slave'. A controller endpoint sets the data source via DataSourceContextHolder, queries the table, and removes the data source.

@GetMapping("/getData.do/{datasourceName}")
public String getMasterData(@PathVariable("datasourceName") String datasourceName) {
    DataSourceContextHolder.setDataSource(datasourceName);
    TestUser testUser = testUserMapper.selectOne(null);
    DataSourceContextHolder.removeDataSource();
    return testUser.getUserName();
}

Results confirm switching works: passing "master" returns 'master', passing "slave" returns 'slave'. The author notes that MyBatis-Plus uses a stack-based ThreadLocal to support nested data source switching.

Optimizations

2.5.1 Annotation-Based Switching

To avoid repetitive boilerplate, a @DS annotation and an AOP aspect are introduced.

Annotation Definition

@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Inherited
public @interface DS {
    String value() default "master";
}

AOP Aspect

@Aspect
@Component
@Slf4j
public class DSAspect {

    @Pointcut("@annotation(com.jiashn.dynamic_datasource.dynamic.aop.DS)")
    public void dynamicDataSource() {}

    @Around("dynamicDataSource()")
    public Object datasourceAround(ProceedingJoinPoint point) throws Throwable {
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        DS ds = method.getAnnotation(DS.class);
        if (Objects.nonNull(ds)) {
            DataSourceContextHolder.setDataSource(ds.value());
        }
        try {
            return point.proceed();
        } finally {
            DataSourceContextHolder.removeDataSource();
        }
    }
}

Test

Two endpoints: /getMasterData.do (no annotation, uses default master) and /getSlaveData.do (annotated with @DS("slave")). Results match expectations.

2.5.2 Dynamic Data Source Registration at Runtime

Business requirement: load additional data sources from a database table at startup.

Data Source Entity

@Data
@Accessors(chain = true)
public class DataSourceEntity {
    private String url;
    private String userName;
    private String passWord;
    private String driverClassName;
    private String key;
}

Enhanced DynamicDataSource

The class now holds a reference to the target data source map and provides a createDataSource method that validates connections, creates Druid data sources, and updates the map. It also calls afterPropertiesSet() to refresh the resolved data sources.

@Slf4j
public class DynamicDataSource extends AbstractRoutingDataSource {

    private final Map<Object, Object> targetDataSourceMap;

    public DynamicDataSource(DataSource defaultDataSource, Map<Object, Object> targetDataSources) {
        super.setDefaultTargetDataSource(defaultDataSource);
        super.setTargetDataSources(targetDataSources);
        this.targetDataSourceMap = targetDataSources;
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceContextHolder.getDataSource();
    }

    public void createDataSource(List<DataSourceEntity> dataSources) {
        try {
            if (CollectionUtils.isNotEmpty(dataSources)) {
                for (DataSourceEntity ds : dataSources) {
                    Class.forName(ds.getDriverClassName());
                    DriverManager.getConnection(ds.getUrl(), ds.getUserName(), ds.getPassWord());
                    DruidDataSource dataSource = new DruidDataSource();
                    BeanUtils.copyProperties(ds, dataSource);
                    dataSource.setTestOnBorrow(true);
                    dataSource.setTestWhileIdle(true);
                    dataSource.setValidationQuery("select 1");
                    dataSource.init();
                    this.targetDataSourceMap.put(ds.getKey(), dataSource);
                }
                super.setTargetDataSources(this.targetDataSourceMap);
                super.afterPropertiesSet();
                return Boolean.TRUE;
            }
        } catch (ClassNotFoundException | SQLException e) {
            log.error("---程序报错---:{}", e.getMessage());
        }
        return Boolean.FALSE;
    }

    public boolean existsDataSource(String key) {
        return Objects.nonNull(this.targetDataSourceMap.get(key));
    }
}

Loading Data Sources on Startup

A CommandLineRunner reads from a test_db_info table, maps rows to DataSourceEntity, and calls dynamicDataSource.createDataSource().

@Component
public class LoadDataSourceRunner implements CommandLineRunner {
    @Resource
    private DynamicDataSource dynamicDataSource;
    @Resource
    private TestDbInfoMapper testDbInfoMapper;

    @Override
    public void run(String... args) throws Exception {
        List<TestDbInfo> testDbInfos = testDbInfoMapper.selectList(null);
        if (CollectionUtils.isNotEmpty(testDbInfos)) {
            List<DataSourceEntity> ds = new ArrayList<>();
            for (TestDbInfo testDbInfo : testDbInfos) {
                DataSourceEntity sourceEntity = new DataSourceEntity();
                BeanUtils.copyProperties(testDbInfo, sourceEntity);
                sourceEntity.setKey(testDbInfo.getName());
                ds.add(sourceEntity);
            }
            dynamicDataSource.createDataSource(ds);
        }
    }
}

Test

After startup, the dynamically added data source 'add_slave' is available and returns the expected data.

Conclusion

The article provides a complete, elegant implementation of dynamic data source switching in Spring Boot, progressing from basic ThreadLocal routing to annotation-driven AOP and runtime data source registration. The GitHub repository is available at github.com/lovejiashn/… .