动态数据源 @DS 注解源码解析

参考：动态数据源切换——@DS 注解源码解析

前言

借助 dynamic-datasource 可实现多数据源读写，其核心注解@DS用来动态切换数据源。

下面介绍@DS注解的实现原理。

如何使用

在 pom 中引入依赖：

<!-- spring-boot 1.5.x 2.x.x -->
<dependency>
    <groupId>com.baomidou</groupId>
    <artifactId>dynamic-datasource-spring-boot-starter</artifactId>
    <version>${version}</version>
</dependency>

配置数据源：

spring:
  datasource:
    dynamic:
      enabled: true        # 启用动态数据源，默认 true
      primary: master      # 设置默认的数据源或者数据源组，默认值即为 master
      strict: false        # 严格匹配数据源，默认 false. true 未匹配到指定数据源时抛异常，false 使用默认数据源
      grace-destroy: false # 是否优雅关闭数据源，默认为 false，设置为 true 时，关闭数据源时如果数据源中还存在活跃连接，至多等待 10s 后强制关闭
      datasource:
        master:
          url: jdbc:mysql://xx.xx.xx.xx:3306/dynamic
          username: root
          password: 123456
          driver-class-name: com.mysql.jdbc.Driver # 3.2.0 开始支持 SPI 可省略此配置
        slave_1:
          url: jdbc:mysql://xx.xx.xx.xx:3307/dynamic
          username: root
          password: 123456
          driver-class-name: com.mysql.jdbc.Driver
        slave_2:
          url: jdbc:mysql://xx.xx.xx.xx:3308/dynamic
          username: root
          password: 123456
          driver-class-name: com.mysql.jdbc.Driver
        # 以上会配置一个默认库 master，一个组 slave 下有两个子库 slave_1 slave_2

使用@DS注解切换数据源，@DS注解可加在类或者方法上，方法上注解优先于类上注解。对于没有使用@DS注解的类或方法，使用默认数据源。以下是官方示例：

@Service
@DS("slave")
public class UserServiceImpl implements UserService {

    @Autowired
    private JdbcTemplate jdbcTemplate;

    public List selectAll() {
        return jdbcTemplate.queryForList("select * from user");
    }

    @Override
    @DS("slave_1")
    public List selectByCondition() {
        return jdbcTemplate.queryForList("select * from user where age >10");
    }
}

实现原理

进入自动配置类DynamicDataSourceAutoConfiguration，可以看到其注册了两个切面通知：

DynamicDataSourceAnnotationInterceptor拦截器用来增强带有@DS注解的方法，获取注解中的值（数据源名），交给DynamicDataSourceContextHolder管理：

DynamicDataSourceContextHolder中维护了一个ThreadLocal<Deque<String>>，其中的 Deque 是用来存放数据源名称的栈：

public final class DynamicDataSourceContextHolder {

    private static final ThreadLocal<Deque<String>> LOOKUP_KEY_HOLDER = new NamedThreadLocal<Deque<String>>("dynamic-datasource") {
        @Override
        protected Deque<String> initialValue() {
            return new ArrayDeque<>();
        }
    };

    private DynamicDataSourceContextHolder() {
    }

    /**
     * 获得当前线程数据源
     */
    public static String peek() {
        return LOOKUP_KEY_HOLDER.get().peek();
    }

    /**
     * 设置当前线程数据源
     */
    public static String push(String ds) {
        String dataSourceStr = StringUtils.isEmpty(ds) ? "" : ds;
        LOOKUP_KEY_HOLDER.get().push(dataSourceStr);
        return dataSourceStr;
    }

    /**
     * 清空当前线程数据源
     */
    public static void poll() {
        Deque<String> deque = LOOKUP_KEY_HOLDER.get();
        deque.poll();
        if (deque.isEmpty()) {
            LOOKUP_KEY_HOLDER.remove();
        }
    }

    /**
     * 强制清空本地线程
     */
    public static void clear() {
        LOOKUP_KEY_HOLDER.remove();
    }
}

