基于AbstractRoutingDataSource的多数据源切换实现

最近由于项目需要，在同一个工程中需要访问两个数据源，在调用不同的dao层的时候访问不同的数据源。当然，原因是基础组件那边导致的，这里就不详细描述，只讲解决方法。

首先我们要了解一下AbstractRoutingDataSource是个什么东西，参考这篇文章：https://blog.csdn.net/u011463444/article/details/72842500

我们看AbstractRoutingDataSource类对应的注释：

Abstract {@link javax.sql.DataSource} implementation that routes {@link #getConnection()}
calls to one of various target DataSources based on a lookup key. The latter is usually
(but not necessarily) determined through some thread-bound transaction context.

大概意思就是getConnection()根据查找lookup key键对不同目标数据源的调用，通常是通过(但不一定)某些线程绑定的事物上下文来实现。通过这我们知道可以实现： 
- 多数据源的动态切换，在程序运行时，把数据源数据源动态织入到程序中，灵活的进行数据源切换。 
- 基于多数据源的动态切换，我们可以实现读写分离，这么做缺点也很明显，无法动态的增加数据源。

对于事务，只支持单库事务，也就是说切换数据源要在开启事务之前执行。 
spring DataSourceTransactionManager进行事务管理，开启事务，会将数据源缓存到DataSourceTransactionObject对象中进行后续的commit rollback等事务操作。

从源码的determineTargetDataSource可以知道，我们只需要实现determineCurrentLookupKey()方法即可：

protected DataSource determineTargetDataSource() {
		Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
		Object lookupKey = determineCurrentLookupKey();
		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;
	}

那么关注点主要在以下几点：

1、继承AbstractRoutingDataSource类，实现determineCurrentLookupKey()方法

2、通过AOP实现在执行SQL前切换数据源

下面是代码

首先我们要定义一个注解，并在需要切换数据源的方法（或类）上使用

@Target({ElementType.TYPE, ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface DataSource {
    String value();
}

这里遇到了一个使用上的问题，就是mybatis的dao层使用AOP失败的问题，之前的文章已经说明了原因，这里不再复述，笔者是在dao层上又加了一层做AOP。

接下来是一个用来获取数据源的上下文，一般实现用threadlocal

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

    public DataSourceContextHolder() {
    }

    public static void setSourceKey(String dbType) {
        CONTEXT_HOLDER.set(dbType);
    }

    public static String getSourceKey() {
        return (String)CONTEXT_HOLDER.get();
    }

    public static void clearSourceKey() {
        CONTEXT_HOLDER.remove();
    }
}

接下来是一个继承AbstractRoutingDataSource的代理数据源，用来设置数据源

public class ProxyDataSource extends AbstractRoutingDataSource {
    public ProxyDataSource() {
    }

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

那么还需要一个切面，用来在执行SQL前告诉DataSourceContextHolder当前线程切换到哪个数据源上

@Aspect
@Component
public class DataSourceInterceptor {
    private static final Logger LOGGER = LoggerFactory.getLogger(DataSourceInterceptor.class);

    public DataSourceInterceptor() {
    }

    @Pointcut("@annotation(org.dal.multisource.DataSource)")
    public void dataSourceAspect() {
    }

    @Around("dataSourceAspect()")
    public Object preInterceptor(ProceedingJoinPoint proceedingJoinPoint) throws Throwable {
        Class target = proceedingJoinPoint.getTarget().getClass();
        MethodSignature methodSignature = (MethodSignature)proceedingJoinPoint.getSignature();
        Class[] interfaces = target.getInterfaces();
        Class[] object;
        if(interfaces != null && interfaces.length > 0) {
            object = target.getInterfaces();
            int e = object.length;

            for(int i = 0; i < e; ++i) {
                Class clazz = object[i];
                this.resolveDataSource(clazz, methodSignature.getMethod());
            }
        } else {
            this.resolveDataSource(target, methodSignature.getMethod());
        }

        Object result;
        try {
            result = proceedingJoinPoint.proceed();
        } catch (Throwable e) {
            LOGGER.error("eval method error", e);
            throw e;
        } finally {
            DataSourceContextHolder.clearSourceKey();
        }

        return result;
    }

    private void resolveDataSource(Class<?> clazz, Method method) {
        try {
            DataSource dataSource;
            if(method.isAnnotationPresent(DataSource.class)) {
                dataSource = method.getAnnotation(DataSource.class);
                DataSourceContextHolder.setSourceKey(dataSource.value());
                return;
            }

            if(clazz.isAnnotationPresent(DataSource.class)) {
                dataSource = clazz.getAnnotation(DataSource.class);
                DataSourceContextHolder.setSourceKey(dataSource.value());
            }
        } catch (Exception e) {
            LOGGER.error("resolve data source error", e);
        }

    }
}

接下来是数据源的配置，Spring提供了多种配置方式，网上也以XML配置居多，这里使用的是@Configuration注解配置的方式。

@Configuration
@EnableTransactionManagement(proxyTargetClass = true)
@MapperScan(basePackages = "org.dal.dao", sqlSessionFactoryRef =
        "sqlSessionFactory")
public class DBConfig {

    @Autowired
    private String jdbcRef;
    @Autowired
    private String jdbcRef1;
    @Autowired
    private int dbMaxPoolSize;

    @Bean(destroyMethod = "close")
    public DataSource dataSource1() {
        DataSource dataSource = new DataSource();
        dataSource.setJdbcRef(jdbcRef);
        dataSource.setPoolType("hikaricp");
        dataSource.setMinPoolSize(5);
        dataSource.setMaxPoolSize(dbMaxPoolSize);
        dataSource.setCheckoutTimeout(1000);
        dataSource.setPreferredTestQuery("SELECT 1");
        dataSource.init();
        return dataSource;
    }

    @Bean(destroyMethod = "close")
    public DataSource dataSource2() {
        DataSource dataSource = new DataSource();
        dataSource.setJdbcRef(jdbcRef1);
        dataSource.setPoolType("hikaricp");
        dataSource.setMinPoolSize(5);
        dataSource.setMaxPoolSize(dbMaxPoolSize);
        dataSource.setCheckoutTimeout(1000);
        dataSource.setPreferredTestQuery("SELECT 1");
        dataSource.init();
        return dataSource;
    }

    @Primary
    @Bean
    public ProxyDataSource proxyDataSource() {
        ProxyDataSource proxyDataSource = new ProxyDataSource();
        Map<Object, Object> targetDataSourcesMap = Maps.newHashMap();
        targetDataSourcesMap.put("dataSource1", dataSource1());
        targetDataSourcesMap.put("dataSource2t", dataSource2());
        proxyDataSource.setTargetDataSources(targetDataSourcesMap);
        proxyDataSource.setDefaultTargetDataSource(dataSourceMagneto());
        return proxyDataSource;
    }

    @Bean
    public DataSourceTransactionManager transactionManager() {
        return new DataSourceTransactionManager(proxyDataSource());
    }

    @Bean
    public SqlSessionFactoryBean sqlSessionFactoryBean() {
        try {
            SqlSessionFactoryBean sqlSessionFactoryBean = new SqlSessionFactoryBean();
            sqlSessionFactoryBean.setDataSource(proxyDataSource());
            sqlSessionFactoryBean.setConfigLocation(new ClassPathResource("mybatis/sqlmap-config.xml"));
            return sqlSessionFactoryBean;
        } catch (Exception e) {
            throw new RuntimeException("创建SqlSessionFactoryBean异常", e);
        }
    }

    @Bean
    public SqlSessionFactory sqlSessionFactory() {
        try {
            return sqlSessionFactoryBean().getObject();
        } catch (Exception e) {
            throw new RuntimeException("创建SqlSessionFactory异常", e);
        }
    }

    @Bean
    public SqlSessionTemplate sqlSessionTemplate() {
        return new SqlSessionTemplate(sqlSessionFactory());
    }

}

上面的配置可能会有一些问题，因为不能直接使用工程中的代码，具体的配置可能需要调整，不过整体的思路是这样的。

-------------分割线--------------------

这个实现是有问题的，因为事务开启后就不能切换数据源了，即使设置了切换也不行，只能事务开启前切换数据源。目前使用编程式事务，执行方法前先走AOP切换数据源，然后方法中开启事务，这样虽然用起来很麻烦，但是至少保证切换数据源后的事务是没问题的，但是跨数据源的事务还是没有办法，据说可以通过JTA实现，暂时还没尝试。

--------------分割线--------------------

另一个问题，笔者使用了编程式事务，但是在事务中的一个判断条件后，没有提交也没有回滚，直接就返回了，这样会造成一个bug，当线程池中的同一个线程被复用时，切换数据源的操作将会失效，虽然这是一个低级失误，但是背后的原理应该有深度，不过目前还没找到合理的解释，只能理解为事务是和数据源绑定的，也就是开启事务后数据源不可变，但是由于事务在线程中没有结束，所以后续复用该线程的时候变更数据源失效。