Spring多数据源配置

项目团队最近需要更换框架，临时搭建一套组合框架。小项目，两个数据库：业务库，配置库。根据实际业务，动态切换。
之前对这块配置处理没有什么了解，看了一些资料以及以前框架的实现，了解了下思路，做个笔记整理：
1、自定义一个DataSource，Map<String,javax.sql.DataSource>存放所有数据源
2、重写getConnection(),根据key值从map中获取DataSource，并返回连接
3、定义一个上下文DataRouteContext 利用ThreadLocal 存储当前线程当前需要获取的数据源key值
4、自定义一个注解@DataSource("key")
5、自定义一个Aspect，拦截所有实现@DataSource注解的方法，将key值set到DataSourceContext

问题：如何把数据源key标记还原

代码实现

自定义数据源

public class DynamicDataSource extends AbstractDataSource {

    /** 默认数据源 */
    private DataSource defaultSource;
    /** 扩展数据源 */
    private Map<String,DataSource> sourceMap;


    public void setSourceMap(Map<String, DataSource> map) {
        if (sourceMap != null)
            throw new IllegalArgumentException("already injected by groupList");
        this.sourceMap = map;
    }

    public void setDefaultSource(DataSource defaultSource) {
        if(defaultSource == null)
            throw new IllegalArgumentException("defaultSource not null");
        this.defaultSource = defaultSource;
    }

    /**
     * <p>Attempts to establish a connection with the data source that
     * this <code>DataSource</code> object represents.
     *
     * @return a connection to the data source
     * @throws SQLException if a database access error occurs
     */
    @Override
    public Connection getConnection() throws SQLException {
        return getConnection(null,null);
    }

    /**
     * <p>Attempts to establish a connection with the data source that
     * this <code>DataSource</code> object represents.
     *
     * @param username the database user on whose behalf the connection is
     *                 being made
     * @param password the user's password
     * @return a connection to the data source
     * @throws SQLException if a database access error occurs
     * @since 1.4
     */
    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        DataSource ds = null;
        String routeName = DataRouteContext.getRoute();

        if (routeName != null) {
            DataRouteLogger.info("dataSource changed , current dataSource is:"+routeName);
            ds = sourceMap.get(routeName);
        } else {
            DataRouteLogger.info("current dataSource is:defaultSource");
            ds = this.defaultSource;
        }
        if (ds == null){
            DataRouteLogger.error("dataSource is:" + routeName + " not found");
            throw new IllegalArgumentException("dataSource is: " + routeName + "not found");
        }
        if(username == null || password == null) {
            return ds.getConnection();
        }
        return ds.getConnection(username, password);

    }
}

对应XML配置

<bean id="dataSource" class="wiki.zhanglong.pay.framework.core.dataroute.DynamicDataSource">
        <!-- 业务库 -->
        <property name="defaultSource" ref="busDataSource" />
        <!-- 路由库配置-->
        <property name="sourceMap">
            <map>
                <!-- 配置库-->
                <entry key="configDataSource" value-ref="configDataSource"/>
                <!-- 如果还要用到其它的库，在此添加-->
            </map>
        </property>

    </bean>

    具体的单个数据源配置省略

上下文DataRouteContext

public class DataRouteContext {

    private static ThreadLocal<Deque<String>> route = new ThreadLocal<>();

    public static String getRoute(){
        Deque<String> deque = route.get();
        if (deque == null || deque.size() == 0) {
            return null;
        }
        return deque.pop();

    }

    public static void setRoute(String routeName){
        Deque<String> deque = route.get();
        if (deque == null) {
            deque = new LinkedList<>();
            route.set(deque);
        }
        deque.push(routeName);
    }

    public static void reset(){
        Deque<String> routeDeque = route.get();
        if (routeDeque != null && routeDeque.size() > 0) {
            routeDeque.clear();
        }
    }

}

路由注解

@Target({ TYPE, METHOD })
@Retention(RUNTIME)
public @interface DataRoute {
    String value();
}

Aspect

@Aspect
@Component
@Order(1)
public class DataRouteAspect {

//    @Around("execution(public * *(..)) && @annotation(dataRoute))")
    @Around("@annotation(dataRoute)")
    public Object setRouteName(ProceedingJoinPoint jp, DataRoute dataRoute) throws Throwable {
        String routeKey = dataRoute.value();
        DataRouteLogger.info("Aspect 数据路由设置为："+routeKey);
        if (StringUtils.isNotBlank(routeKey)) {
            DataRouteContext.setRoute(routeKey);
        }
        return jp.proceed();;
    }
}

使用方式

在需要切换数据源的方法上，配置路由注解

    @DataRoute("configDataSource")
    @Override
    public User selectUser(String userId) {
        //数据库操作
        User user = userDAO.selectUserById(userId);
        return user;
    }

BUG修复

在PageHelper插件使用中，上述代码还有点小BUG，修改进行改造。
Spring 多数据源配置(2)