在DynamicDataSourceAutoConfiguration中，会在容器中注册DynamicRoutingDataSource：

@Bean
@ConditionalOnMissingBean
public DataSource dataSource() {
    DynamicRoutingDataSource dataSource = new DynamicRoutingDataSource();
    dataSource.setPrimary(properties.getPrimary());
    dataSource.setStrict(properties.getStrict());
    dataSource.setStrategy(properties.getStrategy());
    dataSource.setP6spy(properties.getP6spy());
    dataSource.setSeata(properties.getSeata());
    return dataSource;
}

DynamicRoutingDataSource继承自AbstractRoutingDataSource，AbstractRoutingDataSource继承自AbstractDataSource，AbstractDataSource实现了DataSource。

DynamicRoutingDataSource最终会被注入到 MyBatis 等 ORM 框架中，比如MybatisAutoConfiguration用其初始化 SqlSessionFactory：

@Bean
@ConditionalOnMissingBean
public SqlSessionFactory sqlSessionFactory(DataSource dataSource) throws Exception {
  SqlSessionFactoryBean factory = new SqlSessionFactoryBean();
  factory.setDataSource(dataSource);
  // ...
}

继而用来初始化 Executor:

// org.apache.ibatis.session.defaults.DefaultSqlSessionFactory#openSessionFromDataSource
private SqlSession openSessionFromDataSource(ExecutorType execType, TransactionIsolationLevel level, boolean autoCommit) {
  Transaction tx = null;
  try {
    final Environment environment = configuration.getEnvironment();
    final TransactionFactory transactionFactory = getTransactionFactoryFromEnvironment(environment);
    // ⭐ environment.getDataSource() 会获取到 DynamicRoutingDataSource
    tx = transactionFactory.newTransaction(environment.getDataSource(), level, autoCommit);
    final Executor executor = configuration.newExecutor(tx, execType);
    return new DefaultSqlSession(configuration, executor, autoCommit);
  } catch (Exception e) {
    closeTransaction(tx); // may have fetched a connection so lets call close()
    throw ExceptionFactory.wrapException("Error opening session.  Cause: " + e, e);
  } finally {
    ErrorContext.instance().reset();
  }
}

Executor 在执行 SQL 时会从中获取数据库连接：

// org.apache.ibatis.executor.BaseExecutor#getConnection
protected Connection getConnection(Log statementLog) throws SQLException {
  // ⭐ 会从 DynamicRoutingDataSource 中获取数据源
  Connection connection = transaction.getConnection();
  if (statementLog.isDebugEnabled()) {
    return ConnectionLogger.newInstance(connection, statementLog, queryStack);
  } else {
    return connection;
  }
}

回到DynamicRoutingDataSource，其InitializingBean方法会加载 yml 配置文件中的数据源，并调用addDataSource(String ds, DataSource dataSource)方法将数据源存入dataSourceMap中：

前面说了，在获取数据库连接时（执行DataSource#getConnection方法），最终会调用到DynamicRoutingDataSource#getConnection，进入该方法：

可以看到上面的DynamicDataSourceContextHolder.peek()会取到栈顶的数据源名称，所以DynamicRoutingDataSource#getConnection会返回栈顶的数据源。

这里解释一下为什么DynamicDataSourceContextHolder要用栈：为了支持嵌套切换，如 ABC 三个 service 都是不同的数据源，其中 A 的某个业务要调 B 的方法，B 的方法需要调用 C 的方法，一级一级调用切换，形成了链，此时就需要用栈保证先进后出。

在上面的getDataSource()中可以看到，如果上次入栈的数据源名为空，就取默认数据源，否则从dataSourceMap中获取对应的数据源，此时完成数据源的切换。

最后，在执行完invocation.proceed()之后，保证出栈，以免影响后面的操作：

另，参考上面的代码，可以通过如下方式来实现手动切换数据源：

try {
    DynamicDataSourceContextHolder.push("ds1");
} finally {
    DynamicDataSourceContextHolder.poll();
}